more stuff in regs

src/ion/fast_alloc.rs

@@ -159,11 +159,11 @@ impl PRNG {
     }
 
     fn val(&mut self) -> u64 {
-        let e = self.a - PRNG::rot(self.b, 27);
+        let e = self.a.wrapping_sub(PRNG::rot(self.b, 27));
         self.a = self.b ^ PRNG::rot(self.c, 17);
-        self.b = self.c + self.d;
+        self.b = self.c.wrapping_add(self.d);
-        self.c = self.d + e;
+        self.c = self.d.wrapping_add(e);
-        self.d = e + self.a;
+        self.d = e.wrapping_add(self.a);
         self.d
     }
 
@@ -463,15 +463,26 @@ impl<'a, F: Function> FastAllocState<'a, F> {
         }
     }
 
-    pub fn clear_reftype_vregs(&mut self) {
+    pub fn clear_reftype_vregs(&mut self, cur_inst: Inst) {
         if self.reftype_vregs_in_pregs_count == 0 {
             return;
         }
 
         for i in 0..self.pregs.len() {
             if let Some(vreg) = self.pregs[i].vreg.clone() {
-                let vreg = vreg.vreg();
+                let vreg_idx = vreg.vreg();
-                if self.vregs[vreg].reftype {
+                if self.vregs[vreg_idx].reftype {
+                    debug_assert_eq!(self.vregs[vreg_idx].preg.unwrap().index(), i);
+                    if self.vregs[vreg_idx].slot_idx.is_none() {
+                        // TODO: we unconditionally save a vreg even if it is killed, move that logic down later
+                        self.alloc_stack_slot(vreg);
+                        self.move_to_stack(
+                            self.vregs[vreg_idx].preg.unwrap(),
+                            vreg,
+                            ProgPoint::before(cur_inst),
+                        );
+                    }
+
                     self.clear_preg_idx(i);
                 }
             }
@@ -680,7 +691,7 @@ fn allocate_block_insts<'a, F: Function>(
         let mut late_write_regs = PRegSet::empty();
 
         if req_refs_on_stack {
-            state.clear_reftype_vregs();
+            state.clear_reftype_vregs(inst);
             let pos = ProgPoint::before(inst);
             trace!("Calculating Stackmap for {:?}", pos);
 
@@ -994,8 +1005,8 @@ fn allocate_block_insts<'a, F: Function>(
                                 state.move_to_stack(tmp_reg, vreg, ProgPoint::after(inst));
                                 regs_allocated.add(tmp_reg);
                             } else {
-                                state.alloc_stack_slot(vreg);
+                                //state.alloc_stack_slot(vreg);
-                                state.move_to_stack(reg, vreg, ProgPoint::after(inst));
+                                //state.move_to_stack(reg, vreg, ProgPoint::after(inst));
                                 regs_allocated.add(reg);
                             }
                             trace!("Chose {} for operand {}", reg, i);
@@ -1011,7 +1022,19 @@ fn allocate_block_insts<'a, F: Function>(
                                     Allocation::stack(SpillSlot::new(*slot as usize));
                                 trace!("Chose slot {} for operand {}", slot, i);
                             } else {
-                                return Err(RegAllocError::SSA(vreg, inst));
+                                let cur_preg = state.vregs[vreg.vreg()].preg.unwrap();
+                                if op.constraint() == OperandConstraint::Any {
+                                    // TODO: have reg be a higher priority than stack?
+                                    state.allocs[alloc_idx + i] = Allocation::reg(cur_preg);
+                                    trace!("Chose {} for operand {}", cur_preg, i);
+                                } else {
+                                    // move to stack
+                                    let slot: u32 = state.alloc_stack_slot(vreg);
+                                    state.move_to_stack(cur_preg, vreg, ProgPoint::before(inst));
+                                    state.allocs[alloc_idx + i] =
+                                        Allocation::stack(SpillSlot::new(slot as usize));
+                                    trace!("Chose slot {} for operand {}", slot, i);
+                                }
                             }
                         }
                         OperandKind::Def => {
@@ -1097,14 +1120,29 @@ fn allocate_block_insts<'a, F: Function>(
                                 return Err(RegAllocError::TooManyLiveRegs);
                             }
 
-                            // need to move from stack to reg
+                            if let Some(cur_preg) = state.vregs[vreg.vreg()].preg {
-                            state.move_to_preg(vreg, reg, ProgPoint::before(inst));
+                                // TODO: make sure this is not used in another use?
+                                state.clear_preg(cur_preg);
+                                state.edits.push((
+                                    ProgPoint::before(inst),
+                                    Edit::Move {
+                                        from: Allocation::reg(cur_preg),
+                                        to: Allocation::reg(reg),
+                                    },
+                                ));
+                                debug_assert!(state.pregs[reg.index()].vreg.is_none());
+                                state.pregs[reg.index()].vreg = Some(vreg);
+                                state.vregs[vreg.vreg()].preg = Some(reg);
+                            } else {
+                                // need to move from stack to reg
+                                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.alloc_stack_slot(vreg);
-                            state.move_to_stack(reg, vreg, ProgPoint::after(inst));
+                            //state.move_to_stack(reg, vreg, ProgPoint::after(inst));
                         }
 
                         allocated = true;
@@ -1255,8 +1293,8 @@ fn allocate_block_insts<'a, F: Function>(
                             // early def
                             state.vregs[vreg.vreg()].def_block = Some(block);
                             state.assign_preg(preg, vreg);
-                            state.alloc_stack_slot(vreg);
+                            //state.alloc_stack_slot(vreg);
-                            state.move_to_stack(preg, vreg, ProgPoint::after(inst));
+                            //state.move_to_stack(preg, vreg, ProgPoint::after(inst));
                         }
 
                         trace!("Chose {} for operand {}", preg, i);
@@ -1370,8 +1408,8 @@ fn allocate_block_insts<'a, F: Function>(
 
                         state.clear_preg(reg);
                         state.assign_preg(reg, vreg);
-                        state.alloc_stack_slot(vreg);
+                        //state.alloc_stack_slot(vreg);
-                        state.move_to_stack(reg, vreg, ProgPoint::after(inst));
+                        //state.move_to_stack(reg, vreg, ProgPoint::after(inst));
                         trace!("Chose {} for operand {}", reg, i);
                         allocated = true;
                         break;
@@ -1446,8 +1484,8 @@ fn allocate_block_insts<'a, F: Function>(
                         state.allocs[alloc_idx + i] = Allocation::reg(reg);
 
                         state.assign_preg(reg, vreg);
-                        state.alloc_stack_slot(vreg);
+                        //state.alloc_stack_slot(vreg);
-                        state.move_to_stack(reg, vreg, ProgPoint::after(inst));
+                        //state.move_to_stack(reg, vreg, ProgPoint::after(inst));
                         trace!("Chose {} for operand {}", reg, i);
                     } else {
                         if ffa_reg_pool == PRegSet::empty() {
@@ -1500,8 +1538,8 @@ fn allocate_block_insts<'a, F: Function>(
                             // early def
                             state.vregs[vreg.vreg()].def_block = Some(block);
                             state.assign_preg(preg, vreg);
-                            state.alloc_stack_slot(vreg);
+                            //state.alloc_stack_slot(vreg);
-                            state.move_to_stack(preg, vreg, ProgPoint::after(inst));
+                            //state.move_to_stack(preg, vreg, ProgPoint::after(inst));
                         }
 
                         trace!("Chose {} for operand {}", preg, i);
@@ -1574,55 +1612,10 @@ fn allocate_block_insts<'a, F: Function>(
         state.cur_inst_pos += 1;
     }
 
-    // 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();
-                // 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((
-                        ProgPoint::before(Inst::new(block_last_inst_idx)),
-                        Edit::Move {
-                            from: Allocation::reg(preg),
-                            to: Allocation::stack(SpillSlot::new(slot as usize)),
-                        },
-                    ));
-                    state.vregs[idx].slot_idx = Some(slot);
-                }
-
-                state.vregs[idx].preg = None;
-                state.pregs[i].vreg = None;
-            }
-        }
-    }
-
     Ok(())
 }
 
-fn handle_out_block_params<'a, F: Function>(
+fn handle_out_block_params2<'a, F: Function>(
     state: &mut FastAllocState<'a, F>,
     const_state: &ReadOnlyData,
     block: Block,
@@ -1638,6 +1631,42 @@ fn handle_out_block_params<'a, F: Function>(
         return Ok(());
     }
 
+    let succs = state.func.block_succs(block);
+    if succs.len() == 1 {
+        let succ = succs[0];
+
+        if state.blocks[succ.index()].params_allocated {
+        } else {
+        }
+    }
+
+    todo!("")
+}
+
+fn handle_out_block_params<'a, F: Function>(
+    state: &mut FastAllocState<'a, F>,
+    const_state: &ReadOnlyData,
+    block: Block,
+) -> Result<(), RegAllocError> {
+    let mut 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());
+
+        // TODO: we need preg maps for each bblock but rn we need to clear the pregs so that the allocation of the next block is not mistaken
+        for i in 0..state.pregs.len() {
+            if state.pregs[i].vreg.is_some() {
+                state.clear_preg_idx(i);
+            }
+        }
+
+        return Ok(());
+    }
+
     let mut pregs_used_by_br = PRegSet::empty();
     {
         let alloc_start = state.inst_alloc_offsets[last_inst.index()] as usize;
@@ -1677,7 +1706,56 @@ fn handle_out_block_params<'a, F: Function>(
 
     let succs = state.func.block_succs(block);
     if succs.len() == 1 && state.blocks[succs[0].index()].params_allocated {
-        trace!("Only one allocated successor, moving allocations");
+        trace!("Only one allocated successor, moving allocations. Moving all liveouts to stack");
+        // done here as there are cases where there are multiple branches at the end of a bblock so we cannot insert any moves
+        // in the case of multiple successors
+        // 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();
+                    // 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((
+                            ProgPoint::before(Inst::new(last_inst.index())),
+                            Edit::Move {
+                                from: Allocation::reg(preg),
+                                to: Allocation::stack(SpillSlot::new(slot as usize)),
+                            },
+                        ));
+                        state.vregs[idx].slot_idx = Some(slot);
+                    }
+
+                    state.vregs[idx].preg = None;
+                    state.pregs[i].vreg = None;
+                }
+            }
+        }
+
         let succ = succs[0];
         // move values to the already allocated places
         let in_params = state.func.block_params(succ);
@@ -2110,14 +2188,75 @@ fn handle_out_block_params<'a, F: Function>(
     } else {
         trace!("Successors not allocated. Creating allocation");
 
-        let mut allocs = SmallVec::<[(VReg, u32); 4]>::new();
+        // 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();
+                    // 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());
+                        for succ in state.func.block_succs(block) {
+                            let insert_pos = if succs.len() == 1 {
+                                ProgPoint::before(last_inst)
+                            } else {
+                                ProgPoint::before(state.func.block_insns(*succ).first())
+                            };
+
+                            state.edits.push((
+                                insert_pos,
+                                Edit::Move {
+                                    from: Allocation::reg(preg),
+                                    to: Allocation::stack(SpillSlot::new(slot as usize)),
+                                },
+                            ));
+                        }
+                        state.vregs[idx].slot_idx = Some(slot);
+                    }
+
+                    state.vregs[idx].preg = None;
+                    state.pregs[i].vreg = None;
+                }
+            }
+        }
+
         // set incoming block params of successor to the current stack slot
         for (i, &succ) in state.func.block_succs(block).iter().enumerate() {
+            let mut allocs = SmallVec::<[(VReg, u32); 4]>::new();
             trace!("Creating block {}", succ.index());
             if state.blocks[succ.index()].params_allocated {
                 return Err(RegAllocError::CritEdge(block, succ));
             }
 
+            let insert_pos = if succs.len() == 1 {
+                ProgPoint::before(last_inst)
+            } else {
+                ProgPoint::before(state.func.block_insns(succ).first())
+            };
+
             // we allocate the params here
             // TODO: can there be a problem if the same successor occurs multiple times?
             state.blocks[succ.index()].params_allocated = true;
@@ -2132,6 +2271,7 @@ fn handle_out_block_params<'a, F: Function>(
                 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());
+                debug_assert!(state.vregs[out_vreg.vreg()].preg.is_none());
                 let out_slot_idx = state.vregs[out_vreg.vreg()].slot_idx.unwrap();
 
                 if out_vreg == VReg::invalid() {
@@ -2197,14 +2337,14 @@ fn handle_out_block_params<'a, F: Function>(
                     in_vreg
                 );
                 state.edits.push((
-                    ProgPoint::before(last_inst),
+                    insert_pos,
                     Edit::Move {
                         from: Allocation::stack(SpillSlot::new(out_slot_idx as usize)),
                         to: Allocation::reg(tmp_reg),
                     },
                 ));
                 state.edits.push((
-                    ProgPoint::before(last_inst),
+                    insert_pos,
                     Edit::Move {
                         from: Allocation::reg(tmp_reg),
                         to: Allocation::stack(SpillSlot::new(slot as usize)),