WIP

src/ion/fast_alloc2.rs

@@ -940,7 +940,7 @@ impl<'a, F: Function> FastAlloc<'a, F> {
                                         Some(slot) => slot,
                                         None => self.alloc_stack_slot(op.vreg()),
                                     };
-                                    self.move_to_stack(op.vreg(), ProgPoint::before(inst));
+                                    self.move_to_stack(op.vreg().vreg(), ProgPoint::before(inst));
                                     trace!(" -> Allocated op {} to slot {}", op_idx, slot);
                                 } else {
                                     self.allocs[alloc_idx + op_idx] = Allocation::reg(preg);
@@ -1406,7 +1406,197 @@ impl<'a, F: Function> FastAlloc<'a, F> {
         block: Block,
         succ: Block,
     ) -> Result<(), RegAllocError> {
-        todo!()
+        trace!("  -> Block has a single allocated edge. Moving");
+        let succ_params = self.func.block_params(succ);
+        let last_inst = self.func.block_insns(block).last();
+        debug_assert!(self.func.is_branch(last_inst));
+        let out_vregs = self.func.branch_blockparams(block, last_inst, 0);
+        debug_assert_eq!(succ_params.len(), out_vregs.len());
+
+        // first handle all the stack moves
+        // then registers
+        let succ_liveins = &self.liveins[succ.index()];
+        let succ_livein_loc_idx = self.livein_loc_lookup[succ.index()] as usize;
+
+        // live vregs that need to be in certain registers
+        let mut live_vregs_in_pregs = Vec::new();
+
+        let mut loc_idx = succ_livein_loc_idx;
+        for live_vreg in succ_liveins.iter() {
+            let loc = self.livein_locs[loc_idx];
+            match loc.slot() {
+                Some(idx) => {
+                    if self.vregs[live_vreg].stack_slot().is_none() {
+                        // inlined because i dont have the nerve to deal with the borrow checker
+                        // self.move_to_stack(live_vreg, ProgPoint::before(last_inst));
+
+                        // must exist
+                        let slot = self.vregs[live_vreg].allocated_stack_slot().unwrap();
+
+                        let preg = if let Some(preg) = self.vregs[live_vreg].preg() {
+                            preg
+                        } else {
+                            panic!(
+                                "Trying to move {} to stack but it is not in a preg",
+                                live_vreg
+                            );
+                        };
+
+                        trace!("Moving {} from {} to slot {}", live_vreg, preg, slot);
+                        self.edits.push((
+                            ProgPoint::before(last_inst),
+                            Edit::Move {
+                                from: Allocation::reg(preg),
+                                to: Allocation::stack(SpillSlot::new(slot as usize)),
+                            },
+                        ));
+                        self.vregs[live_vreg].set_stack_slot_valid();
+                    } else {
+                        debug_assert_eq!(self.vregs[live_vreg].stack_slot().unwrap(), idx);
+                    }
+                }
+                None => {
+                    // don't need to reset stack alloc since that is handled when allocating the block
+                }
+            }
+            if let Some(preg) = loc.preg() {
+                live_vregs_in_pregs.push((live_vreg, preg));
+            } else {
+                self.vregs[live_vreg].cur_use_idx -= 1;
+            }
+            loc_idx += 1;
+        }
+
+        for i in 0..out_vregs.len() {
+            let out_vreg = out_vregs[i];
+            let succ_vreg = succ_params[i];
+            let loc = self.livein_locs[loc_idx];
+            match loc.slot() {
+                Some(slot) => {
+                    if self.vregs[out_vreg.vreg()].preg().is_none() {
+                        // TODO: we could also use a single temp register
+                        let preg = self.find_free_reg(
+                            out_vreg.class(),
+                            PRegSet::empty(),
+                            Some(ProgPoint::before(last_inst)),
+                            block,
+                            self.cur_inst_pos + 1,
+                        )?;
+                        self.move_vreg_to_preg(
+                            out_vreg,
+                            preg,
+                            ProgPoint::before(last_inst),
+                            block,
+                            self.cur_inst_pos + 1,
+                        );
+                    }
+
+                    let preg = self.vregs[out_vreg.vreg()].preg().unwrap();
+                    // since this is the stack slot of
+                    trace!("Moving {} from {} to slot {}", out_vreg, preg, slot);
+                    self.edits.push((
+                        ProgPoint::before(last_inst),
+                        Edit::Move {
+                            from: Allocation::reg(preg),
+                            to: Allocation::stack(SpillSlot::new(slot as usize)),
+                        },
+                    ));
+                    self.vregs[succ_vreg.vreg()].set_stack_slot_valid();
+                }
+                None => {}
+            }
+            if loc.preg().is_none() {
+                self.vregs[out_vreg.vreg()].cur_use_idx -= 1;
+            }
+            loc_idx += 1;
+        }
+
+        // move to pregs, note: this can unnecessarily spill
+        for &(live_vreg, target_preg) in &live_vregs_in_pregs {
+            self.vregs[live_vreg].cur_use_idx -= 1;
+            if let Some(cur_preg) = self.vregs[live_vreg].preg() {
+                if cur_preg == target_preg {
+                    continue;
+                }
+            }
+
+            if let Some(cur_vreg) = self.pregs[target_preg.index()].vreg() {
+                if !self.vreg_killed(cur_vreg.vreg(), block, self.cur_inst_pos + 1, true) {
+                    if self.vregs[cur_vreg.vreg()].stack_slot().is_none() {
+                        self.alloc_and_move_to_stack(cur_vreg, ProgPoint::before(last_inst));
+                    }
+                    self.clear_preg(target_preg.index());
+                }
+            }
+
+            if let Some(cur_preg) = self.vregs[live_vreg].preg() {
+                self.vregs[live_vreg].set_preg(target_preg);
+                self.pregs[cur_preg.index()].clear_vreg();
+                self.pregs[target_preg.index()].set_vreg(VReg::new(live_vreg, target_preg.class()));
+                self.edits.push((
+                    ProgPoint::before(last_inst),
+                    Edit::Move {
+                        from: Allocation::reg(cur_preg),
+                        to: Allocation::reg(target_preg),
+                    },
+                ));
+            } else {
+                self.move_vreg_to_preg(
+                    VReg::new(live_vreg, target_preg.class()),
+                    target_preg,
+                    ProgPoint::before(last_inst),
+                    block,
+                    self.cur_inst_pos + 1,
+                );
+            }
+        }
+
+        loc_idx -= out_vregs.len();
+        for &out_vreg in out_vregs {
+            let loc = self.livein_locs[loc_idx];
+            if let Some(target_preg) = loc.preg() {
+                self.vregs[out_vreg.vreg()].cur_use_idx -= 1;
+                if let Some(cur_preg) = self.vregs[out_vreg.vreg()].preg() {
+                    if cur_preg == target_preg {
+                        loc_idx += 1;
+                        continue;
+                    }
+                }
+
+                if let Some(cur_vreg) = self.pregs[target_preg.index()].vreg() {
+                    if !self.vreg_killed(cur_vreg.vreg(), block, self.cur_inst_pos + 1, true) {
+                        if self.vregs[cur_vreg.vreg()].stack_slot().is_none() {
+                            self.alloc_and_move_to_stack(cur_vreg, ProgPoint::before(last_inst));
+                        }
+                        self.clear_preg(target_preg.index());
+                    }
+                }
+
+                if let Some(cur_preg) = self.vregs[out_vreg.vreg()].preg() {
+                    self.vregs[out_vreg.vreg()].set_preg(target_preg);
+                    self.pregs[cur_preg.index()].clear_vreg();
+                    self.pregs[target_preg.index()].set_vreg(out_vreg);
+                    self.edits.push((
+                        ProgPoint::before(last_inst),
+                        Edit::Move {
+                            from: Allocation::reg(cur_preg),
+                            to: Allocation::reg(target_preg),
+                        },
+                    ));
+                } else {
+                    self.move_vreg_to_preg(
+                        out_vreg,
+                        target_preg,
+                        ProgPoint::before(last_inst),
+                        block,
+                        self.cur_inst_pos + 1,
+                    );
+                }
+            }
+            loc_idx += 1;
+        }
+
+        return Ok(());
     }
 
     // TODO: check if the reftype count is correctly managed
@@ -1797,19 +1987,19 @@ impl<'a, F: Function> FastAlloc<'a, F> {
     // panics if there is a stack slot already or the vreg is not in a preg
     fn alloc_and_move_to_stack(&mut self, vreg: VReg, pos: ProgPoint) {
         self.alloc_stack_slot(vreg);
-        self.move_to_stack(vreg, pos);
+        self.move_to_stack(vreg.vreg(), pos);
     }
 
     // Moves a vreg to stack
     // panics if there is no stack slot or the vreg is not in a preg
-    fn move_to_stack(&mut self, vreg: VReg, pos: ProgPoint) {
+    fn move_to_stack(&mut self, vreg: usize, pos: ProgPoint) {
-        let slot = if let Some(slot) = self.vregs[vreg.vreg()].allocated_stack_slot() {
+        let slot = if let Some(slot) = self.vregs[vreg].allocated_stack_slot() {
             slot
         } else {
             panic!("Trying to move {} to stack but it has no slot", vreg);
         };
 
-        let preg = if let Some(preg) = self.vregs[vreg.vreg()].preg() {
+        let preg = if let Some(preg) = self.vregs[vreg].preg() {
             preg
         } else {
             panic!("Trying to move {} to stack but it is not in a preg", vreg);
@@ -1823,7 +2013,7 @@ impl<'a, F: Function> FastAlloc<'a, F> {
                 to: Allocation::stack(SpillSlot::new(slot as usize)),
             },
         ));
-        self.vregs[vreg.vreg()].set_stack_slot_valid();
+        self.vregs[vreg].set_stack_slot_valid();
     }
 
     // Allocates a stack slot for a vreg