Merge pull request #134 from Wodann/fix/issue-133

fix(codegen): only generate exposed functions

crates/mun_codegen/src/code_gen.rs

@@ -6,16 +6,18 @@ use inkwell::targets::TargetData;
 use inkwell::{
     module::{Linkage, Module},
     passes::{PassManager, PassManagerBuilder},
-    targets::{CodeModel, FileType, InitializationConfig, RelocMode, Target},
+    targets::{CodeModel, FileType, InitializationConfig, RelocMode, Target, TargetMachine},
     types::StructType,
     values::{BasicValue, GlobalValue, IntValue, PointerValue, UnnamedAddress},
     AddressSpace, OptimizationLevel,
 };
+use mun_target::spec;
 use std::io::{self, Write};
 use std::{
     path::{Path, PathBuf},
     sync::Arc,
 };
+use tempfile::NamedTempFile;
 
 mod linker;
 pub mod symbols;
@@ -42,6 +44,51 @@ impl From<LinkerError> for CodeGenerationError {
     }
 }
 
+pub struct ObjectFile {
+    target: spec::Target,
+    src_path: RelativePathBuf,
+    obj_file: NamedTempFile,
+}
+
+impl ObjectFile {
+    pub fn new(
+        target: &spec::Target,
+        target_machine: &TargetMachine,
+        src_path: RelativePathBuf,
+        module: Arc<inkwell::module::Module>,
+    ) -> Result<Self, failure::Error> {
+        let obj = target_machine
+            .write_to_memory_buffer(&module, FileType::Object)
+            .map_err(|e| CodeGenerationError::CodeGenerationError(e.to_string()))?;
+
+        let mut obj_file = tempfile::NamedTempFile::new()
+            .map_err(CodeGenerationError::CouldNotCreateObjectFile)?;
+        obj_file
+            .write(obj.as_slice())
+            .map_err(CodeGenerationError::CouldNotCreateObjectFile)?;
+
+        Ok(Self {
+            target: target.clone(),
+            src_path,
+            obj_file,
+        })
+    }
+
+    pub fn into_shared_object(self, out_dir: Option<&Path>) -> Result<PathBuf, failure::Error> {
+        // Construct a linker for the target
+        let mut linker = linker::create_with_target(&self.target);
+        linker.add_object(self.obj_file.path())?;
+
+        let output_path = assembly_output_path(&self.src_path, out_dir);
+
+        // Link the object
+        linker.build_shared_object(&output_path)?;
+        linker.finalize()?;
+
+        Ok(output_path)
+    }
+}
+
 /// A struct that can be used to build an LLVM `Module`.
 pub struct ModuleBuilder<'a, D: IrDatabase> {
     db: &'a D,
@@ -52,8 +99,8 @@ pub struct ModuleBuilder<'a, D: IrDatabase> {
 }
 
 impl<'a, D: IrDatabase> ModuleBuilder<'a, D> {
-    /// Construct module for the given `hir::FileId` at the specified output file location.
-    pub fn new(db: &'a mut D, file_id: FileId) -> Result<Self, failure::Error> {
+    /// Constructs module for the given `hir::FileId` at the specified output file location.
+    pub fn new(db: &'a D, file_id: FileId) -> Result<Self, failure::Error> {
         let target = db.target();
 
         // Construct a module for the assembly
@@ -91,8 +138,8 @@ impl<'a, D: IrDatabase> ModuleBuilder<'a, D> {
         })
     }
 
-    /// Construct a shared object at the specified output file location.
-    pub fn finalize(&self, out_dir: Option<&Path>) -> Result<PathBuf, failure::Error> {
+    /// Constructs an object file.
+    pub fn build(self) -> Result<ObjectFile, failure::Error> {
         let group_ir = self.db.group_ir(self.file_id);
         let file = self.db.file_ir(self.file_id);
 
@@ -120,39 +167,18 @@ impl<'a, D: IrDatabase> ModuleBuilder<'a, D> {
         // Debug print the IR
         //println!("{}", assembly_module.print_to_string().to_string());
 
-        // Generate object file
-        let obj_file = {
-            let obj = self
-                .target_machine
-                .write_to_memory_buffer(&self.assembly_module, FileType::Object)
-                .map_err(|e| CodeGenerationError::CodeGenerationError(e.to_string()))?;
-            let mut obj_file = tempfile::NamedTempFile::new()
-                .map_err(CodeGenerationError::CouldNotCreateObjectFile)?;
-            obj_file
-                .write(obj.as_slice())
-                .map_err(CodeGenerationError::CouldNotCreateObjectFile)?;
-            obj_file
-        };
-
-        let target = self.db.target();
-
-        // Construct a linker for the target
-        let mut linker = linker::create_with_target(&target);
-        linker.add_object(obj_file.path())?;
-
-        let output_path = assembly_output_path(self.db, self.file_id, out_dir);
-
-        // Link the object
-        linker.build_shared_object(&output_path)?;
-        linker.finalize()?;
-
-        Ok(output_path)
+        ObjectFile::new(
+            &self.db.target(),
+            &self.target_machine,
+            self.db.file_relative_path(self.file_id),
+            self.assembly_module,
+        )
     }
 }
+
 /// Computes the output path for the assembly of the specified file.
-fn assembly_output_path<D: IrDatabase>(db: &D, file_id: FileId, out_dir: Option<&Path>) -> PathBuf {
-    let relative_path: RelativePathBuf = db.file_relative_path(file_id);
-    let original_filename = Path::new(relative_path.file_name().unwrap());
+fn assembly_output_path(src_path: &RelativePathBuf, out_dir: Option<&Path>) -> PathBuf {
+    let original_filename = Path::new(src_path.file_name().unwrap());
 
     // Add the `munlib` suffix to the original filename
     let output_file_name = original_filename.with_extension("munlib");

crates/mun_codegen/src/ir/body.rs

@@ -775,7 +775,7 @@ impl<'a, 'b, D: IrDatabase> BodyIrGenerator<'a, 'b, D> {
 
     /// Generates IR for a function call.
     fn gen_call(&mut self, function: hir::Function, args: &[BasicValueEnum]) -> CallSiteValue {
-        if self.should_use_dispatch_table() {
+        if self.dispatch_table.contains(function) && self.should_use_dispatch_table() {
             let ptr_value = self.dispatch_table.gen_function_lookup(
                 self.db,
                 self.external_globals.dispatch_table,

crates/mun_codegen/src/ir/dispatch_table.rs

@@ -60,6 +60,11 @@ pub struct DispatchableFunction {
 }
 
 impl DispatchTable {
+    /// Returns whether the `DispatchTable` contains the specified `function`.
+    pub fn contains(&self, function: hir::Function) -> bool {
+        self.function_to_idx.contains_key(&function)
+    }
+
     /// Returns a slice containing all the functions in the dispatch table.
     pub fn entries(&self) -> &[DispatchableFunction] {
         &self.entries
@@ -174,7 +179,7 @@ impl<'a, D: IrDatabase> DispatchTableBuilder<'a, D> {
     pub fn new(
         db: &'a D,
         module: &'a Module,
-        intrinsics: BTreeMap<FunctionPrototype, FunctionType>,
+        intrinsics: &BTreeMap<FunctionPrototype, FunctionType>,
     ) -> Self {
         let mut table = DispatchTableBuilder {
             db,
@@ -192,17 +197,17 @@ impl<'a, D: IrDatabase> DispatchTableBuilder<'a, D> {
             table.ensure_table_ref();
 
             // Use a `BTreeMap` to guarantee deterministically ordered output
-            for (prototype, ir_type) in intrinsics.into_iter() {
+            for (prototype, ir_type) in intrinsics.iter() {
                 let index = table.entries.len();
                 table.entries.push(TypedDispatchableFunction {
                     function: DispatchableFunction {
                         prototype: prototype.clone(),
                         hir: None,
                     },
-                    ir_type,
+                    ir_type: *ir_type,
                 });
 
-                table.prototype_to_idx.insert(prototype, index);
+                table.prototype_to_idx.insert(prototype.clone(), index);
             }
         }
         table

crates/mun_codegen/src/ir/file.rs

@@ -25,13 +25,15 @@ pub(crate) fn ir_query(db: &impl IrDatabase, file_id: FileId) -> Arc<FileIR> {
 
     let group_ir = db.group_ir(file_id);
 
-    // Generate all function and wrapper function signatures.
+    // Generate all exposed function and wrapper function signatures.
     // Use a `BTreeMap` to guarantee deterministically ordered output.ures
     let mut functions = HashMap::new();
     let mut wrapper_functions = BTreeMap::new();
     for def in db.module_data(file_id).definitions() {
         if let ModuleDef::Function(f) = def {
-            if !f.is_extern(db) {
+            if (!f.data(db).visibility().is_private() || group_ir.dispatch_table.contains(*f))
+                && !f.is_extern(db)
+            {
                 let fun = function::gen_signature(
                     db,
                     *f,

crates/mun_codegen/src/ir/file_group.rs

@@ -61,9 +61,28 @@ pub(crate) fn ir_query(db: &impl IrDatabase, file_id: hir::FileId) -> Arc<FileGr
         }
     }
 
+    // Collect all exposed functions' bodies.
+    let mut dispatch_table_builder = DispatchTableBuilder::new(db, &llvm_module, &intrinsics_map);
+    for def in db.module_data(file_id).definitions() {
+        if let ModuleDef::Function(f) = def {
+            if !f.data(db).visibility().is_private() && !f.is_extern(db) {
+                let body = f.body(db);
+                let infer = f.infer(db);
+                dispatch_table_builder.collect_body(&body, &infer);
+            }
+        }
+    }
+
+    let dispatch_table = dispatch_table_builder.build();
+
     let abi_types = gen_abi_types(&db.context());
-    let mut type_table_builder =
-        TypeTableBuilder::new(db, &llvm_module, &abi_types, intrinsics_map.keys());
+    let mut type_table_builder = TypeTableBuilder::new(
+        db,
+        &llvm_module,
+        &abi_types,
+        intrinsics_map.keys(),
+        &dispatch_table,
+    );
 
     // Collect all used types
     for def in db.module_data(file_id).definitions() {
@@ -80,20 +99,6 @@ pub(crate) fn ir_query(db: &impl IrDatabase, file_id: hir::FileId) -> Arc<FileGr
 
     let type_table = type_table_builder.build();
 
-    // Collect all function and wrapper function signatures.
-    let mut dispatch_table_builder = DispatchTableBuilder::new(db, &llvm_module, intrinsics_map);
-    for def in db.module_data(file_id).definitions() {
-        if let ModuleDef::Function(f) = def {
-            if !f.is_extern(db) {
-                let body = f.body(db);
-                let infer = f.infer(db);
-                dispatch_table_builder.collect_body(&body, &infer);
-            }
-        }
-    }
-
-    let dispatch_table = dispatch_table_builder.build();
-
     // Create the allocator handle global value
     let allocator_handle_type = if needs_alloc {
         let allocator_handle_type = db.context().i8_type().ptr_type(AddressSpace::Generic);

crates/mun_codegen/src/ir/type_table.rs

@@ -1,7 +1,10 @@
 use crate::code_gen::{
     gen_global, gen_string_array, gen_struct_ptr_array, gen_u16_array, intern_string,
 };
-use crate::ir::{abi_types::AbiTypes, dispatch_table::FunctionPrototype};
+use crate::ir::{
+    abi_types::AbiTypes,
+    dispatch_table::{DispatchTable, FunctionPrototype},
+};
 use crate::type_info::{TypeGroup, TypeInfo};
 use crate::IrDatabase;
 use hir::{Body, ExprId, InferenceResult};
@@ -83,6 +86,7 @@ pub(crate) struct TypeTableBuilder<'a, D: IrDatabase> {
     target_data: Arc<TargetData>,
     module: &'a Module,
     abi_types: &'a AbiTypes,
+    dispatch_table: &'a DispatchTable,
     entries: BTreeSet<TypeInfo>, // Use a `BTreeSet` to guarantee deterministically ordered output
 }
 
@@ -93,12 +97,14 @@ impl<'a, D: IrDatabase> TypeTableBuilder<'a, D> {
         module: &'a Module,
         abi_types: &'a AbiTypes,
         intrinsics: impl Iterator<Item = &'f FunctionPrototype>,
+        dispatch_table: &'a DispatchTable,
     ) -> Self {
         let mut builder = Self {
             db,
             target_data: db.target_data(),
             module,
             abi_types,
+            dispatch_table,
             entries: BTreeSet::new(),
         };
 
@@ -136,7 +142,7 @@ impl<'a, D: IrDatabase> TypeTableBuilder<'a, D> {
     /// Collects unique `TypeInfo` from the specified function signature and body.
     pub fn collect_fn(&mut self, hir_fn: hir::Function) {
         // Collect type info for exposed function
-        if !hir_fn.data(self.db).visibility().is_private() || hir_fn.is_extern(self.db) {
+        if !hir_fn.data(self.db).visibility().is_private() || self.dispatch_table.contains(hir_fn) {
             let fn_sig = hir_fn.ty(self.db).callable_sig(self.db).unwrap();
 
             // Collect argument types

crates/mun_codegen/src/snapshots/test__binary_expressions.snap

@@ -1,11 +1,15 @@
 ---
 source: crates/mun_codegen/src/test.rs
-expression: "fn add(a:int, b:int)->int {\n  a+b\n}\n\nfn subtract(a:int, b:int)->int {\n  a-b\n}\n\nfn multiply(a:int, b:int)->int {\n  a*b\n}"
+expression: "pub fn add(a:int, b:int) -> int {\n  a+b\n}\n\npub fn subtract(a:int, b:int) -> int {\n  a-b\n}\n\npub fn multiply(a:int, b:int) -> int {\n  a*b\n}"
 ---
 ; == FILE IR =====================================
 ; ModuleID = 'main.mun'
 source_filename = "main.mun"
 
+%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
+
+@global_type_table = external global [1 x %struct.MunTypeInfo addrspace(4)*]
+
 define i64 @add(i64, i64) {
 body:
   %add = add i64 %0, %1
@@ -29,3 +33,9 @@ body:
 ; ModuleID = 'group_name'
 source_filename = "group_name"
 
+%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
+
+@"type_info::<core::i64>::name" = private unnamed_addr constant [10 x i8] c"core::i64\00"
+@"type_info::<core::i64>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"G\13;t\97j8\18\D7M\83`\1D\C8\19%", [10 x i8]* @"type_info::<core::i64>::name", i32 64, i8 8, i8 0 }
+@global_type_table = global [1 x %struct.MunTypeInfo addrspace(4)*] [%struct.MunTypeInfo addrspace(4)* @"type_info::<core::i64>"]
+

crates/mun_codegen/src/snapshots/test__conditional_return_expr.snap

@@ -1,11 +1,15 @@
 ---
 source: crates/mun_codegen/src/test.rs
-expression: "fn main(a:int)->int {\n    if a > 4 {\n        return a;\n    }\n    a - 1\n}"
+expression: "pub fn main(a:int) -> int {\n    if a > 4 {\n        return a;\n    }\n    a - 1\n}"
 ---
 ; == FILE IR =====================================
 ; ModuleID = 'main.mun'
 source_filename = "main.mun"
 
+%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
+
+@global_type_table = external global [1 x %struct.MunTypeInfo addrspace(4)*]
+
 define i64 @main(i64) {
 body:
   %greater = icmp sgt i64 %0, 4
@@ -24,3 +28,9 @@ if_merge:                                         ; preds = %body
 ; ModuleID = 'group_name'
 source_filename = "group_name"
 
+%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
+
+@"type_info::<core::i64>::name" = private unnamed_addr constant [10 x i8] c"core::i64\00"
+@"type_info::<core::i64>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"G\13;t\97j8\18\D7M\83`\1D\C8\19%", [10 x i8]* @"type_info::<core::i64>::name", i32 64, i8 8, i8 0 }
+@global_type_table = global [1 x %struct.MunTypeInfo addrspace(4)*] [%struct.MunTypeInfo addrspace(4)* @"type_info::<core::i64>"]
+

crates/mun_codegen/src/snapshots/test__equality_operands.snap

@@ -1,11 +1,15 @@
 ---
 source: crates/mun_codegen/src/test.rs
-expression: "fn equals(a:int, b:int)->bool                { a == b }\nfn not_equals(a:int, b:int)->bool            { a != b }\nfn less(a:int, b:int)->bool                  { a < b }\nfn less_equal(a:int, b:int)->bool            { a <= b }\nfn greater(a:int, b:int)->bool               { a > b }\nfn greater_equal(a:int, b:int)->bool         { a >= b }\nfn equalsf(a:float, b:float)->bool           { a == b }\nfn not_equalsf(a:float, b:float)->bool       { a != b }\nfn lessf(a:float, b:float)->bool             { a < b }\nfn less_equalf(a:float, b:float)->bool       { a <= b }\nfn greaterf(a:float, b:float)->bool          { a > b }\nfn greater_equalf(a:float, b:float)->bool    { a >= b }"
+expression: "pub fn equals(a:int, b:int) -> bool                { a == b }\npub fn not_equals(a:int, b:int) -> bool            { a != b }\npub fn less(a:int, b:int) -> bool                  { a < b }\npub fn less_equal(a:int, b:int) -> bool            { a <= b }\npub fn greater(a:int, b:int) -> bool               { a > b }\npub fn greater_equal(a:int, b:int) -> bool         { a >= b }\npub fn equalsf(a:float, b:float) -> bool           { a == b }\npub fn not_equalsf(a:float, b:float) -> bool       { a != b }\npub fn lessf(a:float, b:float) -> bool             { a < b }\npub fn less_equalf(a:float, b:float) -> bool       { a <= b }\npub fn greaterf(a:float, b:float) -> bool          { a > b }\npub fn greater_equalf(a:float, b:float) -> bool    { a >= b }"
 ---
 ; == FILE IR =====================================
 ; ModuleID = 'main.mun'
 source_filename = "main.mun"
 
+%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
+
+@global_type_table = external global [3 x %struct.MunTypeInfo addrspace(4)*]
+
 define i1 @equals(i64, i64) {
 body:
   %eq = icmp eq i64 %0, %1
@@ -83,3 +87,13 @@ body:
 ; ModuleID = 'group_name'
 source_filename = "group_name"
 
+%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
+
+@"type_info::<core::i64>::name" = private unnamed_addr constant [10 x i8] c"core::i64\00"
+@"type_info::<core::i64>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"G\13;t\97j8\18\D7M\83`\1D\C8\19%", [10 x i8]* @"type_info::<core::i64>::name", i32 64, i8 8, i8 0 }
+@"type_info::<core::f64>::name" = private unnamed_addr constant [10 x i8] c"core::f64\00"
+@"type_info::<core::f64>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"`\DBF\9C?YJ%G\AD4\9F\D5\92%A", [10 x i8]* @"type_info::<core::f64>::name", i32 64, i8 8, i8 0 }
+@"type_info::<core::bool>::name" = private unnamed_addr constant [11 x i8] c"core::bool\00"
+@"type_info::<core::bool>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"x\82\81m t7\03\CB\F8k\81-;\C9\84", [11 x i8]* @"type_info::<core::bool>::name", i32 1, i8 1, i8 0 }
+@global_type_table = global [3 x %struct.MunTypeInfo addrspace(4)*] [%struct.MunTypeInfo addrspace(4)* @"type_info::<core::i64>", %struct.MunTypeInfo addrspace(4)* @"type_info::<core::f64>", %struct.MunTypeInfo addrspace(4)* @"type_info::<core::bool>"]
+