fix(codegen): only generate exposed functions

This also refactors codegen tests to always include symbol generation in
snapshot tests.

fixes issue #133

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__extern_fn.snap

@@ -1,6 +1,6 @@
 ---
 source: crates/mun_codegen/src/test.rs
-expression: "extern fn add(a:int, b:int) -> int;\nfn main() {\n    add(3,4);\n}"
+expression: "extern fn add(a:int, b:int) -> int;\npub fn main() {\n    add(3,4);\n}"
 ---
 ; == FILE IR =====================================
 ; ModuleID = 'main.mun'
@@ -24,11 +24,11 @@ body:
 ; ModuleID = 'group_name'
 source_filename = "group_name"
 
-%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
 %DispatchTable = type { i64 (i64, i64)* }
+%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
 
+@dispatchTable = global %DispatchTable zeroinitializer
 @"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>"]
-@dispatchTable = global %DispatchTable zeroinitializer
 

crates/mun_codegen/src/snapshots/test__fibonacci.snap

@@ -37,13 +37,13 @@ if_merge:                                         ; preds = %body, %else
 ; ModuleID = 'group_name'
 source_filename = "group_name"
 
-%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
 %DispatchTable = type { i64 (i64)* }
+%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
 
+@dispatchTable = global %DispatchTable { i64 (i64)* @fibonacci }
 @"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>"]
-@dispatchTable = global %DispatchTable { i64 (i64)* @fibonacci }
 
 declare i64 @fibonacci(i64)
 

crates/mun_codegen/src/snapshots/test__field_crash.snap

@@ -1,6 +1,6 @@
 ---
 source: crates/mun_codegen/src/test.rs
-expression: "struct(gc) Foo { a: int };\n\nfn main(c:int)->int {\n    let b = Foo { a: c + 5 }\n    b.a\n}"
+expression: "struct(gc) Foo { a: int };\n\npub fn main(c:int) -> int {\n    let b = Foo { a: c + 5 }\n    b.a\n}"
 ---
 ; == FILE IR =====================================
 ; ModuleID = 'main.mun'
@@ -43,10 +43,11 @@ body:
 ; ModuleID = 'group_name'
 source_filename = "group_name"
 
+%DispatchTable = type { i8* addrspace(4)* (i8 addrspace(4)*, i8*)* }
 %struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
 %struct.MunStructInfo = type { i8 addrspace(4)* addrspace(4)*, %struct.MunTypeInfo addrspace(4)* addrspace(4)*, i16 addrspace(4)*, i16, i8 }
-%DispatchTable = type { i8* addrspace(4)* (i8 addrspace(4)*, i8*)* }
 
+@dispatchTable = global %DispatchTable zeroinitializer
 @"type_info::<Foo>::name" = private unnamed_addr constant [4 x i8] c"Foo\00"
 @"struct_info::<Foo>::field_names" = private unnamed_addr constant [2 x i8] c"a\00"
 @0 = private unnamed_addr constant [1 x i8 addrspace(4)*] [i8 addrspace(4)* @"struct_info::<Foo>::field_names"]
@@ -62,6 +63,5 @@ source_filename = "group_name"
 @"type_info::<*mut core::void>::name" = private unnamed_addr constant [16 x i8] c"*mut core::void\00"
 @"type_info::<*mut core::void>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"\F0Y\22\FC\95\9E\7F\CE\08T\B1\A2\CD\A7\FAz", [16 x i8]* @"type_info::<*mut core::void>::name", i32 64, i8 8, i8 0 }
 @global_type_table = global [5 x %struct.MunTypeInfo addrspace(4)*] [%struct.MunTypeInfo addrspace(4)* @"type_info::<Foo>", %struct.MunTypeInfo addrspace(4)* @"type_info::<*const TypeInfo>", %struct.MunTypeInfo addrspace(4)* @"type_info::<core::i64>", %struct.MunTypeInfo addrspace(4)* @"type_info::<*const *mut core::void>", %struct.MunTypeInfo addrspace(4)* @"type_info::<*mut core::void>"]
-@dispatchTable = global %DispatchTable zeroinitializer
 @allocatorHandle = unnamed_addr global i8* null
 

crates/mun_codegen/src/snapshots/test__field_expr.snap

@@ -57,12 +57,13 @@ body:
 ; ModuleID = 'group_name'
 source_filename = "group_name"
 
-%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
-%struct.MunStructInfo = type { i8 addrspace(4)* addrspace(4)*, %struct.MunTypeInfo addrspace(4)* addrspace(4)*, i16 addrspace(4)*, i16, i8 }
 %DispatchTable = type { i8* addrspace(4)* (i8 addrspace(4)*, i8*)*, i64 (%Foo)*, %Foo (%Bar)* }
 %Foo = type { i64 }
 %Bar = type { double, %Foo }
+%struct.MunTypeInfo = type { [16 x i8], i8 addrspace(4)*, i32, i8, i8 }
+%struct.MunStructInfo = type { i8 addrspace(4)* addrspace(4)*, %struct.MunTypeInfo addrspace(4)* addrspace(4)*, i16 addrspace(4)*, i16, i8 }
 
+@dispatchTable = global %DispatchTable { i8* addrspace(4)* (i8 addrspace(4)*, i8*)* null, i64 (%Foo)* @foo_a, %Foo (%Bar)* @bar_1 }
 @"type_info::<Foo>::name" = private unnamed_addr constant [4 x i8] c"Foo\00"
 @"struct_info::<Foo>::field_names" = private unnamed_addr constant [2 x i8] c"a\00"
 @0 = private unnamed_addr constant [1 x i8 addrspace(4)*] [i8 addrspace(4)* @"struct_info::<Foo>::field_names"]
@@ -87,7 +88,6 @@ source_filename = "group_name"
 @"type_info::<*mut core::void>::name" = private unnamed_addr constant [16 x i8] c"*mut core::void\00"
 @"type_info::<*mut core::void>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"\F0Y\22\FC\95\9E\7F\CE\08T\B1\A2\CD\A7\FAz", [16 x i8]* @"type_info::<*mut core::void>::name", i32 64, i8 8, i8 0 }
 @global_type_table = global [7 x %struct.MunTypeInfo addrspace(4)*] [%struct.MunTypeInfo addrspace(4)* @"type_info::<Foo>", %struct.MunTypeInfo addrspace(4)* @"type_info::<*const TypeInfo>", %struct.MunTypeInfo addrspace(4)* @"type_info::<core::i64>", %struct.MunTypeInfo addrspace(4)* @"type_info::<core::f64>", %struct.MunTypeInfo addrspace(4)* @"type_info::<*const *mut core::void>", %struct.MunTypeInfo addrspace(4)* @"type_info::<Bar>", %struct.MunTypeInfo addrspace(4)* @"type_info::<*mut core::void>"]
-@dispatchTable = global %DispatchTable { i8* addrspace(4)* (i8 addrspace(4)*, i8*)* null, i64 (%Foo)* @foo_a, %Foo (%Bar)* @bar_1 }
 @allocatorHandle = unnamed_addr global i8* null
 
 declare i64 @foo_a(%Foo)