From 5832cb758964402e8bb3b60b85e2cc36b368e2e7 Mon Sep 17 00:00:00 2001
From: oflatt <oflatt@gmail.com>
Date: Wed, 30 Oct 2024 14:31:42 -0700
Subject: [PATCH] fix bug with return types

---
 bril-rs/brillvm/src/llvm.rs | 37 ++++++++++++++++++++++++++++++-------
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/bril-rs/brillvm/src/llvm.rs b/bril-rs/brillvm/src/llvm.rs
index df637b24..79ebd773 100644
--- a/bril-rs/brillvm/src/llvm.rs
+++ b/bril-rs/brillvm/src/llvm.rs
@@ -96,10 +96,19 @@ impl<'a, 'b> Heap<'a, 'b> {
         name: &'b String,
         ty: &Type,
     ) -> WrappedPointer<'a> {
-        self.map
+        let result = self
+            .map
             .entry(name)
             .or_insert_with(|| WrappedPointer::new(builder, context, name, ty))
-            .clone()
+            .clone();
+        if result.ty != *ty {
+            println!(
+                "`{}` had type `{}` but is now being assigned type `{}`",
+                name, result.ty, ty
+            );
+            unimplemented!("brillvm does not currently support variables within a function having different types. Implementing this might require a control flow analysis? Feel free to try and implement this.")
+        }
+        result
     }
 
     fn get(&self, name: &String) -> WrappedPointer<'a> {
@@ -1483,7 +1492,7 @@ pub fn create_module_from_program<'a>(
                     }
                 });
 
-                (llvm_func, instrs, block, heap)
+                (llvm_func, instrs, block, heap, return_type)
             },
         )
         .collect(); // Important to collect, can't be done lazily because we need all functions to be loaded in before a call instruction of a function is processed.
@@ -1493,9 +1502,13 @@ pub fn create_module_from_program<'a>(
     let mut ticks_start_ref = None;
     funcs
         .into_iter()
-        .for_each(|(llvm_func, instrs, mut block, heap)| {
+        .for_each(|(llvm_func, instrs, mut block, heap, return_type)| {
             let mut last_instr = None;
 
+            // Maps labels to llvm blocks for jumps
+            let mut block_map = HashMap::new();
+
+
             // If there are actually instructions, proceed
             if !instrs.is_empty() {
                 builder.position_at_end(block);
@@ -1525,8 +1538,6 @@ pub fn create_module_from_program<'a>(
                     func.add_attribute(AttributeLoc::Function, context.create_enum_attribute(3, 1));*/
                 }
 
-                // Maps labels to llvm blocks for jumps
-                let mut block_map = HashMap::new();
                 let mut index = 0;
                 while index < instrs.len() {
                     // for main, we expect the last instruction to be a print
@@ -1687,7 +1698,19 @@ pub fn create_module_from_program<'a>(
 
             // Make sure every function is terminated with a return if not already
             if !is_terminating_instr(&last_instr) {
-                builder.build_return(None).unwrap();
+                if return_type.is_none() {
+                    builder.build_return(None).unwrap();
+                } else {
+                    // This block did not have a terminating instruction
+                    // Returning void is ill-typed for this function
+                    // This code should be unreachable in well-formed Bril
+                    // Let's just arbitrarily jump to avoid needing to
+                    // instantiate a valid return value.
+                    assert!(!block_map.is_empty());
+                    builder
+                        .build_unconditional_branch(*block_map.values().next().unwrap())
+                        .unwrap();
+                }
             }
         });