fix stackmaps

src/ion/fast_alloc.rs

@@ -9,14 +9,16 @@ use crate::{
     OperandKind, OperandPos, Output, PReg, PRegSet, ProgPoint, RegAllocError, RegClass, SpillSlot,
     VReg,
 };
-use crate::{postorder, InstPosition};
+use crate::{domtree, postorder, InstPosition};
 
+use super::data_structures::u64_key;
 use super::Stats;
 
 #[derive(Default, Clone, Copy)]
 struct VRegData {
     pub preg: Option<PReg>,
     pub slot_idx: Option<u32>,
+    pub def_block: Option<Block>,
     pub reftype: bool,
 }
 
@@ -34,6 +36,7 @@ struct BlockData {
 
 struct ReadOnlyData {
     pub postorder: Vec<Block>,
+    pub idom: Vec<Block>,
     pub reg_order_int: Vec<PReg>,
     pub reg_order_float: Vec<PReg>,
 }
@@ -60,12 +63,22 @@ impl ReadOnlyData {
             reg_order
         };
 
+        let postorder = postorder::calculate(func.num_blocks(), func.entry_block(), |b| {
+            func.block_succs(b)
+        });
+
+        let idom = domtree::calculate(
+            func.num_blocks(),
+            |b| func.block_preds(b),
+            &postorder,
+            func.entry_block(),
+        );
+
         Self {
             reg_order_int,
             reg_order_float,
-            postorder: postorder::calculate(func.num_blocks(), func.entry_block(), |b| {
-                func.block_succs(b)
-            }),
+            postorder,
+            idom,
         }
     }
 
@@ -402,14 +415,26 @@ pub fn run<F: Function>(func: &F, mach_env: &MachineEnv) -> Result<Output, RegAl
     // we do not iterate the blocks in their index order so the order of edits might not be sorted by progpoint
     // however it should be nearly sorted
     state.edits.sort_by_key(|entry| entry.0);
+    // these might also not be sorted
+    //state.safepoint_slots.sort_by_key(|entry| entry.0);
+    state
+        .safepoint_slots
+        .sort_unstable_by_key(|(progpoint, slot)| u64_key(progpoint.to_index(), slot.bits()));
 
-    trace!("Edits:");
-    for edit in &state.edits {
-        match edit.1 {
-            Edit::Move { from, to } => {
-                trace!("At {:?} from {} to {}", edit.0, from, to);
+    if trace_enabled!() {
+        trace!("Edits:");
+        for edit in &state.edits {
+            match edit.1 {
+                Edit::Move { from, to } => {
+                    trace!("At {:?} from {} to {}", edit.0, from, to);
+                }
             }
         }
+
+        trace!("StackMaps:");
+        for entry in &state.safepoint_slots {
+            trace!("At {:?} at {}", entry.0, entry.1);
+        }
     }
 
     Ok(Output {
@@ -469,14 +494,19 @@ fn allocate_block_insts<'a, F: Function>(
         if req_refs_on_stack {
             state.clear_reftype_vregs();
             let pos = ProgPoint::before(inst);
+            trace!("Calculating Stackmap for {:?}", pos);
 
             for vreg in state.reftype_vregs {
                 let data = &state.vregs[vreg.vreg()];
                 if let Some(slot) = data.slot_idx {
-                    trace!("Marking vreg {} as saved on stack at {}", vreg, slot);
-                    state
-                        .safepoint_slots
-                        .push((pos, Allocation::stack(SpillSlot::new(slot as usize))));
+                    if domtree::dominates(&const_state.idom, data.def_block.unwrap(), block) {
+                        trace!("Marking vreg {} as saved on stack at {}", vreg, slot);
+                        state
+                            .safepoint_slots
+                            .push((pos, Allocation::stack(SpillSlot::new(slot as usize))));
+                    } else {
+                        trace!("Skipping {} as it does not dominate", vreg);
+                    }
                 }
             }
         }
@@ -560,6 +590,7 @@ fn allocate_block_insts<'a, F: Function>(
                                 late_write_disallow_regs.add(reg);
                             }
 
+                            state.vregs[vreg.vreg()].def_block = Some(block);
                             state.allocs[alloc_idx + i] = Allocation::reg(reg);
                             state.assign_preg(reg, vreg);
 
@@ -634,6 +665,7 @@ fn allocate_block_insts<'a, F: Function>(
                             }
                         }
                         OperandKind::Def => {
+                            state.vregs[vreg.vreg()].def_block = Some(block);
                             let slot = state.alloc_stack_slot(vreg);
                             state.allocs[alloc_idx + i] =
                                 Allocation::stack(SpillSlot::new(slot as usize));
@@ -684,6 +716,7 @@ fn allocate_block_insts<'a, F: Function>(
                             state.move_to_preg(vreg, reg, ProgPoint::before(inst));
                         } else {
                             // early def
+                            state.vregs[vreg.vreg()].def_block = Some(block);
                             state.assign_preg(reg, vreg);
                             state.alloc_stack_slot(vreg);
                             state.move_to_stack(reg, vreg, ProgPoint::after(inst));
@@ -724,6 +757,8 @@ fn allocate_block_insts<'a, F: Function>(
                 continue;
             }
 
+            state.vregs[vreg.vreg()].def_block = Some(block);
+
             match op.constraint() {
                 OperandConstraint::Reg => {
                     // find first non-allocated register
@@ -874,6 +909,13 @@ fn handle_out_block_params<'a, F: Function>(
             let out_slot = state.vregs[out_vreg.vreg()].slot_idx.unwrap();
             let in_slot = state.vregs[in_vreg.vreg()].slot_idx.unwrap();
 
+            trace!(
+                "Move {} from slot {} to slot {} for {}",
+                out_vreg,
+                out_slot,
+                in_slot,
+                in_vreg
+            );
             state.edits.push((
                 ProgPoint::before(last_inst),
                 Edit::Move {
@@ -915,6 +957,8 @@ fn handle_out_block_params<'a, F: Function>(
                 debug_assert!(state.vregs[out_vreg.vreg()].slot_idx.is_some());
                 let out_slot_idx = state.vregs[out_vreg.vreg()].slot_idx.unwrap();
 
+                state.vregs[in_vreg.vreg()].def_block = Some(succ);
+
                 // TODO: if out_vreg dies at this edge, we could reuse its stack slot
                 // TODO: we should also be able to reuse the slot if the successor only has one predecessor (us); check with AE
                 let mut no_alias = false;
@@ -947,6 +991,13 @@ fn handle_out_block_params<'a, F: Function>(
                     tmp_reg_float
                 };
                 let slot = state.create_stack_slot(out_vreg.class());
+                trace!(
+                    "Moving {} from slot {} to slot {} for {}",
+                    out_vreg,
+                    out_slot_idx,
+                    slot,
+                    in_vreg
+                );
                 state.edits.push((
                     ProgPoint::before(last_inst),
                     Edit::Move {