Use regular queue for identity queue. (#2588)

chriseth · georgwiese · web-flow · commit 28307fba8528 · 2025-03-31T17:49:03.000Z
This avoids cloning QueueItems and processes the items in a queue
instead of always looking at the "smallest" activated item first.

---------

Co-authored-by: Georg Wiese &lt;georgwiese@gmail.com&gt;
diff --git a/executor/src/witgen/jit/block_machine_processor.rs b/executor/src/witgen/jit/block_machine_processor.rs
@@ -512,29 +512,31 @@ main_binary::operation_id[3] = params[0];
 main_binary::A[3] = params[1];
 main_binary::B[3] = params[2];
 main_binary::operation_id[2] = main_binary::operation_id[3];
-main_binary::operation_id[1] = main_binary::operation_id[2];
-main_binary::operation_id[0] = main_binary::operation_id[1];
-main_binary::operation_id_next[-1] = main_binary::operation_id[0];
-call_var(9, -1, 0) = main_binary::operation_id_next[-1];
-main_binary::operation_id_next[0] = main_binary::operation_id[1];
-call_var(9, 0, 0) = main_binary::operation_id_next[0];
-main_binary::operation_id_next[1] = main_binary::operation_id[2];
-call_var(9, 1, 0) = main_binary::operation_id_next[1];
 2**0 * main_binary::A[2] + 2**24 * main_binary::A_byte[2] := main_binary::A[3];
-call_var(9, 2, 1) = main_binary::A_byte[2];
-2**0 * main_binary::A[1] + 2**16 * main_binary::A_byte[1] := main_binary::A[2];
-call_var(9, 1, 1) = main_binary::A_byte[1];
-2**0 * main_binary::A[0] + 2**8 * main_binary::A_byte[0] := main_binary::A[1];
-call_var(9, 0, 1) = main_binary::A_byte[0];
-main_binary::A_byte[-1] = main_binary::A[0];
-call_var(9, -1, 1) = main_binary::A_byte[-1];
 2**0 * main_binary::B[2] + 2**24 * main_binary::B_byte[2] := main_binary::B[3];
+main_binary::operation_id_next[2] = main_binary::operation_id[3];
+call_var(9, 2, 1) = main_binary::A_byte[2];
 call_var(9, 2, 2) = main_binary::B_byte[2];
+call_var(9, 2, 0) = main_binary::operation_id_next[2];
+main_binary::operation_id[1] = main_binary::operation_id[2];
+2**0 * main_binary::A[1] + 2**16 * main_binary::A_byte[1] := main_binary::A[2];
 2**0 * main_binary::B[1] + 2**16 * main_binary::B_byte[1] := main_binary::B[2];
-call_var(9, 1, 2) = main_binary::B_byte[1];
+main_binary::operation_id_next[1] = main_binary::operation_id[2];
+main_binary::operation_id[0] = main_binary::operation_id[1];
+main_binary::operation_id_next[0] = main_binary::operation_id[1];
+2**0 * main_binary::A[0] + 2**8 * main_binary::A_byte[0] := main_binary::A[1];
+call_var(9, 1, 1) = main_binary::A_byte[1];
 2**0 * main_binary::B[0] + 2**8 * main_binary::B_byte[0] := main_binary::B[1];
-call_var(9, 0, 2) = main_binary::B_byte[0];
+call_var(9, 1, 2) = main_binary::B_byte[1];
+call_var(9, 1, 0) = main_binary::operation_id_next[1];
+main_binary::operation_id_next[-1] = main_binary::operation_id[0];
+call_var(9, 0, 0) = main_binary::operation_id_next[0];
+main_binary::A_byte[-1] = main_binary::A[0];
+call_var(9, 0, 1) = main_binary::A_byte[0];
 main_binary::B_byte[-1] = main_binary::B[0];
+call_var(9, 0, 2) = main_binary::B_byte[0];
+call_var(9, -1, 0) = main_binary::operation_id_next[-1];
+call_var(9, -1, 1) = main_binary::A_byte[-1];
 call_var(9, -1, 2) = main_binary::B_byte[-1];
 machine_call(2, [Known(call_var(9, -1, 0)), Known(call_var(9, -1, 1)), Known(call_var(9, -1, 2)), Unknown(call_var(9, -1, 3))]);
 main_binary::C_byte[-1] = call_var(9, -1, 3);
@@ -545,8 +547,6 @@ main_binary::C[1] = (main_binary::C[0] + (main_binary::C_byte[0] * 256));
 machine_call(2, [Known(call_var(9, 1, 0)), Known(call_var(9, 1, 1)), Known(call_var(9, 1, 2)), Unknown(call_var(9, 1, 3))]);
 main_binary::C_byte[1] = call_var(9, 1, 3);
 main_binary::C[2] = (main_binary::C[1] + (main_binary::C_byte[1] * 65536));
-main_binary::operation_id_next[2] = main_binary::operation_id[3];
-call_var(9, 2, 0) = main_binary::operation_id_next[2];
 machine_call(2, [Known(call_var(9, 2, 0)), Known(call_var(9, 2, 1)), Known(call_var(9, 2, 2)), Unknown(call_var(9, 2, 3))]);
 main_binary::C_byte[2] = call_var(9, 2, 3);
 main_binary::C[3] = (main_binary::C[2] + (main_binary::C_byte[2] * 16777216));
@@ -611,12 +611,12 @@ params[1] = Sub::b[0];"
 params[1] = SubM::b[0];
 params[2] = SubM::c[0];
 call_var(1, 0, 0) = SubM::c[0];
-machine_call(2, [Known(call_var(1, 0, 0))]);
+SubM::c[1] = SubM::c[0];
 SubM::b[1] = SubM::b[0];
 call_var(1, 1, 0) = SubM::b[1];
-SubM::c[1] = SubM::c[0];
-machine_call(2, [Known(call_var(1, 1, 0))]);
-SubM::a[1] = ((SubM::b[1] * 256) + SubM::c[1]);"
+SubM::a[1] = ((SubM::b[1] * 256) + SubM::c[1]);
+machine_call(2, [Known(call_var(1, 0, 0))]);
+machine_call(2, [Known(call_var(1, 1, 0))]);"
         );
     }
 
@@ -670,15 +670,15 @@ machine_call(3, [Known(call_var(3, 0, 0))]);"
             code,
             "S::a[0] = params[0];
 S::b[0] = 0;
-params[1] = 0;
 S::b[1] = 0;
 S::c[0] = 1;
-params[2] = 1;
 S::b[2] = 0;
 S::c[1] = 1;
 S::b[3] = 8;
 S::c[2] = 1;
-S::c[3] = 9;"
+S::c[3] = 9;
+params[1] = 0;
+params[2] = 1;"
         );
     }
 
diff --git a/executor/src/witgen/jit/identity_queue.rs b/executor/src/witgen/jit/identity_queue.rs
@@ -1,5 +1,5 @@
 use std::{
-    collections::{BTreeSet, HashMap},
+    collections::{BTreeSet, HashMap, VecDeque},
     rc::Rc,
 };
 
@@ -20,50 +20,116 @@ use super::{prover_function_heuristics::ProverFunction, variable::Variable};
 /// updates this list based on the occurrence of updated variables
 /// in identities.
 #[derive(Clone)]
-pub struct IdentityQueue<'a, T: FieldElement> {
-    queue: BTreeSet<QueueItem<'a, T>>,
-    occurrences: Rc<HashMap<Variable, Vec<QueueItem<'a, T>>>>,
+pub struct IdentityQueue<'ast, 'queue, T: FieldElement> {
+    items: &'queue Vec<QueueItem<'ast, T>>,
+    in_queue: Vec<bool>,
+    identity_queue: VecDeque<usize>,
+    /// This is a priority queue because we always want to process
+    /// the "first" machine call that received a variable update
+    /// in the hope that this results in the machine calls to be
+    /// in source order, where possible.
+    machine_call_queue: BTreeSet<usize>,
+    prover_function_queue: VecDeque<usize>,
+    /// Maps a variable to a list of indices in `items`, pointing to the items where they are referenced.
+    occurrences: Rc<HashMap<Variable, Vec<usize>>>,
 }
 
-impl<'a, T: FieldElement> IdentityQueue<'a, T> {
-    pub fn new(items: impl IntoIterator<Item = QueueItem<'a, T>>) -> Self {
-        let queue: BTreeSet<_> = items.into_iter().collect();
+impl<'ast, 'queue, T: FieldElement> IdentityQueue<'ast, 'queue, T> {
+    /// Creates a new queue based on the given identities.
+    /// The order of identities in this queue matters for the
+    /// order in which they are processed.
+    pub fn new(items: &'queue Vec<QueueItem<'ast, T>>) -> Self {
         let mut references = ReferencesComputer::default();
         let occurrences = Rc::new(
-            queue
+            items
                 .iter()
-                .flat_map(|item| {
+                .enumerate()
+                .flat_map(|(id, item)| {
                     references
                         .references(item)
                         .iter()
-                        .map(|v| (v.clone(), item.clone()))
+                        .map(|v| (v.clone(), id))
                         .collect_vec()
                 })
                 .into_group_map(),
         );
-        Self { queue, occurrences }
+        Self {
+            items,
+            in_queue: vec![true; items.len()],
+            identity_queue: collect_filtered_indices(items, &is_polynomial_identity_or_assignment),
+            machine_call_queue: collect_filtered_indices(items, &is_submachine_call),
+            prover_function_queue: collect_filtered_indices(items, &is_prover_function),
+            occurrences,
+        }
     }
 
     /// Returns the next identity to be processed and its row and
     /// removes it from the queue.
-    pub fn next(&mut self) -> Option<QueueItem<'a, T>> {
-        self.queue.pop_first()
+    pub fn next(&mut self) -> Option<&'queue QueueItem<'ast, T>> {
+        self.identity_queue
+            .pop_front()
+            .or_else(|| self.machine_call_queue.pop_first())
+            .or_else(|| self.prover_function_queue.pop_front())
+            .map(|id| {
+                self.in_queue[id] = false;
+                &self.items[id]
+            })
     }
 
     pub fn variables_updated(&mut self, variables: impl IntoIterator<Item = Variable>) {
         // Note that this will usually re-add the item that caused the update,
         // which is fine, since there are situations where we can further process
         // it from an update (for example a range constraint).
-        self.queue.extend(
-            variables
-                .into_iter()
-                .flat_map(|var| self.occurrences.get(&var))
-                .flatten()
-                .cloned(),
-        )
+        for id in variables
+            .into_iter()
+            .flat_map(|var| self.occurrences.get(&var))
+            .flatten()
+        {
+            if !self.in_queue[*id] {
+                self.in_queue[*id] = true;
+                if is_polynomial_identity_or_assignment(&self.items[*id]) {
+                    self.identity_queue.push_back(*id);
+                } else if is_submachine_call(&self.items[*id]) {
+                    self.machine_call_queue.insert(*id);
+                } else {
+                    assert!(is_prover_function(&self.items[*id]));
+                    self.prover_function_queue.push_back(*id);
+                }
+            }
+        }
     }
 }
 
+fn is_polynomial_identity_or_assignment<T: FieldElement>(item: &QueueItem<'_, T>) -> bool {
+    match item {
+        QueueItem::Identity(Identity::Polynomial(..), _)
+        | QueueItem::VariableAssignment(..)
+        | QueueItem::ConstantAssignment(..) => true,
+        QueueItem::Identity(Identity::BusSend(..), _) | QueueItem::ProverFunction(..) => false,
+        QueueItem::Identity(Identity::Connect(..), _) => unreachable!(),
+    }
+}
+
+fn is_submachine_call<T: FieldElement>(item: &QueueItem<'_, T>) -> bool {
+    matches!(item, QueueItem::Identity(Identity::BusSend(..), _))
+}
+
+fn is_prover_function<T: FieldElement>(item: &QueueItem<'_, T>) -> bool {
+    matches!(item, QueueItem::ProverFunction(..))
+}
+
+/// Filters a slice by a boolean selector and returns a collection of the matching indices.
+fn collect_filtered_indices<D, F, C: FromIterator<usize>>(items: &[D], mut filter: F) -> C
+where
+    F: FnMut(&D) -> bool,
+{
+    items
+        .iter()
+        .enumerate()
+        .filter_map(|(index, item)| if filter(item) { Some(index) } else { None })
+        .collect()
+}
+
 #[derive(Clone)]
 pub enum QueueItem<'a, T: FieldElement> {
     Identity(&'a Identity<T>, i32),
diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs
@@ -99,15 +99,17 @@ impl<'a, T: FieldElement> Processor<'a, T> {
             }
         }));
         let branch_depth = 0;
-        let identity_queue = IdentityQueue::new(queue_items);
+        // Sort the queue so that we have proper source order.
+        queue_items.sort();
+        let identity_queue = IdentityQueue::new(&queue_items);
         self.generate_code_for_branch(can_process, witgen, identity_queue, branch_depth)
     }
 
     fn generate_code_for_branch<FixedEval: FixedEvaluator<T>>(
         &self,
         can_process: impl CanProcessCall<T>,
         mut witgen: WitgenInference<'a, T, FixedEval>,
-        mut identity_queue: IdentityQueue<'a, T>,
+        mut identity_queue: IdentityQueue<'a, '_, T>,
         branch_depth: usize,
     ) -> Result<ProcessorResult<T>, Error<'a, T, FixedEval>> {
         if self
@@ -267,16 +269,16 @@ impl<'a, T: FieldElement> Processor<'a, T> {
         &self,
         can_process: impl CanProcessCall<T>,
         witgen: &mut WitgenInference<'a, T, FixedEval>,
-        mut identity_queue: IdentityQueue<'a, T>,
+        mut identity_queue: IdentityQueue<'a, '_, T>,
     ) -> Result<(), affine_symbolic_expression::Error> {
         while let Some(item) = identity_queue.next() {
             let updated_vars = match item {
                 QueueItem::Identity(identity, row_offset) => match identity {
                     Identity::Polynomial(PolynomialIdentity { expression, .. }) => {
-                        witgen.process_equation_on_row(expression, None, 0.into(), row_offset)
+                        witgen.process_equation_on_row(expression, None, 0.into(), *row_offset)
                     }
                     Identity::BusSend(bus_send) => {
-                        witgen.process_call(can_process.clone(), bus_send, row_offset)
+                        witgen.process_call(can_process.clone(), bus_send, *row_offset)
                     }
                     Identity::Connect(..) => Ok(vec![]),
                 },
@@ -293,7 +295,7 @@ impl<'a, T: FieldElement> Processor<'a, T> {
                     assignment.row_offset,
                 ),
                 QueueItem::ProverFunction(prover_function, row_offset) => {
-                    witgen.process_prover_function(&prover_function, row_offset)
+                    witgen.process_prover_function(prover_function, *row_offset)
                 }
             }?;
             identity_queue.variables_updated(updated_vars);
@@ -358,7 +360,7 @@ impl<'a, T: FieldElement> Processor<'a, T> {
         incomplete_machine_calls: &[(&Identity<T>, i32)],
         can_process: impl CanProcessCall<T>,
         witgen: &mut WitgenInference<'a, T, FixedEval>,
-        mut identity_queue: IdentityQueue<'a, T>,
+        mut identity_queue: IdentityQueue<'a, '_, T>,
     ) -> bool {
         let missing_sends_in_block = incomplete_machine_calls
             .iter()
diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs
@@ -263,8 +263,8 @@ call_var(1, 0, 1) = VM::instr_add[0];
 call_var(1, 0, 2) = VM::instr_mul[0];
 call_var(1, 0, 0) = VM::pc[0];
 VM::pc[1] = (VM::pc[0] + 1);
-call_var(1, 1, 0) = VM::pc[1];
 VM::B[1] = VM::B[0];
+call_var(1, 1, 0) = VM::pc[1];
 machine_call(1, [Known(call_var(1, 1, 0)), Unknown(call_var(1, 1, 1)), Unknown(call_var(1, 1, 2))]);
 VM::instr_add[1] = call_var(1, 1, 1);
 VM::instr_mul[1] = call_var(1, 1, 2);
@@ -306,9 +306,9 @@ call_var(2, 0, 1) = VM::instr_add[0];
 call_var(2, 0, 2) = VM::instr_mul[0];
 call_var(2, 0, 0) = VM::pc[0];
 VM::pc[1] = VM::pc[0];
-call_var(2, 1, 0) = VM::pc[1];
 VM::instr_add[1] = 0;
 call_var(2, 1, 1) = 0;
+call_var(2, 1, 0) = VM::pc[1];
 call_var(2, 1, 2) = 1;
 machine_call(1, [Known(call_var(2, 1, 0)), Known(call_var(2, 1, 1)), Unknown(call_var(2, 1, 2))]);
 VM::instr_mul[1] = 1;"
