fix(codegen): crash on missing nested private function

crates/mun_codegen/src/ir/dispatch_table.rs

@@ -237,7 +237,7 @@ impl<'a, D: IrDatabase> DispatchTableBuilder<'a, D> {
         }
 
         // Recurse further
-        expr.walk_child_exprs(|expr_id| self.collect_expr(expr_id, body, infer))
+        expr.walk_child_exprs(|expr_id| self.collect_expr(expr_id, body, infer));
     }
 
     /// Collects function call expression from the given expression.
@@ -285,6 +285,14 @@ impl<'a, D: IrDatabase> DispatchTableBuilder<'a, D> {
             });
             self.prototype_to_idx.insert(prototype, index);
             self.function_to_idx.insert(function, index);
+
+            // Recurse further
+            let fn_body = function.body(self.db);
+            self.collect_expr(
+                fn_body.body_expr(),
+                &fn_body,
+                function.infer(self.db).as_ref(),
+            );
         }
     }
 

crates/mun_codegen/src/snapshots/test__nested_private_extern_fn.snap

@@ -0,0 +1,43 @@
+---
+source: crates/mun_codegen/src/test.rs
+expression: "extern fn extern_fn() -> f32;\n\nfn private_fn() -> f32 {\n    extern_fn()\n}\n\npub fn main() -> f32 {\n    private_fn()\n}"
+---
+; == FILE IR =====================================
+; ModuleID = 'main.mun'
+source_filename = "main.mun"
+
+%DispatchTable = type { float ()*, float ()* }
+%struct.MunTypeInfo = type { [16 x i8], i8*, i32, i8, i8 }
+
+@dispatchTable = external global %DispatchTable
+@global_type_table = external global [1 x %struct.MunTypeInfo*]
+
+define float @private_fn() {
+body:
+  %extern_fn_ptr = load float ()*, float ()** getelementptr inbounds (%DispatchTable, %DispatchTable* @dispatchTable, i32 0, i32 1)
+  %extern_fn = call float %extern_fn_ptr()
+  ret float %extern_fn
+}
+
+define float @main() {
+body:
+  %private_fn_ptr = load float ()*, float ()** getelementptr inbounds (%DispatchTable, %DispatchTable* @dispatchTable, i32 0, i32 0)
+  %private_fn = call float %private_fn_ptr()
+  ret float %private_fn
+}
+
+
+; == GROUP IR ====================================
+; ModuleID = 'group_name'
+source_filename = "group_name"
+
+%DispatchTable = type { float ()*, float ()* }
+%struct.MunTypeInfo = type { [16 x i8], i8*, i32, i8, i8 }
+
+@dispatchTable = global %DispatchTable { float ()* @private_fn, float ()* null }
+@"type_info::<core::f32>::name" = private unnamed_addr constant [10 x i8] c"core::f32\00"
+@"type_info::<core::f32>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"P\19b7\A8k\F2\81P\FB\83\F5P\B0\82!", [10 x i8]* @"type_info::<core::f32>::name", i32 32, i8 4, i8 0 }
+@global_type_table = constant [1 x %struct.MunTypeInfo*] [%struct.MunTypeInfo* @"type_info::<core::f32>"]
+
+declare float @private_fn()
+

crates/mun_codegen/src/snapshots/test__nested_private_fn.snap

@@ -6,7 +6,7 @@ expression: "fn nested_private_fn() -> i32 {\n    1\n}\n\nfn private_fn() -> i32
 ; ModuleID = 'main.mun'
 source_filename = "main.mun"
 
-%DispatchTable = type { i32 ()* }
+%DispatchTable = type { i32 ()*, i32 ()* }
 %struct.MunTypeInfo = type { [16 x i8], i8*, i32, i8, i8 }
 
 @dispatchTable = external global %DispatchTable
@@ -19,7 +19,8 @@ body:
 
 define i32 @private_fn() {
 body:
-  %nested_private_fn = call i32 @nested_private_fn()
+  %nested_private_fn_ptr = load i32 ()*, i32 ()** getelementptr inbounds (%DispatchTable, %DispatchTable* @dispatchTable, i32 0, i32 1)
+  %nested_private_fn = call i32 %nested_private_fn_ptr()
   ret i32 %nested_private_fn
 }
 
@@ -35,13 +36,15 @@ body:
 ; ModuleID = 'group_name'
 source_filename = "group_name"
 
-%DispatchTable = type { i32 ()* }
+%DispatchTable = type { i32 ()*, i32 ()* }
 %struct.MunTypeInfo = type { [16 x i8], i8*, i32, i8, i8 }
 
-@dispatchTable = global %DispatchTable { i32 ()* @private_fn }
+@dispatchTable = global %DispatchTable { i32 ()* @private_fn, i32 ()* @nested_private_fn }
 @"type_info::<core::i32>::name" = private unnamed_addr constant [10 x i8] c"core::i32\00"
 @"type_info::<core::i32>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"\17yzt\19\D62\17\D25\95C\17\88[\FA", [10 x i8]* @"type_info::<core::i32>::name", i32 32, i8 4, i8 0 }
 @global_type_table = constant [1 x %struct.MunTypeInfo*] [%struct.MunTypeInfo* @"type_info::<core::i32>"]
 
 declare i32 @private_fn()
 
+declare i32 @nested_private_fn()
+

crates/mun_codegen/src/snapshots/test__nested_private_recursive_fn.snap

@@ -0,0 +1,43 @@
+---
+source: crates/mun_codegen/src/test.rs
+expression: "fn private_fn() -> f32 {\n    private_fn()\n}\n\npub fn main() -> f32 {\n    private_fn()\n}"
+---
+; == FILE IR =====================================
+; ModuleID = 'main.mun'
+source_filename = "main.mun"
+
+%DispatchTable = type { float ()* }
+%struct.MunTypeInfo = type { [16 x i8], i8*, i32, i8, i8 }
+
+@dispatchTable = external global %DispatchTable
+@global_type_table = external global [1 x %struct.MunTypeInfo*]
+
+define float @private_fn() {
+body:
+  %private_fn_ptr = load float ()*, float ()** getelementptr inbounds (%DispatchTable, %DispatchTable* @dispatchTable, i32 0, i32 0)
+  %private_fn = call float %private_fn_ptr()
+  ret float %private_fn
+}
+
+define float @main() {
+body:
+  %private_fn_ptr = load float ()*, float ()** getelementptr inbounds (%DispatchTable, %DispatchTable* @dispatchTable, i32 0, i32 0)
+  %private_fn = call float %private_fn_ptr()
+  ret float %private_fn
+}
+
+
+; == GROUP IR ====================================
+; ModuleID = 'group_name'
+source_filename = "group_name"
+
+%DispatchTable = type { float ()* }
+%struct.MunTypeInfo = type { [16 x i8], i8*, i32, i8, i8 }
+
+@dispatchTable = global %DispatchTable { float ()* @private_fn }
+@"type_info::<core::f32>::name" = private unnamed_addr constant [10 x i8] c"core::f32\00"
+@"type_info::<core::f32>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"P\19b7\A8k\F2\81P\FB\83\F5P\B0\82!", [10 x i8]* @"type_info::<core::f32>::name", i32 32, i8 4, i8 0 }
+@global_type_table = constant [1 x %struct.MunTypeInfo*] [%struct.MunTypeInfo* @"type_info::<core::f32>"]
+
+declare float @private_fn()
+

crates/mun_codegen/src/snapshots/test__nested_private_recursive_fn_with_args.snap

@@ -0,0 +1,47 @@
+---
+source: crates/mun_codegen/src/test.rs
+expression: "extern fn other() -> i32;\n\nfn private_fn(a: i32) -> f32 {\n    private_fn(a)\n}\n\npub fn main() -> f32 {\n    private_fn(other())\n}"
+---
+; == FILE IR =====================================
+; ModuleID = 'main.mun'
+source_filename = "main.mun"
+
+%DispatchTable = type { float (i32)*, i32 ()* }
+%struct.MunTypeInfo = type { [16 x i8], i8*, i32, i8, i8 }
+
+@dispatchTable = external global %DispatchTable
+@global_type_table = external global [2 x %struct.MunTypeInfo*]
+
+define float @private_fn(i32) {
+body:
+  %private_fn_ptr = load float (i32)*, float (i32)** getelementptr inbounds (%DispatchTable, %DispatchTable* @dispatchTable, i32 0, i32 0)
+  %private_fn = call float %private_fn_ptr(i32 %0)
+  ret float %private_fn
+}
+
+define float @main() {
+body:
+  %other_ptr = load i32 ()*, i32 ()** getelementptr inbounds (%DispatchTable, %DispatchTable* @dispatchTable, i32 0, i32 1)
+  %other = call i32 %other_ptr()
+  %private_fn_ptr = load float (i32)*, float (i32)** getelementptr inbounds (%DispatchTable, %DispatchTable* @dispatchTable, i32 0, i32 0)
+  %private_fn = call float %private_fn_ptr(i32 %other)
+  ret float %private_fn
+}
+
+
+; == GROUP IR ====================================
+; ModuleID = 'group_name'
+source_filename = "group_name"
+
+%DispatchTable = type { float (i32)*, i32 ()* }
+%struct.MunTypeInfo = type { [16 x i8], i8*, i32, i8, i8 }
+
+@dispatchTable = global %DispatchTable { float (i32)* @private_fn, i32 ()* null }
+@"type_info::<core::i32>::name" = private unnamed_addr constant [10 x i8] c"core::i32\00"
+@"type_info::<core::i32>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"\17yzt\19\D62\17\D25\95C\17\88[\FA", [10 x i8]* @"type_info::<core::i32>::name", i32 32, i8 4, i8 0 }
+@"type_info::<core::f32>::name" = private unnamed_addr constant [10 x i8] c"core::f32\00"
+@"type_info::<core::f32>" = private unnamed_addr constant %struct.MunTypeInfo { [16 x i8] c"P\19b7\A8k\F2\81P\FB\83\F5P\B0\82!", [10 x i8]* @"type_info::<core::f32>::name", i32 32, i8 4, i8 0 }
+@global_type_table = constant [2 x %struct.MunTypeInfo*] [%struct.MunTypeInfo* @"type_info::<core::i32>", %struct.MunTypeInfo* @"type_info::<core::f32>"]
+
+declare float @private_fn(i32)
+

crates/mun_codegen/src/test.rs

@@ -824,6 +824,55 @@ fn nested_private_fn() {
     );
 }
 
+#[test]
+fn nested_private_extern_fn() {
+    test_snapshot(
+        r#"
+    extern fn extern_fn() -> f32;
+
+    fn private_fn() -> f32 {
+        extern_fn()
+    }
+
+    pub fn main() -> f32 {
+        private_fn()
+    }
+    "#,
+    )
+}
+
+#[test]
+fn nested_private_recursive_fn() {
+    test_snapshot(
+        r#"
+    fn private_fn() -> f32 {
+        private_fn()
+    }
+
+    pub fn main() -> f32 {
+        private_fn()
+    }
+    "#,
+    )
+}
+
+#[test]
+fn nested_private_recursive_fn_with_args() {
+    test_snapshot(
+        r#"
+    extern fn other() -> i32;
+
+    fn private_fn(a: i32) -> f32 {
+        private_fn(a)
+    }
+
+    pub fn main() -> f32 {
+        private_fn(other())
+    }
+    "#,
+    )
+}
+
 fn test_snapshot(text: &str) {
     test_snapshot_with_optimization(text, OptimizationLevel::Default);
 }