Handle assignments in identity queue as well (#2453)

chriseth · web-flow · commit f198341ae819 · 2025-02-07T11:02:09.000Z
diff --git a/ast/src/analyzed/mod.rs b/ast/src/analyzed/mod.rs
@@ -1209,6 +1209,16 @@ impl From<&AlgebraicReference> for AlgebraicReferenceThin {
     }
 }
 
+impl AlgebraicReferenceThin {
+    pub fn with_name(&self, name: String) -> AlgebraicReference {
+        AlgebraicReference {
+            name,
+            poly_id: self.poly_id,
+            next: self.next,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Eq, Serialize, Deserialize, JsonSchema)]
 pub struct AlgebraicReference {
     /// Name of the polynomial - just for informational purposes.
diff --git a/executor/src/witgen/jit/identity_queue.rs b/executor/src/witgen/jit/identity_queue.rs
@@ -5,17 +5,17 @@ use std::{
 
 use itertools::Itertools;
 use powdr_ast::{
-    analyzed::{
-        AlgebraicExpression as Expression, AlgebraicReference, AlgebraicReferenceThin,
-        PolynomialType,
-    },
+    analyzed::{AlgebraicExpression as Expression, AlgebraicReferenceThin, PolynomialType},
     parsed::visitor::{AllChildren, Children},
 };
 use powdr_number::FieldElement;
 
 use crate::witgen::{data_structures::identity::Identity, FixedData};
 
-use super::variable::Variable;
+use super::{
+    variable::Variable,
+    witgen_inference::{Assignment, VariableOrValue},
+};
 
 /// Keeps track of identities that still need to be processed and
 /// updates this list based on the occurrence of updated variables
@@ -27,98 +27,120 @@ pub struct IdentityQueue<'a, T: FieldElement> {
 }
 
 impl<'a, T: FieldElement> IdentityQueue<'a, T> {
-    pub fn new(fixed_data: &'a FixedData<'a, T>, identities: &[(&'a Identity<T>, i32)]) -> Self {
-        let occurrences = compute_occurrences_map(fixed_data, identities).into();
-        Self {
-            queue: identities
-                .iter()
-                .map(|(id, row)| QueueItem(id, *row))
-                .collect(),
-            occurrences,
-        }
+    pub fn new(
+        fixed_data: &'a FixedData<'a, T>,
+        identities: &[(&'a Identity<T>, i32)],
+        assignments: &[Assignment<'a, T>],
+    ) -> Self {
+        let queue: BTreeSet<_> = identities
+            .iter()
+            .map(|(id, row)| QueueItem::Identity(id, *row))
+            .chain(assignments.iter().map(|a| QueueItem::Assignment(a.clone())))
+            .collect();
+        let occurrences = compute_occurrences_map(fixed_data, &queue).into();
+        Self { queue, occurrences }
     }
 
     /// Returns the next identity to be processed and its row and
     /// removes it from the queue.
-    pub fn next(&mut self) -> Option<(&'a Identity<T>, i32)> {
-        self.queue.pop_first().map(|QueueItem(id, row)| (id, row))
+    pub fn next(&mut self) -> Option<QueueItem<'a, T>> {
+        self.queue.pop_first()
     }
 
     pub fn variables_updated(
         &mut self,
         variables: impl IntoIterator<Item = Variable>,
-        skip_identity: Option<(&'a Identity<T>, i32)>,
+        skip_item: Option<QueueItem<'a, T>>,
     ) {
         self.queue.extend(
             variables
                 .into_iter()
                 .flat_map(|var| self.occurrences.get(&var))
                 .flatten()
-                .filter(|QueueItem(id, row)| match skip_identity {
-                    Some((id2, row2)) => (id.id(), *row) != (id2.id(), row2),
+                .filter(|item| match &skip_item {
+                    Some(it) => *item != it,
                     None => true,
-                }),
+                })
+                .cloned(),
         )
     }
 }
 
-/// Sorts identities by row and then by ID.
-#[derive(Clone, Copy)]
-struct QueueItem<'a, T>(&'a Identity<T>, i32);
-
-impl<T> QueueItem<'_, T> {
-    fn key(&self) -> (i32, u64) {
-        let QueueItem(id, row) = self;
-        (*row, id.id())
-    }
+#[derive(Clone)]
+pub enum QueueItem<'a, T: FieldElement> {
+    Identity(&'a Identity<T>, i32),
+    Assignment(Assignment<'a, T>),
 }
 
-impl<T> Ord for QueueItem<'_, T> {
+/// Sorts identities by row and then by ID, preceded by assignments.
+impl<T: FieldElement> Ord for QueueItem<'_, T> {
     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-        self.key().cmp(&other.key())
+        match (self, other) {
+            (QueueItem::Identity(id1, row1), QueueItem::Identity(id2, row2)) => {
+                (row1, id1.id()).cmp(&(row2, id2.id()))
+            }
+            (QueueItem::Assignment(a1), QueueItem::Assignment(a2)) => a1.cmp(a2),
+            (QueueItem::Assignment(_), QueueItem::Identity(_, _)) => std::cmp::Ordering::Less,
+            (QueueItem::Identity(_, _), QueueItem::Assignment(_)) => std::cmp::Ordering::Greater,
+        }
     }
 }
 
-impl<T> PartialOrd for QueueItem<'_, T> {
+impl<T: FieldElement> PartialOrd for QueueItem<'_, T> {
     fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
         Some(self.cmp(other))
     }
 }
 
-impl<T> PartialEq for QueueItem<'_, T> {
+impl<T: FieldElement> PartialEq for QueueItem<'_, T> {
     fn eq(&self, other: &Self) -> bool {
-        self.key() == other.key()
+        self.cmp(other) == std::cmp::Ordering::Equal
     }
 }
 
-impl<T> Eq for QueueItem<'_, T> {}
+impl<T: FieldElement> Eq for QueueItem<'_, T> {}
 
-/// Computes a map from each variable to the identity-row-offset pairs it occurs in.
-fn compute_occurrences_map<'a, T: FieldElement>(
+/// Computes a map from each variable to the queue items it occurs in.
+fn compute_occurrences_map<'b, 'a: 'b, T: FieldElement>(
     fixed_data: &'a FixedData<'a, T>,
-    identities: &[(&'a Identity<T>, i32)],
+    items: &BTreeSet<QueueItem<'a, T>>,
 ) -> HashMap<Variable, Vec<QueueItem<'a, T>>> {
-    let mut references_per_identity = HashMap::new();
     let mut intermediate_cache = HashMap::new();
-    for id in identities.iter().map(|(id, _)| *id).unique_by(|id| id.id()) {
+
+    // Compute references only once per identity.
+    let mut references_per_identity = HashMap::new();
+    for id in items
+        .iter()
+        .filter_map(|item| match item {
+            QueueItem::Identity(id, _) => Some(id),
+            _ => None,
+        })
+        .unique_by(|id| id.id())
+    {
         references_per_identity.insert(
-            id,
+            id.id(),
             references_in_identity(id, fixed_data, &mut intermediate_cache),
         );
     }
-    identities
+
+    items
         .iter()
-        .flat_map(|(id, row)| {
-            references_per_identity[id].iter().map(move |reference| {
-                let name = fixed_data.column_name(&reference.poly_id).to_string();
-                let fat_ref = AlgebraicReference {
-                    name,
-                    poly_id: reference.poly_id,
-                    next: reference.next,
-                };
-                let var = Variable::from_reference(&fat_ref, *row);
-                (var, QueueItem(*id, *row))
-            })
+        .flat_map(|item| {
+            let variables = match item {
+                QueueItem::Identity(id, row) => {
+                    references_in_identity(id, fixed_data, &mut intermediate_cache)
+                        .into_iter()
+                        .map(|r| {
+                            let name = fixed_data.column_name(&r.poly_id).to_string();
+                            Variable::from_reference(&r.with_name(name), *row)
+                        })
+                        .collect_vec()
+                }
+                QueueItem::Assignment(a) => {
+                    variables_in_assignment(a, fixed_data, &mut intermediate_cache)
+                }
+            };
+            variables.into_iter().map(move |v| (v, item.clone()))
         })
         .into_group_map()
 }
@@ -184,3 +206,22 @@ fn references_in_expression<'a, T: FieldElement>(
         )
         .unique()
 }
+
+/// Returns a vector of all variables that occur in the assignment.
+fn variables_in_assignment<'a, T: FieldElement>(
+    assignment: &Assignment<'a, T>,
+    fixed_data: &'a FixedData<'a, T>,
+    intermediate_cache: &mut HashMap<AlgebraicReferenceThin, Vec<AlgebraicReferenceThin>>,
+) -> Vec<Variable> {
+    let rhs_var = match &assignment.rhs {
+        VariableOrValue::Variable(v) => Some(v.clone()),
+        VariableOrValue::Value(_) => None,
+    };
+    references_in_expression(assignment.lhs, fixed_data, intermediate_cache)
+        .map(|r| {
+            let name = fixed_data.column_name(&r.poly_id).to_string();
+            Variable::from_reference(&r.with_name(name), assignment.row_offset)
+        })
+        .chain(rhs_var)
+        .collect()
+}
diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs
@@ -18,7 +18,7 @@ use crate::witgen::{
 use super::{
     affine_symbolic_expression,
     effect::{format_code, Effect},
-    identity_queue::IdentityQueue,
+    identity_queue::{IdentityQueue, QueueItem},
     prover_function_heuristics::ProverFunction,
     variable::{Cell, MachineCallVariable, Variable},
     witgen_inference::{BranchResult, CanProcessCall, FixedEvaluator, Value, WitgenInference},
@@ -127,7 +127,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Processor<'a, T, FixedEv
             }
         }
         let branch_depth = 0;
-        let identity_queue = IdentityQueue::new(self.fixed_data, &self.identities);
+        let identity_queue = IdentityQueue::new(self.fixed_data, &self.identities, &[]);
         self.generate_code_for_branch(can_process, witgen, identity_queue, branch_depth)
     }
 
@@ -298,7 +298,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Processor<'a, T, FixedEv
         loop {
             let identity = identity_queue.next();
             let updated_vars = match identity {
-                Some((identity, row_offset)) => match identity {
+                Some(QueueItem::Identity(identity, row_offset)) => match identity {
                     Identity::Polynomial(PolynomialIdentity { id, expression, .. }) => {
                         witgen.process_polynomial_identity(*id, expression, row_offset)
                     }
@@ -317,6 +317,9 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Processor<'a, T, FixedEv
                 },
                 // TODO Also add prover functions to the queue (activated by their variables)
                 // and sort them so that they are always last.
+                Some(QueueItem::Assignment(_assignment)) => {
+                    todo!()
+                }
                 None => self.process_prover_functions(witgen),
             }?;
             if updated_vars.is_empty() && identity.is_none() {
diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs
@@ -680,14 +680,14 @@ fn is_known_zero<T: FieldElement>(x: &Option<AffineSymbolicExpression<T, Variabl
 /// An equality constraint between an algebraic expression evaluated
 /// on a certain row offset and a variable or fixed constant value.
 #[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Debug)]
-struct Assignment<'a, T: FieldElement> {
-    lhs: &'a Expression<T>,
-    row_offset: i32,
-    rhs: VariableOrValue<T, Variable>,
+pub struct Assignment<'a, T: FieldElement> {
+    pub lhs: &'a Expression<T>,
+    pub row_offset: i32,
+    pub rhs: VariableOrValue<T, Variable>,
 }
 
 #[derive(Clone, derive_more::Display, Ord, PartialOrd, Eq, PartialEq, Debug)]
-enum VariableOrValue<T, V> {
+pub enum VariableOrValue<T, V> {
     Variable(V),
     Value(T),
 }
