From d31dbaaa16cd01c0c50ca581a1bc26106eb84f90 Mon Sep 17 00:00:00 2001
From: T0b1 <T0b1-iOS@users.noreply.github.com>
Date: Sun, 16 Apr 2023 02:03:50 +0200
Subject: [PATCH] calculate use positions

---
 src/ion/fast_alloc.rs | 56 ++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 53 insertions(+), 3 deletions(-)

diff --git a/src/ion/fast_alloc.rs b/src/ion/fast_alloc.rs
index 8a5b389..bfb0234 100644
--- a/src/ion/fast_alloc.rs
+++ b/src/ion/fast_alloc.rs
@@ -14,12 +14,16 @@ use crate::{domtree, postorder, InstPosition};
 use super::data_structures::u64_key;
 use super::Stats;
 
-#[derive(Default, Clone, Copy)]
+#[derive(Default, Clone)]
 struct VRegData {
     pub preg: Option<PReg>,
     pub slot_idx: Option<u32>,
     pub def_block: Option<Block>,
     pub reftype: bool,
+
+    // use information
+    pub cur_use_idx: u32,
+    pub uses: SmallVec<[u32; 8]>,
 }
 
 #[derive(Default, Clone, Copy)]
@@ -395,18 +399,20 @@ pub fn run<F: Function>(func: &F, mach_env: &MachineEnv) -> Result<Output, RegAl
     let mut state = FastAllocState::init(func, mach_env);
     let const_state = ReadOnlyData::init(func, mach_env);
 
+    calc_use_positions(&mut state, &const_state);
+
     state.blocks[func.entry_block().index()].params_allocated = true;
 
     let len = const_state.postorder.len();
     for i in 0..len {
         let block = const_state.postorder[len - 1 - i];
         if state.blocks[block.index()].reg_allocated {
-            trace!("Block {} already allocated. Skipping", i);
+            trace!("Block {} already allocated. Skipping", block.index());
             continue;
         }
         state.blocks[block.index()].reg_allocated = true;
 
-        trace!("Allocating block {}", i);
+        trace!("Allocating block {}", block.index());
 
         allocate_block_insts(&mut state, &const_state, block)?;
         handle_out_block_params(&mut state, &const_state, block)?;
@@ -1023,3 +1029,47 @@ fn handle_out_block_params<'a, F: Function>(
 
     Ok(())
 }
+
+fn calc_use_positions<'a, F: Function>(
+    state: &mut FastAllocState<'a, F>,
+    const_state: &ReadOnlyData,
+) {
+    // we use a pseudo-counter to have a uniform position for instructions
+    let mut cur_pos = 0u32;
+    let len = const_state.postorder.len();
+    for i in 0..len {
+        let block = const_state.postorder[len - 1 - i];
+        trace!("Calculating uses for block {}", block.index());
+
+        let insts = state.func.block_insns(block);
+        for inst in insts.clone().iter() {
+            let operands = state.func.inst_operands(inst);
+            for op in operands {
+                if op.kind() != OperandKind::Use {
+                    continue;
+                }
+
+                if op.vreg() == VReg::invalid() {
+                    continue;
+                }
+
+                state.vregs[op.vreg().vreg()].uses.push(cur_pos);
+            }
+
+            cur_pos += 1;
+        }
+
+        let last_inst = insts.last();
+        if !state.func.is_branch(last_inst) {
+            continue;
+        }
+
+        for i in 0..state.func.block_succs(block).len() {
+            for vreg in state.func.branch_blockparams(block, last_inst, i) {
+                state.vregs[vreg.vreg()].uses.push(cur_pos);
+            }
+        }
+
+        cur_pos += 1;
+    }
+}