cleanup join points

compiler/gen_dev/src/generic64/mod.rs

@@ -289,7 +289,7 @@ pub struct Backend64Bit<
     free_map: MutMap<*const Stmt<'a>, Vec<'a, Symbol>>,
 
     literal_map: MutMap<Symbol, (*const Literal<'a>, *const Layout<'a>)>,
-    join_map: MutMap<JoinPointId, u64>,
+    join_map: MutMap<JoinPointId, Vec<'a, (u64, u64)>>,
 
     storage_manager: StorageManager<'a, GeneralReg, FloatReg, ASM, CC>,
 }
@@ -550,10 +550,6 @@ impl<
         default_branch: &(BranchInfo<'a>, &'a Stmt<'a>),
         ret_layout: &Layout<'a>,
     ) {
-        // Free everything to the stack to make sure they don't get messed with in the branch.
-        // TODO: look into a nicer solution.
-        self.storage_manager.free_all_to_stack(&mut self.buf);
-
         // Switches are a little complex due to keeping track of jumps.
         // In general I am trying to not have to loop over things multiple times or waste memory.
         // The basic plan is to make jumps to nowhere and then correct them once we know the correct address.
@@ -619,36 +615,34 @@ impl<
     ) {
         // Free everything to the stack to make sure they don't get messed with in the branch.
         // TODO: look into a nicer solution.
-        self.storage_manager.free_all_to_stack(&mut self.buf);
+        // self.storage_manager.free_all_to_stack(&mut self.buf);
 
         // Ensure all the joinpoint parameters have storage locations.
         // On jumps to the joinpoint, we will overwrite those locations as a way to "pass parameters" to the joinpoint.
         self.storage_manager
             .setup_joinpoint(&mut self.buf, id, parameters);
 
-        // Create jump to remaining.
+        self.join_map.insert(*id, bumpalo::vec![in self.env.arena]);
-        let jmp_location = self.buf.len();
+
-        let start_offset = ASM::jmp_imm32(&mut self.buf, 0x1234_5678);
+        // Build remainder of function first. It is what gets run and jumps to join.
+        // self.storage_manager = base_storage;
+        self.build_stmt(remainder, ret_layout);
+
+        let join_location = self.buf.len() as u64;
 
         // Build all statements in body.
-        let mut base_storage = self.storage_manager.clone();
-        self.join_map.insert(*id, self.buf.len() as u64);
         self.build_stmt(body, ret_layout);
-        base_storage.update_stack_size(self.storage_manager.stack_size());
-        base_storage.update_fn_call_stack_size(self.storage_manager.fn_call_stack_size());
 
-        // Overwrite the original jump with the correct offset.
+        // Overwrite the all jumps to the joinpoint with the correct offset.
         let mut tmp = bumpalo::vec![in self.env.arena];
-        self.update_jmp_imm32_offset(
+        for (jmp_location, start_offset) in self
-            &mut tmp,
+            .join_map
-            jmp_location as u64,
+            .remove(id)
-            start_offset as u64,
+            .unwrap_or_else(|| internal_error!("join point not defined"))
-            self.buf.len() as u64,
+        {
-        );
+            tmp.clear();
-
+            self.update_jmp_imm32_offset(&mut tmp, jmp_location, start_offset, join_location);
-        // Build remainder of function.
+        }
-        self.storage_manager = base_storage;
-        self.build_stmt(remainder, ret_layout)
     }
 
     fn build_jump(
@@ -664,15 +658,8 @@ impl<
         let jmp_location = self.buf.len();
         let start_offset = ASM::jmp_imm32(&mut self.buf, 0x1234_5678);
 
-        if let Some(offset) = self.join_map.get(id) {
+        if let Some(vec) = self.join_map.get_mut(id) {
-            let offset = *offset;
+            vec.push((jmp_location as u64, start_offset as u64))
-            let mut tmp = bumpalo::vec![in self.env.arena];
-            self.update_jmp_imm32_offset(
-                &mut tmp,
-                jmp_location as u64,
-                start_offset as u64,
-                offset,
-            );
         } else {
             internal_error!("Jump: unknown point specified to jump to: {:?}", id);
         }

compiler/gen_dev/src/generic64/storage.rs

@@ -1025,12 +1025,20 @@ impl<
                 todo!("joinpoints with borrowed parameters");
             }
             // Claim a location for every join point parameter to be loaded at.
+            // Put everything on the stack for simplicity.
             match layout {
-                single_register_integers!() => {
+                single_register_layouts!() => {
-                    self.claim_general_reg(buf, symbol);
+                    let base_offset = self.claim_stack_size(8);
-                }
+                    self.symbol_storage_map.insert(
-                single_register_floats!() => {
+                        *symbol,
-                    self.claim_float_reg(buf, symbol);
+                        Stack(ReferencedPrimitive {
+                            base_offset,
+                            size: 8,
+                            sign_extend: false,
+                        }),
+                    );
+                    self.allocation_map
+                        .insert(*symbol, Rc::new((base_offset, 8)));
                 }
                 _ => {
                     let stack_size = layout.stack_size(self.target_info);
@@ -1069,17 +1077,8 @@ impl<
                 continue;
             }
             match wanted_storage {
-                Reg(General(reg)) => {
+                Reg(_) => {
-                    // Ensure the reg is free, if not free it.
+                    internal_error!("Register storage is not allowed for jumping to joinpoint")
-                    self.ensure_reg_free(buf, General(*reg));
-                    // Copy the value over to the reg.
-                    self.load_to_specified_general_reg(buf, sym, *reg)
-                }
-                Reg(Float(reg)) => {
-                    // Ensure the reg is free, if not free it.
-                    self.ensure_reg_free(buf, Float(*reg));
-                    // Copy the value over to the reg.
-                    self.load_to_specified_float_reg(buf, sym, *reg)
                 }
                 Stack(ReferencedPrimitive { base_offset, .. } | Complex { base_offset, .. }) => {
                     // TODO: This might be better not to call.
@@ -1090,7 +1089,9 @@ impl<
                 }
                 NoData => {}
                 Stack(Primitive { .. }) => {
-                    internal_error!("Primitive stack storage is not allowed for jumping")
+                    internal_error!(
+                        "Primitive stack storage is not allowed for jumping to joinpoint"
+                    )
                 }
             }
         }

compiler/gen_dev/src/lib.rs

@@ -950,15 +950,16 @@ trait Backend<'a> {
                 parameters,
                 body: continuation,
                 remainder,
-                id,
+                id: JoinPointId(sym),
                 ..
             } => {
-                join_map.insert(*id, parameters);
+                self.set_last_seen(*sym, stmt);
+                join_map.insert(JoinPointId(*sym), parameters);
                 for param in *parameters {
                     self.set_last_seen(param.symbol, stmt);
                 }
-                self.scan_ast(continuation);
                 self.scan_ast(remainder);
+                self.scan_ast(continuation);
             }
             Stmt::Jump(JoinPointId(sym), symbols) => {
                 if let Some(parameters) = join_map.get(&JoinPointId(*sym)) {
@@ -967,7 +968,6 @@ trait Backend<'a> {
                         self.set_last_seen(param.symbol, stmt);
                     }
                 }
-                self.set_last_seen(*sym, stmt);
                 for sym in *symbols {
                     self.set_last_seen(*sym, stmt);
                 }