fixed first reg version

src/ion/fast_alloc.rs

@@ -140,6 +140,38 @@ impl ReadOnlyData {
     }
 }
 
+// https://burtleburtle.net/bob/rand/smallprng.html
+struct PRNG {
+    a: u64,
+    b: u64,
+    c: u64,
+    d: u64,
+}
+
+impl PRNG {
+    fn new(seed: u64) -> Self {
+        Self {
+            a: 0xf1ea5eed,
+            b: seed,
+            c: seed,
+            d: seed,
+        }
+    }
+
+    fn val(&mut self) -> u64 {
+        let e = self.a - PRNG::rot(self.b, 27);
+        self.a = self.b ^ PRNG::rot(self.c, 17);
+        self.b = self.c + self.d;
+        self.c = self.d + e;
+        self.d = e + self.a;
+        self.d
+    }
+
+    fn rot(x: u64, k: u64) -> u64 {
+        (x << k) | (x >> (32 - k))
+    }
+}
+
 struct FastAllocState<'a, F: Function> {
     pub vregs: Vec<VRegData>,
     pub pregs: Vec<PRegData>,
@@ -164,6 +196,7 @@ struct FastAllocState<'a, F: Function> {
 
     pub func: &'a F,
     pub mach_env: &'a MachineEnv,
+    pub prng: PRNG,
 }
 
 impl<'a, F: Function> FastAllocState<'a, F> {
@@ -198,6 +231,11 @@ impl<'a, F: Function> FastAllocState<'a, F> {
             blocks
         };
 
+        trace!(
+            "Num Insts: {} Num Blocks: {}",
+            func.num_insts(),
+            func.num_blocks()
+        );
         let mut inst_alloc_offsets = Vec::with_capacity(func.num_insts());
         inst_alloc_offsets.resize(func.num_insts(), 0);
 
@@ -220,6 +258,12 @@ impl<'a, F: Function> FastAllocState<'a, F> {
             allocs
         };
 
+        trace!("InstAllocOffsets: {:?}", inst_alloc_offsets);
+        trace!("Allocs Len: {}", allocs.len());
+
+        let prng = PRNG::new(
+            (blocks.len() as u64) << 48 | (vregs.len() as u64) << 32 | (allocs.len() as u64),
+        );
         Self {
             vregs,
             pregs,
@@ -245,6 +289,7 @@ impl<'a, F: Function> FastAllocState<'a, F> {
 
             func,
             mach_env,
+            prng,
         }
     }
 
@@ -259,7 +304,10 @@ impl<'a, F: Function> FastAllocState<'a, F> {
     pub fn alloc_stack_slot(&mut self, vreg: VReg) -> u32 {
         let data = &mut self.vregs[vreg.vreg()];
         if data.slot_idx.is_some() {
-            panic!("Trying to allocate already allocated stack slot");
+            panic!(
+                "Trying to allocate already allocated stack slot for {}",
+                vreg
+            );
         }
 
         let size = if vreg.class() == RegClass::Int {
@@ -286,14 +334,21 @@ impl<'a, F: Function> FastAllocState<'a, F> {
     }
 
     pub fn move_to_preg(&mut self, vreg: VReg, preg: PReg, pos: ProgPoint) {
-        if let Some(vreg) = &self.pregs[preg.index()].vreg {
+        trace!("Move {} to {} at {:?}", vreg, preg, pos);
+        /*if let Some(vreg) = &self.pregs[preg.index()].vreg {
             let vdata = &mut self.vregs[vreg.vreg() as usize];
             debug_assert!(vdata.preg.is_some());
             debug_assert_eq!(vdata.preg.unwrap(), preg);
             vdata.preg = None;
-        }
+        }*/
 
         if let Some(cur_preg) = &self.vregs[vreg.vreg()].preg {
+            if *cur_preg == preg {
+                trace!("{} already in target reg", vreg);
+                return;
+            }
+
+            trace!("Moving directly from {} to {}", cur_preg, preg);
             // Do a reg->reg move
             self.edits.push((
                 pos,
@@ -310,9 +365,11 @@ impl<'a, F: Function> FastAllocState<'a, F> {
             pdata.vreg = None;
 
             self.pregs[preg.index()].vreg = Some(vreg);
+            self.vregs[vreg.vreg()].preg = Some(preg);
             return;
         }
 
+        self.clear_preg(preg);
         let vdata = &mut self.vregs[vreg.vreg()];
         let pdata = &mut self.pregs[preg.index()];
 
@@ -320,10 +377,12 @@ impl<'a, F: Function> FastAllocState<'a, F> {
             panic!("Trying to move from vreg that has no stack slot to preg");
         }
 
+        let slot = vdata.slot_idx.unwrap() as usize;
+        trace!("Moving from slot {}", slot);
         self.edits.push((
             pos,
             Edit::Move {
-                from: Allocation::stack(SpillSlot::new(vdata.slot_idx.unwrap() as usize)),
+                from: Allocation::stack(SpillSlot::new(slot)),
                 to: Allocation::reg(preg),
             },
         ));
@@ -336,6 +395,7 @@ impl<'a, F: Function> FastAllocState<'a, F> {
     }
 
     pub fn move_to_stack(&mut self, preg: PReg, vreg: VReg, pos: ProgPoint) {
+        trace!("Move {} of {} to stack at {:?}", preg, vreg, pos);
         let vdata = &mut self.vregs[vreg.vreg()];
         let pdata = &mut self.pregs[preg.index()];
         if pdata.vreg.is_none() || vdata.preg.is_none() {
@@ -358,6 +418,7 @@ impl<'a, F: Function> FastAllocState<'a, F> {
     }
 
     pub fn assign_preg(&mut self, preg: PReg, vreg: VReg) {
+        trace!("Assigning {} to {}", vreg, preg);
         // TODO: somewhere assign_preg is called without making sure the vreg is clear (or inspite of it)
         // need to make sure this is intended behavior
         self.clear_preg(preg);
@@ -374,6 +435,7 @@ impl<'a, F: Function> FastAllocState<'a, F> {
     }
 
     fn clear_preg_idx(&mut self, preg: usize) {
+        trace!("Clearing preg {}", preg);
         let pdata = &mut self.pregs[preg];
         if let Some(vreg) = pdata.vreg {
             let vdata = &mut self.vregs[vreg.vreg()];
@@ -387,8 +449,10 @@ impl<'a, F: Function> FastAllocState<'a, F> {
     }
 
     pub fn clear_vreg_from_reg(&mut self, vreg: VReg) {
+        trace!("Clearing vreg {} from reg", vreg);
         let vdata = &mut self.vregs[vreg.vreg()];
         if let Some(preg) = vdata.preg {
+            debug_assert!(self.pregs[preg.index()].vreg.is_some());
             debug_assert_eq!(self.pregs[preg.index()].vreg.unwrap().vreg(), vreg.vreg());
             self.pregs[preg.index()].vreg = None;
             vdata.preg = None;
@@ -413,6 +477,23 @@ impl<'a, F: Function> FastAllocState<'a, F> {
             }
         }
     }
+
+    pub fn vreg_used_at_cur_inst(&mut self, vreg: VReg) -> bool {
+        let vdata = &self.vregs[vreg.vreg()];
+        if vdata.cur_use_idx as usize >= vdata.uses.len() {
+            return false;
+        }
+        vdata.uses[vdata.cur_use_idx as usize] == self.cur_inst_pos as u32
+    }
+
+    pub fn vreg_next_use(&self, vreg: VReg) -> Option<u32> {
+        let vdata = &self.vregs[vreg.vreg()];
+        if vdata.cur_use_idx as usize >= vdata.uses.len() {
+            return None;
+        }
+
+        Some(vdata.uses[vdata.cur_use_idx as usize])
+    }
 }
 
 pub fn run<F: Function>(func: &F, mach_env: &MachineEnv) -> Result<Output, RegAllocError> {
@@ -437,6 +518,12 @@ pub fn run<F: Function>(func: &F, mach_env: &MachineEnv) -> Result<Output, RegAl
         state.blocks[block.index()].reg_allocated = true;
 
         trace!("Allocating block {}", block.index());
+        trace!("Allocated pregs:");
+        for i in 0..state.pregs.len() {
+            if let Some(vreg) = &state.pregs[i].vreg {
+                trace!("p{}: {}", i, vreg);
+            }
+        }
 
         allocate_block_insts(&mut state, &const_state, block)?;
         handle_out_block_params(&mut state, &const_state, block)?;
@@ -489,19 +576,57 @@ fn vreg_killed<'a, F: Function>(
     block: Block,
     block_last_inst: usize,
     vreg: usize,
+    // if the vreg is used at the current instruction, count it as not killed
+    // TODO: this is currently always true but can be used for optimization later on
+    save_on_current_use: bool,
 ) -> bool {
     let info = &state.vregs[vreg];
     let block_after_pos = state.cur_inst_pos + (block_last_inst - inst.index()) + 1;
     let cur_use_idx = info.cur_use_idx as usize;
+    let cur_pos = if save_on_current_use {
+        state.cur_inst_pos
+    } else {
+        state.cur_inst_pos + 1
+    };
+
+    trace!(
+        "Checking live-status of v{} in {:?} at inst {:?} (CurPos: {} SaveCurPos? {}): Liveout {}, block_after: {}",
+        vreg,
+        block,
+        inst,
+        state.cur_inst_pos,
+        save_on_current_use,
+        state.liveouts[block.index()].get(vreg),
+        block_after_pos
+    );
+    trace!(
+        "Uses of v{}: {:?}. Currently at {}",
+        vreg,
+        info.uses,
+        info.cur_use_idx
+    );
 
     if !state.liveouts[block.index()].get(vreg) {
-        if info.uses.len() >= cur_use_idx {
+        if info.uses.len() <= cur_use_idx {
+            trace!("Uses exhausted, vreg must be dead");
             return true;
         }
-        if info.uses[cur_use_idx] == state.cur_inst_pos as u32 {
+
-            if info.uses.len() >= cur_use_idx + 1
+        if info.uses.len() <= cur_use_idx + 1 {
+            trace!("next use: {}, no use after that", info.uses[cur_use_idx]);
+        } else {
+            trace!(
+                "next use: {}, {}",
+                info.uses[cur_use_idx],
+                info.uses[cur_use_idx + 1]
+            );
+        }
+
+        if !save_on_current_use && info.uses[cur_use_idx] == state.cur_inst_pos as u32 {
+            if info.uses.len() <= cur_use_idx + 1
                 || info.uses[cur_use_idx + 1] >= block_after_pos as u32
             {
+                trace!("v{} is killed", vreg);
                 return true;
             }
         }
@@ -574,25 +699,61 @@ fn allocate_block_insts<'a, F: Function>(
             }
         }
 
+        for preg in clobbers {
+            // TODO: this might save a use that is killed at this inst i think
+            let vreg = if let Some(vreg) = &state.pregs[preg.index()].vreg {
+                *vreg
+            } else {
+                continue;
+            };
+
+            if state.vregs[vreg.vreg()].slot_idx.is_some() {
+                trace!("{} with {} clobbered but saved on stack", preg, vreg);
+                state.clear_preg(preg);
+                continue;
+            }
+
+            if !state.vreg_used_at_cur_inst(vreg)
+                && vreg_killed(state, inst, block, block_last_inst_idx, vreg.vreg(), true)
+            {
+                trace!("{} with {} clobbered but vreg killed", preg, vreg);
+                state.clear_preg(preg);
+                continue;
+            }
+
+            state.alloc_stack_slot(vreg);
+            state.move_to_stack(preg, vreg, ProgPoint::before(inst));
+            state.clear_preg(preg);
+        }
+
         // we allocate fixed defs/uses and stack allocations first
+        // TODO: if a fixed def is before a fixed use with the same preg here, it will incorrectly update state!!!
         trace!("First alloc pass");
         for (i, op) in operands.iter().enumerate() {
             let vreg = op.vreg();
 
             trace!("Operand {}: {}", i, op);
-            if vreg == VReg::invalid() {
+            if op.as_fixed_nonallocatable().is_some() {
                 // it seems cranelift emits fixed reg uses with invalid vregs, handle them here
                 // TODO: treat them like normal vregs by just using last_vreg_index+1 for them?
                 match op.constraint() {
                     OperandConstraint::FixedReg(reg) => {
                         // Save vreg if needed
                         if let Some(vreg) = state.pregs[reg.index()].vreg {
-                            let vreg = vreg.vreg();
+                            let vreg_idx = vreg.vreg();
-                            if state.vregs[vreg].slot_idx.is_none()
+                            if state.vregs[vreg_idx].slot_idx.is_none()
-                                && !vreg_killed(state, inst, block, block_last_inst_idx, vreg)
+                                && (state.vreg_used_at_cur_inst(vreg)
+                                    || !vreg_killed(
+                                        state,
+                                        inst,
+                                        block,
+                                        block_last_inst_idx,
+                                        vreg_idx,
+                                        true,
+                                    ))
                             {
                                 let slot = state.create_stack_slot(reg.class());
-                                state.vregs[vreg].slot_idx = Some(slot);
+                                state.vregs[vreg_idx].slot_idx = Some(slot);
                                 state.edits.push((
                                     ProgPoint::before(inst),
                                     Edit::Move {
@@ -642,18 +803,27 @@ fn allocate_block_insts<'a, F: Function>(
                             if let Some(cur_preg) = state.vregs[vreg.vreg()].preg {
                                 if cur_preg == reg {
                                     trace!("{} already in target {}", vreg, cur_preg);
+                                    state.allocs[alloc_idx + i] = Allocation::reg(cur_preg);
                                     continue;
                                 }
                             }
 
                             // Save vreg if needed
                             if let Some(vreg) = state.pregs[reg.index()].vreg {
-                                let vreg = vreg.vreg();
+                                let vreg_idx = vreg.vreg();
-                                if state.vregs[vreg].slot_idx.is_none()
+                                if state.vregs[vreg_idx].slot_idx.is_none()
-                                    && !vreg_killed(state, inst, block, block_last_inst_idx, vreg)
+                                    && (state.vreg_used_at_cur_inst(vreg)
+                                        || !vreg_killed(
+                                            state,
+                                            inst,
+                                            block,
+                                            block_last_inst_idx,
+                                            vreg_idx,
+                                            true,
+                                        ))
                                 {
                                     let slot = state.create_stack_slot(reg.class());
-                                    state.vregs[vreg].slot_idx = Some(slot);
+                                    state.vregs[vreg_idx].slot_idx = Some(slot);
                                     state.edits.push((
                                         ProgPoint::before(inst),
                                         Edit::Move {
@@ -665,6 +835,7 @@ fn allocate_block_insts<'a, F: Function>(
                             }
 
                             if let Some(cur_preg) = state.vregs[vreg.vreg()].preg {
+                                trace!("Move {} directly from {} to {}", vreg, cur_preg, reg);
                                 // Move from preg to preg
                                 state.edits.push((
                                     ProgPoint::before(inst),
@@ -673,9 +844,12 @@ fn allocate_block_insts<'a, F: Function>(
                                         to: Allocation::reg(reg),
                                     },
                                 ));
+                                debug_assert_eq!(
+                                    state.pregs[cur_preg.index()].vreg.unwrap().vreg(),
+                                    vreg.vreg()
+                                );
                                 state.pregs[cur_preg.index()].vreg = None;
-                                state.vregs[vreg.vreg()].preg = Some(reg);
+                                state.assign_preg(reg, vreg);
-                                state.pregs[reg.index()].vreg = Some(vreg);
                             } else {
                                 state.move_to_preg(vreg, reg, ProgPoint::before(inst));
                             }
@@ -715,12 +889,21 @@ fn allocate_block_insts<'a, F: Function>(
 
                             // Save vreg if needed
                             if let Some(vreg) = state.pregs[reg.index()].vreg {
-                                let vreg = vreg.vreg();
+                                let vreg_idx = vreg.vreg();
-                                if state.vregs[vreg].slot_idx.is_none()
+                                if state.vregs[vreg_idx].slot_idx.is_none()
-                                    && !vreg_killed(state, inst, block, block_last_inst_idx, vreg)
+                                    && (op.pos() != OperandPos::Late
+                                        && state.vreg_used_at_cur_inst(vreg)
+                                        || !vreg_killed(
+                                            state,
+                                            inst,
+                                            block,
+                                            block_last_inst_idx,
+                                            vreg_idx,
+                                            true,
+                                        ))
                                 {
                                     let slot = state.create_stack_slot(reg.class());
-                                    state.vregs[vreg].slot_idx = Some(slot);
+                                    state.vregs[vreg_idx].slot_idx = Some(slot);
                                     state.edits.push((
                                         ProgPoint::before(inst),
                                         Edit::Move {
@@ -769,18 +952,21 @@ fn allocate_block_insts<'a, F: Function>(
                                 let tmp_reg = pregs.into_iter().next().unwrap();
                                 if let Some(vreg) = state.pregs[tmp_reg.index()].vreg {
                                     // Save vreg if needed
-                                    let vreg = vreg.vreg();
+                                    let vreg_idx = vreg.vreg();
-                                    if state.vregs[vreg].slot_idx.is_none()
+                                    if state.vregs[vreg_idx].slot_idx.is_none()
-                                        && !vreg_killed(
+                                        && (op.pos() != OperandPos::Late
+                                            && state.vreg_used_at_cur_inst(vreg)
+                                            || !vreg_killed(
                                                 state,
                                                 inst,
                                                 block,
                                                 block_last_inst_idx,
-                                            vreg,
+                                                vreg_idx,
-                                        )
+                                                true,
+                                            ))
                                     {
                                         let slot = state.create_stack_slot(reg.class());
-                                        state.vregs[vreg].slot_idx = Some(slot);
+                                        state.vregs[vreg_idx].slot_idx = Some(slot);
                                         state.edits.push((
                                             ProgPoint::before(inst),
                                             Edit::Move {
@@ -851,12 +1037,24 @@ fn allocate_block_insts<'a, F: Function>(
             trace!("Operand {}: {}", i, op);
 
             let vreg = op.vreg();
-            if vreg == VReg::invalid() {
+            if op.as_fixed_nonallocatable().is_some() {
                 continue;
             }
 
             match op.constraint() {
                 OperandConstraint::Reg => {
+                    // Are we alredy in a reg?
+                    if let Some(cur_preg) = &state.vregs[vreg.vreg()].preg {
+                        assert_eq!(op.kind(), OperandKind::Use);
+                        // Late uses need to survive the instruction
+                        if op.pos() == OperandPos::Early || !clobbers.contains(*cur_preg) {
+                            trace!("{} already in reg {}. Using that", vreg, cur_preg);
+                            state.allocs[alloc_idx + i] = Allocation::reg(*cur_preg);
+                            regs_allocated.add(*cur_preg);
+                            continue;
+                        }
+                    }
+
                     // find first non-allocated register
                     let reg_order = const_state.reg_order(op.class());
                     let mut allocated = false;
@@ -865,15 +1063,34 @@ fn allocate_block_insts<'a, F: Function>(
                             continue;
                         }
 
-                        if state.pregs[reg.index()].vreg.is_some() {
+                        if let Some(cur_vreg) = &state.pregs[reg.index()].vreg {
+                            // we can override the reg if the vreg was killed already
+                            if !vreg_killed(
+                                state,
+                                inst,
+                                block,
+                                block_last_inst_idx,
+                                cur_vreg.vreg(),
+                                true,
+                            ) {
                                 continue;
                             }
+                            state.clear_preg(reg);
+                        }
 
                         // reg should not contain anything
                         debug_assert!(state.pregs[reg.index()].vreg.is_none());
 
+                        if op.kind() == OperandKind::Use
+                            && op.pos() == OperandPos::Late
+                            && clobbers.contains(reg)
+                        {
+                            continue;
+                        }
+
                         state.allocs[alloc_idx + i] = Allocation::reg(reg);
                         regs_allocated.add(reg);
+                        trace!("Chose {} for operand {}", reg, i);
                         if op.kind() == OperandKind::Use {
                             if req_refs_on_stack && state.vregs[vreg.vreg()].reftype {
                                 panic!("reftype required to be in reg at safepoint");
@@ -890,7 +1107,6 @@ fn allocate_block_insts<'a, F: Function>(
                             state.move_to_stack(reg, vreg, ProgPoint::after(inst));
                         }
 
-                        trace!("Chose {} for operand {}", reg, i);
                         allocated = true;
                         break;
                     }
@@ -903,15 +1119,22 @@ fn allocate_block_insts<'a, F: Function>(
                     // No register available
                     // TODO: first evict pregs that already have a stack slot even if they are used earlier?
                     let mut evict_candidate = None;
+                    let mut ffa_reg_pool = PRegSet::empty();
                     for &reg in reg_order {
                         if regs_allocated.contains(reg) {
                             continue;
                         }
 
+                        if op.kind() == OperandKind::Use
+                            && op.pos() == OperandPos::Late
+                            && clobbers.contains(reg)
+                        {
+                            continue;
+                        }
+
                         debug_assert!(state.pregs[reg.index()].vreg.is_some());
-                        let vreg = state.pregs[reg.index()].vreg.unwrap().vreg();
+                        let vreg = state.pregs[reg.index()].vreg.unwrap();
-                        let next_use =
+                        if let Some(next_use) = state.vreg_next_use(vreg) {
-                            state.vregs[vreg].uses[state.vregs[vreg].cur_use_idx as usize];
                             if next_use == state.cur_inst_pos as u32 {
                                 continue;
                             }
@@ -923,14 +1146,31 @@ fn allocate_block_insts<'a, F: Function>(
                             } else {
                                 evict_candidate = Some((reg, next_use));
                             }
+                        } else {
+                            // see further below
+                            ffa_reg_pool.add(reg);
+
+                            //panic!("preg should have already been chosen")
                         }
+                    }
+
+                    // TODO: we need some logic to shuffle assignments around if there is a late use that needs to survive a clobber
+                    // and another reg is available but taken by an early use so it would not be an eviction_candidate
 
                     if let Some((reg, next_use)) = evict_candidate {
                         // Save vreg if needed
+                        {
                             let vreg = state.pregs[reg.index()].vreg.unwrap();
-                        log::trace!("Evicting {} with v{}", reg, vreg);
+                            trace!("Evicting {} with v{}", reg, vreg);
                             if state.vregs[vreg.vreg()].slot_idx.is_none()
-                            && !vreg_killed(state, inst, block, block_last_inst_idx, vreg.vreg())
+                                && !vreg_killed(
+                                    state,
+                                    inst,
+                                    block,
+                                    block_last_inst_idx,
+                                    vreg.vreg(),
+                                    true,
+                                )
                             {
                                 let slot = state.create_stack_slot(reg.class());
                                 state.vregs[vreg.vreg()].slot_idx = Some(slot);
@@ -942,6 +1182,7 @@ fn allocate_block_insts<'a, F: Function>(
                                     },
                                 ));
                             }
+                        }
                         state.clear_preg(reg);
 
                         state.allocs[alloc_idx + i] = Allocation::reg(reg);
@@ -964,9 +1205,62 @@ fn allocate_block_insts<'a, F: Function>(
 
                         trace!("Chose {} for operand {}", reg, i);
                     } else {
+                        if ffa_reg_pool == PRegSet::empty() {
                             panic!("Out of registers: {:?}", regs_allocated);
                             return Err(RegAllocError::TooManyLiveRegs);
                         }
+
+                        let preg = 'block: {
+                            let len = ffa_reg_pool.bits.count_ones() as usize;
+                            let mut idx = (state.prng.val() as usize % 128) % len;
+                            for preg in ffa_reg_pool.into_iter() {
+                                if idx == 0 {
+                                    break 'block preg;
+                                }
+
+                                idx -= 1;
+                            }
+
+                            panic!("I can't do math");
+                        };
+
+                        trace!("Chose {} from ffa_reg_pool", preg);
+                        {
+                            let vreg = state.pregs[preg.index()].vreg.unwrap();
+                            // need to save vreg if it does not have a slot
+                            if state.vregs[vreg.vreg()].slot_idx.is_none() {
+                                let slot = state.create_stack_slot(preg.class());
+                                state.vregs[vreg.vreg()].slot_idx = Some(slot);
+                                state.edits.push((
+                                    ProgPoint::before(inst),
+                                    Edit::Move {
+                                        from: Allocation::reg(preg),
+                                        to: Allocation::stack(SpillSlot::new(slot as usize)),
+                                    },
+                                ));
+                            }
+                        }
+                        state.clear_preg(preg);
+                        state.allocs[alloc_idx + i] = Allocation::reg(preg);
+                        regs_allocated.add(preg);
+                        if op.kind() == OperandKind::Use {
+                            if req_refs_on_stack && state.vregs[vreg.vreg()].reftype {
+                                panic!("reftype required to be in reg at safepoint");
+                                return Err(RegAllocError::TooManyLiveRegs);
+                            }
+
+                            // need to move from stack to reg
+                            state.move_to_preg(vreg, preg, ProgPoint::before(inst));
+                        } else {
+                            // early def
+                            state.vregs[vreg.vreg()].def_block = Some(block);
+                            state.assign_preg(preg, vreg);
+                            state.alloc_stack_slot(vreg);
+                            state.move_to_stack(preg, vreg, ProgPoint::after(inst));
+                        }
+
+                        trace!("Chose {} for operand {}", preg, i);
+                    }
                 }
                 OperandConstraint::Reuse(_) => {
                     panic!("Illegal register constraint reuse for early def or use");
@@ -975,6 +1269,62 @@ fn allocate_block_insts<'a, F: Function>(
             }
         }
 
+        // advance use_idx
+        for op in operands {
+            if op.kind() != OperandKind::Use || op.as_fixed_nonallocatable().is_some() {
+                continue;
+            }
+
+            let vreg_idx = op.vreg().vreg();
+            let info = &mut state.vregs[vreg_idx];
+            info.cur_use_idx += 1;
+
+            if vreg_killed(state, inst, block, block_last_inst_idx, vreg_idx, true) {
+                // TODO: clear stack slot
+                state.clear_vreg_from_reg(op.vreg());
+            }
+            /*let block_after_pos = state.cur_inst_pos + (block_last_inst_idx - inst.index()) + 1;
+            // check if vreg dies here
+            if !state.liveouts[block.index()].get(vreg_idx)
+                && (info.uses.len() <= info.cur_use_idx as usize
+                    || info.uses[info.cur_use_idx as usize] > block_after_pos as u32)
+            {
+                // TODO: clear stack slot
+                state.clear_vreg_from_reg(op.vreg());
+            }*/
+        }
+
+        // TODO: this is currently a fix for the register state for uses that are clobbered as it is incorrectly set
+        // but this is inefficient as we could check for this when handling uses
+        trace!("Late clobber handling");
+        for preg in clobbers {
+            // TODO: this might save a use that is killed at this inst i think
+            let vreg = if let Some(vreg) = &state.pregs[preg.index()].vreg {
+                *vreg
+            } else {
+                continue;
+            };
+
+            if state.vregs[vreg.vreg()].slot_idx.is_some() {
+                trace!("{} with {} clobbered but saved on stack", preg, vreg);
+                state.clear_preg(preg);
+                continue;
+            }
+
+            // we don't care if the reg is used at the current inst
+            if vreg_killed(state, inst, block, block_last_inst_idx, vreg.vreg(), true) {
+                trace!("{} with {} clobbered but vreg killed", preg, vreg);
+                state.clear_preg(preg);
+                continue;
+            }
+
+            // TODO: this should not be hit i think as all we should be clearing here are use assignments and the vregs
+            // that need to be saved should have been saved at the check before
+            state.alloc_stack_slot(vreg);
+            state.move_to_stack(preg, vreg, ProgPoint::before(inst));
+            state.clear_preg(preg);
+        }
+
         // alloc non-fixed late defs and reuse
         trace!("Third alloc pass");
         for (i, op) in operands.iter().enumerate() {
@@ -984,7 +1334,7 @@ fn allocate_block_insts<'a, F: Function>(
 
             trace!("Operand {}: {}", i, op);
             let vreg = op.vreg();
-            if vreg == VReg::invalid() {
+            if op.as_fixed_nonallocatable().is_some() {
                 continue;
             }
 
@@ -999,9 +1349,20 @@ fn allocate_block_insts<'a, F: Function>(
                         if regs_allocated.contains(reg) || late_write_disallow_regs.contains(reg) {
                             continue;
                         }
-                        if state.pregs[reg.index()].vreg.is_some() {
+                        if let Some(cur_vreg) = &state.pregs[reg.index()].vreg {
+                            // we can override the reg if the vreg was killed already
+                            if !vreg_killed(
+                                state,
+                                inst,
+                                block,
+                                block_last_inst_idx,
+                                cur_vreg.vreg(),
+                                true,
+                            ) {
                                 continue;
                             }
+                            state.clear_preg(reg);
+                        }
 
                         // reg should not contain anything
                         regs_allocated.add(reg);
@@ -1024,15 +1385,15 @@ fn allocate_block_insts<'a, F: Function>(
 
                     // TODO: first evict pregs that already have a stack slot even if they are used earlier?
                     let mut evict_candidate = None;
+                    let mut ffa_reg_pool = PRegSet::empty();
                     for &reg in reg_order {
                         if regs_allocated.contains(reg) || late_write_disallow_regs.contains(reg) {
                             continue;
                         }
 
                         debug_assert!(state.pregs[reg.index()].vreg.is_some());
-                        let vreg = state.pregs[reg.index()].vreg.unwrap().vreg();
+                        let vreg = state.pregs[reg.index()].vreg.unwrap();
-                        let next_use =
+                        if let Some(next_use) = state.vreg_next_use(vreg) {
-                            state.vregs[vreg].uses[state.vregs[vreg].cur_use_idx as usize];
                             if next_use == state.cur_inst_pos as u32 {
                                 continue;
                             }
@@ -1044,14 +1405,29 @@ fn allocate_block_insts<'a, F: Function>(
                             } else {
                                 evict_candidate = Some((reg, next_use));
                             }
+                        } else {
+                            // if we hit this it means that all uses are "before" this one in lowering-order
+                            // we should probably find a nice heuristic for chosing which register to choose
+                            // here. tbf we should probably find an overall better heuristic for chosing which register to evict
+                            // rn just add the reg to a set and we pick a random one later
+                            ffa_reg_pool.add(reg);
+                        }
                     }
 
                     if let Some((reg, next_use)) = evict_candidate {
                         // Save vreg if needed
+                        {
                             let vreg = state.pregs[reg.index()].vreg.unwrap();
-                        log::trace!("Evicting {} with v{}", reg, vreg);
+                            trace!("Evicting {} with {}", reg, vreg);
                             if state.vregs[vreg.vreg()].slot_idx.is_none()
-                            && !vreg_killed(state, inst, block, block_last_inst_idx, vreg.vreg())
+                                && !vreg_killed(
+                                    state,
+                                    inst,
+                                    block,
+                                    block_last_inst_idx,
+                                    vreg.vreg(),
+                                    true,
+                                )
                             {
                                 let slot = state.create_stack_slot(reg.class());
                                 state.vregs[vreg.vreg()].slot_idx = Some(slot);
@@ -1063,6 +1439,7 @@ fn allocate_block_insts<'a, F: Function>(
                                     },
                                 ));
                             }
+                        }
                         state.clear_preg(reg);
 
                         regs_allocated.add(reg);
@@ -1073,9 +1450,62 @@ fn allocate_block_insts<'a, F: Function>(
                         state.move_to_stack(reg, vreg, ProgPoint::after(inst));
                         trace!("Chose {} for operand {}", reg, i);
                     } else {
-                        panic!("out of registers");
+                        if ffa_reg_pool == PRegSet::empty() {
+                            panic!("Out of registers: {:?}", regs_allocated);
                             return Err(RegAllocError::TooManyLiveRegs);
                         }
+
+                        let preg = 'block: {
+                            let len = ffa_reg_pool.bits.count_ones() as usize;
+                            let mut idx = (state.prng.val() as usize % 128) % len;
+                            for preg in ffa_reg_pool.into_iter() {
+                                if idx == 0 {
+                                    break 'block preg;
+                                }
+
+                                idx -= 1;
+                            }
+
+                            panic!("I can't do math");
+                        };
+
+                        trace!("Chose {} from ffa_reg_pool", preg);
+                        {
+                            let vreg = state.pregs[preg.index()].vreg.unwrap();
+                            // need to save vreg if it does not have a slot
+                            if state.vregs[vreg.vreg()].slot_idx.is_none() {
+                                let slot = state.create_stack_slot(preg.class());
+                                state.vregs[vreg.vreg()].slot_idx = Some(slot);
+                                state.edits.push((
+                                    ProgPoint::before(inst),
+                                    Edit::Move {
+                                        from: Allocation::reg(preg),
+                                        to: Allocation::stack(SpillSlot::new(slot as usize)),
+                                    },
+                                ));
+                            }
+                        }
+                        state.clear_preg(preg);
+                        state.allocs[alloc_idx + i] = Allocation::reg(preg);
+                        regs_allocated.add(preg);
+                        if op.kind() == OperandKind::Use {
+                            if req_refs_on_stack && state.vregs[vreg.vreg()].reftype {
+                                panic!("reftype required to be in reg at safepoint");
+                                return Err(RegAllocError::TooManyLiveRegs);
+                            }
+
+                            // need to move from stack to reg
+                            state.move_to_preg(vreg, preg, ProgPoint::before(inst));
+                        } else {
+                            // early def
+                            state.vregs[vreg.vreg()].def_block = Some(block);
+                            state.assign_preg(preg, vreg);
+                            state.alloc_stack_slot(vreg);
+                            state.move_to_stack(preg, vreg, ProgPoint::after(inst));
+                        }
+
+                        trace!("Chose {} for operand {}", preg, i);
+                    }
                 }
                 OperandConstraint::Reuse(idx) => {
                     debug_assert!(state.allocs[alloc_idx + idx].is_reg());
@@ -1085,12 +1515,13 @@ fn allocate_block_insts<'a, F: Function>(
 
                     // Save vreg on stack if it is not killed
                     if let Some(vreg) = state.pregs[preg.index()].vreg {
-                        let vreg = vreg.vreg();
+                        let vreg_idx = vreg.vreg();
-                        if state.vregs[vreg].slot_idx.is_none()
+                        if state.vregs[vreg_idx].slot_idx.is_none()
-                            && !vreg_killed(state, inst, block, block_last_inst_idx, vreg)
+                            && !vreg_killed(state, inst, block, block_last_inst_idx, vreg_idx, true)
                         {
+                            trace!("Saving {}", vreg);
                             let slot = state.create_stack_slot(preg.class());
-                            state.vregs[vreg].slot_idx = Some(slot);
+                            state.vregs[vreg_idx].slot_idx = Some(slot);
                             state.edits.push((
                                 ProgPoint::before(inst),
                                 Edit::Move {
@@ -1112,26 +1543,6 @@ fn allocate_block_insts<'a, F: Function>(
             }
         }
 
-        for op in operands {
-            if op.kind() != OperandKind::Use || op.as_fixed_nonallocatable().is_some() {
-                continue;
-            }
-
-            let vreg_idx = op.vreg().vreg();
-            let info = &mut state.vregs[vreg_idx];
-            info.cur_use_idx += 1;
-
-            let block_after_pos = state.cur_inst_pos + (block_last_inst_idx - inst.index()) + 1;
-            // check if vreg dies here
-            if !state.liveouts[block.index()].get(vreg_idx)
-                && (info.uses.len() >= info.cur_use_idx as usize
-                    || info.uses[info.cur_use_idx as usize] > block_after_pos as u32)
-            {
-                // TODO: clear stack slot
-                state.clear_vreg_from_reg(op.vreg());
-            }
-        }
-
         // fixup edit order
         let mut first_post_pos = None;
         for i in edit_start_idx..state.edits.len() {
@@ -1151,16 +1562,45 @@ fn allocate_block_insts<'a, F: Function>(
             }
         }
 
+        assert!(!state.allocs[alloc_idx..alloc_idx + operands.len()]
+            .iter()
+            .any(|a| a.is_none()));
+
+        trace!(
+            "Instruction Allocs: {:?}",
+            &state.allocs[alloc_idx..alloc_idx + operands.len()]
+        );
+
         state.cur_inst_pos += 1;
     }
 
-    // Move all liveout vregs to a stack slot if they dont have one and clear pregs
+    // Move all liveout/block param vregs to a stack slot if they dont have one and clear pregs
     for i in 0..state.pregs.len() {
         match state.pregs[i].vreg {
             None => {}
             Some(vreg) => {
+                trace!("Clearing {} from p{}", vreg, i);
                 let idx = vreg.vreg();
-                if state.liveouts[block.index()].get(idx) && state.vregs[idx].slot_idx.is_none() {
+                // TODO: obv dont need that if the block param handle funcs can handle reg locations
+                let is_out_param = 'block: {
+                    let last_inst = state.func.block_insns(block).last();
+                    if !state.func.is_branch(last_inst) {
+                        break 'block false;
+                    }
+
+                    for succ_idx in 0..state.func.block_succs(block).len() {
+                        for out_vreg in state.func.branch_blockparams(block, last_inst, succ_idx) {
+                            if *out_vreg == vreg {
+                                break 'block true;
+                            }
+                        }
+                    }
+
+                    false
+                };
+                if (is_out_param || state.liveouts[block.index()].get(idx))
+                    && state.vregs[idx].slot_idx.is_none()
+                {
                     let preg = state.vregs[idx].preg.unwrap();
                     let slot = state.create_stack_slot(preg.class());
                     state.edits.push((
@@ -1187,8 +1627,12 @@ fn handle_out_block_params<'a, F: Function>(
     const_state: &ReadOnlyData,
     block: Block,
 ) -> Result<(), RegAllocError> {
-    trace!("Allocating outgoing blockparams for {}", block.index());
     let last_inst = state.func.block_insns(block).last();
+    trace!(
+        "Allocating outgoing blockparams for {}, last_inst: {}",
+        block.index(),
+        last_inst.index()
+    );
     if !state.func.is_branch(last_inst) {
         trace!("Last inst {} is not a branch", last_inst.index());
         return Ok(());
@@ -1198,10 +1642,11 @@ fn handle_out_block_params<'a, F: Function>(
     {
         let alloc_start = state.inst_alloc_offsets[last_inst.index()] as usize;
         let alloc_end = if last_inst.index() + 1 == state.inst_alloc_offsets.len() {
-            state.inst_alloc_offsets.len()
+            state.allocs.len()
         } else {
             state.inst_alloc_offsets[last_inst.index() + 1] as usize
         };
+        trace!("alloc_start: {}, alloc_end: {}", alloc_start, alloc_end);
         for i in alloc_start..alloc_end {
             if let Some(reg) = state.allocs[i].clone().as_reg() {
                 pregs_used_by_br.add(reg);
@@ -1686,6 +2131,7 @@ fn handle_out_block_params<'a, F: Function>(
                 let out_vreg = out_params[i];
                 let in_vreg = in_params[i];
                 debug_assert!(state.vregs[out_vreg.vreg()].slot_idx.is_some());
+                debug_assert!(state.vregs[in_vreg.vreg()].slot_idx.is_none());
                 let out_slot_idx = state.vregs[out_vreg.vreg()].slot_idx.unwrap();
 
                 if out_vreg == VReg::invalid() {
@@ -1807,6 +2253,12 @@ fn calc_use_positions_and_live_bitmaps<'a, F: Function>(
                     continue;
                 }
 
+                trace!(
+                    "Use of {} at {} (inst {})",
+                    op.vreg(),
+                    cur_pos,
+                    inst.index()
+                );
                 state.vregs[op.vreg().vreg()].uses.push(cur_pos);
             }
 
@@ -1820,6 +2272,12 @@ fn calc_use_positions_and_live_bitmaps<'a, F: Function>(
 
         for i in 0..state.func.block_succs(block).len() {
             for vreg in state.func.branch_blockparams(block, last_inst, i) {
+                trace!(
+                    "Use of {} in blockparam at {} (inst {})",
+                    vreg,
+                    cur_pos,
+                    last_inst.index()
+                );
                 state.vregs[vreg.vreg()].uses.push(cur_pos);
             }
         }
@@ -1955,5 +2413,10 @@ fn calc_live_bitmaps<'a, F: Function>(
         return Err(RegAllocError::EntryLivein);
     }
 
+    for idx in 0..state.blocks.len() {
+        trace!("Livein for block {}: {:?}", idx, state.liveins[idx]);
+        trace!("Liveouts for block {}: {:?}", idx, state.liveouts[idx]);
+    }
+
     Ok(())
 }