From 9f43e3bb925534f613cb2089df6e5fe315fc851a Mon Sep 17 00:00:00 2001
From: Jakub Doka <jakub.doka2@gmail.com>
Date: Sun, 15 Dec 2024 20:49:37 +0100
Subject: [PATCH] refactoring some stuff and loosening a requrement on assert

Signed-off-by: Jakub Doka <jakub.doka2@gmail.com>
---
 lang/src/backend/hbvm/regalloc.rs | 48 ++++++++++++++++---------------
 1 file changed, 25 insertions(+), 23 deletions(-)

diff --git a/lang/src/backend/hbvm/regalloc.rs b/lang/src/backend/hbvm/regalloc.rs
index ce67e1d39..89eccf839 100644
--- a/lang/src/backend/hbvm/regalloc.rs
+++ b/lang/src/backend/hbvm/regalloc.rs
@@ -709,7 +709,7 @@ impl Nodes {
                 );
                 debug_assert_matches!(
                     self[self[user].inputs[3]].kind,
-                    Kind::Stre | Kind::Mem | Kind::Phi
+                    Kind::Stre | Kind::Mem | Kind::Phi | Kind::Join
                 );
                 self[user].inputs.iter().position(|&v| v == val).is_some_and(|v| v < 3)
             }
@@ -749,31 +749,28 @@ impl Nodes {
         nid
     }
 
-    fn uses_of<'a>(
-        &'a self,
-        nid: Nid,
-        types: &'a Types,
-    ) -> impl Iterator<Item = (Nid, Nid)> + use<'a> {
+    fn uses_of(&self, nid: Nid, types: &Types, stack: &mut Vec<Nid>, buf: &mut Vec<(Nid, Nid)>) {
+        debug_assert!(stack.is_empty());
+        debug_assert!(buf.is_empty());
+
         if self[nid].kind.is_cfg() && !matches!(self[nid].kind, Kind::Call { .. }) {
-            return None.into_iter().flatten();
+            return;
         }
 
-        let mut uses = vec![];
-        let mut to_expand = vec![nid];
-        while let Some(exp) = to_expand.pop() {
+        stack.push(nid);
+
+        while let Some(exp) = stack.pop() {
             for &o in self[exp].outputs.iter() {
                 if !self.is_data_dep(exp, o, types) {
                     continue;
                 }
                 if self.is_unlocked(o) {
-                    uses.push((self.use_block_of(exp, o), o));
+                    buf.push((self.use_block_of(exp, o), o));
                 } else {
-                    to_expand.push(o);
+                    stack.push(o);
                 }
             }
         }
-
-        Some(uses).into_iter().flatten()
     }
 }
 
@@ -810,13 +807,14 @@ impl<'a> Regalloc<'a> {
 
         debug_assert!(self.res.dfs_buf.is_empty());
 
+        let mut uses_buf = Vec::new();
         let mut bundle = Bundle::new(self.res.instrs.len());
         self.res.visited.clear(self.nodes.len());
 
         for i in (0..self.res.blocks.len()).rev() {
             for [a, rest @ ..] in self.nodes.phi_inputs_of(self.res.blocks[i].entry) {
                 if self.res.visited.set(a) {
-                    self.append_bundle(a, &mut bundle, None);
+                    self.append_bundle(a, &mut bundle, None, &mut uses_buf);
                 }
 
                 for r in rest {
@@ -828,6 +826,7 @@ impl<'a> Regalloc<'a> {
                         r,
                         &mut bundle,
                         Some(self.res.node_to_reg[a as usize] as usize - 1),
+                        &mut uses_buf,
                     );
                 }
             }
@@ -838,15 +837,22 @@ impl<'a> Regalloc<'a> {
             if self.nodes[inst].has_no_value() || self.res.visited.get(inst) || inst == 0 {
                 continue;
             }
-            self.append_bundle(inst, &mut bundle, None);
+            self.append_bundle(inst, &mut bundle, None, &mut uses_buf);
         }
         self.res.instrs = instrs;
     }
 
-    fn collect_bundle(&mut self, inst: Nid, into: &mut Bundle) {
+    fn append_bundle(
+        &mut self,
+        inst: Nid,
+        tmp: &mut Bundle,
+        prefered: Option<usize>,
+        uses_buf: &mut Vec<(Nid, Nid)>,
+    ) {
         let dom = self.nodes.idom_of(inst);
         self.res.dfs_seem.clear(self.nodes.len());
-        for (cursor, uinst) in self.nodes.uses_of(inst, self.tys) {
+        self.nodes.uses_of(inst, self.tys, &mut self.res.dfs_buf, uses_buf);
+        for (cursor, uinst) in uses_buf.drain(..) {
             if !self.res.dfs_seem.set(uinst) {
                 continue;
             }
@@ -876,13 +882,9 @@ impl<'a> Regalloc<'a> {
                 range.end = new;
                 debug_assert!(range.start < range.end, "{:?} {inst} {uinst}", range);
 
-                into.add(range);
+                tmp.add(range);
             });
         }
-    }
-
-    fn append_bundle(&mut self, inst: Nid, tmp: &mut Bundle, prefered: Option<usize>) {
-        self.collect_bundle(inst, tmp);
 
         if tmp.is_empty() {
             self.res.node_to_reg[inst as usize] = u8::MAX;