ZJIT: Use worklist algorithm for type inference

Replace the RPO fixed-point iteration with a worklist-based approach.
Instead of re-scanning every reachable block on each iteration, only
re-process blocks whose incoming types actually changed. This reduces
compile time on lobsters by ~11% (1,604ms vs 1,808ms).

Author: tekknolagi

zjit/src/bitset.rs

@@ -37,6 +37,16 @@ impl<T: Into<usize> + Copy> BitSet<T> {
         }
     }
 
+    /// Clear a bit. Returns whether the bit was previously set.
+    pub fn remove(&mut self, idx: T) -> bool {
+        debug_assert!(idx.into() < self.num_bits);
+        let entry_idx = idx.into() / ENTRY_NUM_BITS;
+        let bit_idx = idx.into() % ENTRY_NUM_BITS;
+        let was_set = (self.entries[entry_idx] & (1 << bit_idx)) != 0;
+        self.entries[entry_idx] &= !(1 << bit_idx);
+        was_set
+    }
+
     pub fn get(&self, idx: T) -> bool {
         debug_assert!(idx.into() < self.num_bits);
         let entry_idx = idx.into() / ENTRY_NUM_BITS;

zjit/src/hir.rs

@@ -2603,60 +2603,74 @@ impl Function {
         self.copy_param_types();
 
         let mut reachable = BlockSet::with_capacity(self.blocks.len());
+
+        // Maintain both a worklist and a fast membership check to avoid linear search
+        let mut worklist: VecDeque<BlockId> = VecDeque::new();
+        let mut in_worklist = BlockSet::with_capacity(self.blocks.len());
+        macro_rules! worklist_add {
+            ($block:expr) => {
+                if in_worklist.insert($block) {
+                    worklist.push_back($block);
+                }
+            };
+        }
+
         for entry_block in self.entry_blocks() {
             reachable.insert(entry_block);
+            worklist_add!(entry_block);
+        }
+
+        // Helper to propagate types along a branch edge and enqueue the target if anything changed
+        macro_rules! enqueue {
+            ($self:ident, $target:expr) => {
+                let newly_reachable = reachable.insert($target.target);
+                let mut target_changed = newly_reachable;
+                for (idx, arg) in $target.args.iter().enumerate() {
+                    let param = $self.blocks[$target.target.0].params[idx];
+                    let new = self.insn_types[param.0].union(self.insn_types[arg.0]);
+                    if !self.insn_types[param.0].bit_equal(new) {
+                        self.insn_types[param.0] = new;
+                        target_changed = true;
+                    }
+                }
+                if target_changed {
+                    worklist_add!($target.target);
+                }
+            };
         }
 
-        // Walk the graph, computing types until fixpoint
-        let rpo = self.rpo();
-        loop {
-            let mut changed = false;
-            for &block in &rpo {
-                if !reachable.get(block) { continue; }
-                for insn_id in &self.blocks[block.0].insns {
-                    let insn_type = match self.find(*insn_id) {
-                        Insn::IfTrue { val, target: BranchEdge { target, args } } => {
-                            assert!(!self.type_of(val).bit_equal(types::Empty));
-                            if self.type_of(val).could_be(Type::from_cbool(true)) {
-                                reachable.insert(target);
-                                for (idx, arg) in args.iter().enumerate() {
-                                    let param = self.blocks[target.0].params[idx];
-                                    self.insn_types[param.0] = self.type_of(param).union(self.type_of(*arg));
-                                }
-                            }
-                            continue;
-                        }
-                        Insn::IfFalse { val, target: BranchEdge { target, args } } => {
-                            assert!(!self.type_of(val).bit_equal(types::Empty));
-                            if self.type_of(val).could_be(Type::from_cbool(false)) {
-                                reachable.insert(target);
-                                for (idx, arg) in args.iter().enumerate() {
-                                    let param = self.blocks[target.0].params[idx];
-                                    self.insn_types[param.0] = self.type_of(param).union(self.type_of(*arg));
-                                }
-                            }
-                            continue;
+        // Walk the graph, computing types until worklist is empty
+        while let Some(block) = worklist.pop_front() {
+            in_worklist.remove(block);
+            if !reachable.get(block) { continue; }
+
+            for insn_id in &self.blocks[block.0].insns {
+                let insn_type = match self.find(*insn_id) {
+                    Insn::IfTrue { val, target } => {
+                        assert!(!self.type_of(val).bit_equal(types::Empty));
+                        if self.type_of(val).could_be(Type::from_cbool(true)) {
+                            enqueue!(self, target);
                         }
-                        Insn::Jump(BranchEdge { target, args }) => {
-                            reachable.insert(target);
-                            for (idx, arg) in args.iter().enumerate() {
-                                let param = self.blocks[target.0].params[idx];
-                                self.insn_types[param.0] = self.type_of(param).union(self.type_of(*arg));
-                            }
-                            continue;
+                        continue;
+                    }
+                    Insn::IfFalse { val, target } => {
+                        assert!(!self.type_of(val).bit_equal(types::Empty));
+                        if self.type_of(val).could_be(Type::from_cbool(false)) {
+                            enqueue!(self, target);
                         }
-                        insn if insn.has_output() => self.infer_type(*insn_id),
-                        _ => continue,
-                    };
-                    if !self.type_of(*insn_id).bit_equal(insn_type) {
-                        self.insn_types[insn_id.0] = insn_type;
-                        changed = true;
+                        continue;
                     }
+                    Insn::Jump(target) => {
+                        enqueue!(self, target);
+                        continue;
+                    }
+                    insn if insn.has_output() => self.infer_type(*insn_id),
+                    _ => continue,
+                };
+                if !self.type_of(*insn_id).bit_equal(insn_type) {
+                    self.insn_types[insn_id.0] = insn_type;
                 }
             }
-            if !changed {
-                break;
-            }
         }
     }