Merge pull request #685 from powdr-labs/make_identity_expr_generic

Make identities generic over expression type

ast/src/analyzed/display.rs

@@ -7,8 +7,6 @@ use std::fmt::{Display, Formatter, Result};
 
 use itertools::Itertools;
 
-use crate::parsed::display::format_expressions;
-
 use super::*;
 
 impl<T: Display> Display for Analyzed<T> {
@@ -102,7 +100,7 @@ impl<T: Display> Display for RepeatedArray<T> {
     }
 }
 
-impl<T: Display> Display for Identity<T> {
+impl<T: Display> Display for Identity<Expression<T>> {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         match self.kind {
             IdentityKind::Polynomial => {
@@ -120,7 +118,7 @@ impl<T: Display> Display for Identity<T> {
     }
 }
 
-impl<T: Display> Display for SelectedExpressions<T> {
+impl<Expr: Display> Display for SelectedExpressions<Expr> {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         write!(
             f,
@@ -129,7 +127,7 @@ impl<T: Display> Display for SelectedExpressions<T> {
                 .as_ref()
                 .map(|s| format!("{s} "))
                 .unwrap_or_default(),
-            format_expressions(&self.expressions)
+            self.expressions.iter().format(", ")
         )
     }
 }

ast/src/analyzed/mod.rs

@@ -9,11 +9,11 @@ use std::ops::ControlFlow;
 
 use number::DegreeType;
 
-use crate::parsed;
 use crate::parsed::utils::expr_any;
 use crate::parsed::visitor::ExpressionVisitable;
 pub use crate::parsed::BinaryOperator;
 pub use crate::parsed::UnaryOperator;
+use crate::parsed::{self, SelectedExpressions};
 
 #[derive(Debug)]
 pub enum StatementIdentifier {
@@ -29,7 +29,7 @@ pub struct Analyzed<T> {
     pub constants: HashMap<String, T>,
     pub definitions: HashMap<String, (Symbol, Option<FunctionValueDefinition<T>>)>,
     pub public_declarations: HashMap<String, PublicDeclaration>,
-    pub identities: Vec<Identity<T>>,
+    pub identities: Vec<Identity<Expression<T>>>,
     /// The order in which definitions and identities
     /// appear in the source.
     pub source_order: Vec<StatementIdentifier>,
@@ -308,24 +308,26 @@ pub struct PublicDeclaration {
 }
 
 #[derive(Debug, PartialEq, Eq, Clone)]
-pub struct Identity<T> {
+pub struct Identity<Expr> {
     /// The ID is specific to the identity kind.
     pub id: u64,
     pub kind: IdentityKind,
     pub source: SourceRef,
     /// For a simple polynomial identity, the selector contains
     /// the actual expression (see expression_for_poly_id).
-    pub left: SelectedExpressions<T>,
-    pub right: SelectedExpressions<T>,
+    pub left: SelectedExpressions<Expr>,
+    pub right: SelectedExpressions<Expr>,
 }
 
-impl<T> Identity<T> {
+impl<Expr> Identity<Expr> {
     /// Returns the expression in case this is a polynomial identity.
-    pub fn expression_for_poly_id(&self) -> &Expression<T> {
+    pub fn expression_for_poly_id(&self) -> &Expr {
         assert_eq!(self.kind, IdentityKind::Polynomial);
         self.left.selector.as_ref().unwrap()
     }
+}
 
+impl<T> Identity<Expression<T>> {
     pub fn contains_next_ref(&self) -> bool {
         self.left.contains_next_ref() || self.right.contains_next_ref()
     }
@@ -339,9 +341,7 @@ pub enum IdentityKind {
     Connect,
 }
 
-pub type SelectedExpressions<T> = parsed::SelectedExpressions<T, Reference>;
-
-impl<T> SelectedExpressions<T> {
+impl<T> SelectedExpressions<Expression<T>> {
     /// @returns true if the expression contains a reference to a next value of a
     /// (witness or fixed) column
     pub fn contains_next_ref(&self) -> bool {

ast/src/analyzed/visitor.rs

@@ -50,10 +50,10 @@ impl<T> ExpressionVisitable<parsed::Expression<T, Reference>> for Analyzed<T> {
     }
 }
 
-impl<T> ExpressionVisitable<parsed::Expression<T, Reference>> for Identity<T> {
+impl<Expr: ExpressionVisitable<Expr>> ExpressionVisitable<Expr> for Identity<Expr> {
     fn visit_expressions_mut<F, B>(&mut self, f: &mut F, o: VisitOrder) -> ControlFlow<B>
     where
-        F: FnMut(&mut parsed::Expression<T, Reference>) -> ControlFlow<B>,
+        F: FnMut(&mut Expr) -> ControlFlow<B>,
     {
         self.left
             .selector
@@ -67,7 +67,7 @@ impl<T> ExpressionVisitable<parsed::Expression<T, Reference>> for Identity<T> {
 
     fn visit_expressions<F, B>(&self, f: &mut F, o: VisitOrder) -> ControlFlow<B>
     where
-        F: FnMut(&parsed::Expression<T, Reference>) -> ControlFlow<B>,
+        F: FnMut(&Expr) -> ControlFlow<B>,
     {
         self.left
             .selector

ast/src/parsed/display.rs

@@ -416,20 +416,6 @@ impl<T: Display> Display for FunctionDefinition<T> {
     }
 }
 
-impl<T: Display> Display for SelectedExpressions<T> {
-    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
-        write!(
-            f,
-            "{}{{ {} }}",
-            self.selector
-                .as_ref()
-                .map(|s| format!("{s} "))
-                .unwrap_or_default(),
-            format_expressions(&self.expressions)
-        )
-    }
-}
-
 pub fn format_expressions<T: Display, Ref: Display>(expressions: &[Expression<T, Ref>]) -> String {
     format!("{}", expressions.iter().format(", "))
 }

ast/src/parsed/mod.rs

@@ -28,8 +28,16 @@ pub enum PilStatement<T> {
     PolynomialConstantDefinition(usize, String, FunctionDefinition<T>),
     PolynomialCommitDeclaration(usize, Vec<PolynomialName<T>>, Option<FunctionDefinition<T>>),
     PolynomialIdentity(usize, Expression<T>),
-    PlookupIdentity(usize, SelectedExpressions<T>, SelectedExpressions<T>),
-    PermutationIdentity(usize, SelectedExpressions<T>, SelectedExpressions<T>),
+    PlookupIdentity(
+        usize,
+        SelectedExpressions<Expression<T>>,
+        SelectedExpressions<Expression<T>>,
+    ),
+    PermutationIdentity(
+        usize,
+        SelectedExpressions<Expression<T>>,
+        SelectedExpressions<Expression<T>>,
+    ),
     ConnectIdentity(usize, Vec<Expression<T>>, Vec<Expression<T>>),
     ConstantDefinition(usize, String, Expression<T>),
     MacroDefinition(
@@ -43,12 +51,12 @@ pub enum PilStatement<T> {
 }
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
-pub struct SelectedExpressions<T, Ref = ShiftedPolynomialReference<T>> {
-    pub selector: Option<Expression<T, Ref>>,
-    pub expressions: Vec<Expression<T, Ref>>,
+pub struct SelectedExpressions<Expr> {
+    pub selector: Option<Expr>,
+    pub expressions: Vec<Expr>,
 }
 
-impl<T, Ref> Default for SelectedExpressions<T, Ref> {
+impl<Expr> Default for SelectedExpressions<Expr> {
     fn default() -> Self {
         Self {
             selector: Default::default(),

ast/src/parsed/visitor.rs

@@ -256,10 +256,10 @@ impl<T> ExpressionVisitable<Expression<T, ShiftedPolynomialReference<T>>> for Pi
     }
 }
 
-impl<T, Ref> ExpressionVisitable<Expression<T, Ref>> for SelectedExpressions<T, Ref> {
+impl<Expr: ExpressionVisitable<Expr>> ExpressionVisitable<Expr> for SelectedExpressions<Expr> {
     fn visit_expressions_mut<F, B>(&mut self, f: &mut F, o: VisitOrder) -> ControlFlow<B>
     where
-        F: FnMut(&mut Expression<T, Ref>) -> ControlFlow<B>,
+        F: FnMut(&mut Expr) -> ControlFlow<B>,
     {
         self.selector
             .as_mut()
@@ -270,7 +270,7 @@ impl<T, Ref> ExpressionVisitable<Expression<T, Ref>> for SelectedExpressions<T,
 
     fn visit_expressions<F, B>(&self, f: &mut F, o: VisitOrder) -> ControlFlow<B>
     where
-        F: FnMut(&Expression<T, Ref>) -> ControlFlow<B>,
+        F: FnMut(&Expr) -> ControlFlow<B>,
     {
         self.selector
             .as_ref()

backend/src/pilstark/json_exporter/expression_counter.rs

@@ -1,8 +1,11 @@
 use std::collections::HashMap;
 
-use ast::analyzed::{
-    Analyzed, Expression, Identity, PolynomialType, PublicDeclaration, SelectedExpressions,
-    StatementIdentifier, Symbol, SymbolKind,
+use ast::{
+    analyzed::{
+        Analyzed, Identity, PolynomialType, PublicDeclaration, StatementIdentifier, Symbol,
+        SymbolKind,
+    },
+    parsed::SelectedExpressions,
 };
 
 /// Computes expression IDs for each intermediate polynomial.
@@ -32,7 +35,7 @@ trait ExpressionCounter {
     fn expression_count(&self) -> usize;
 }
 
-impl<T> ExpressionCounter for Identity<T> {
+impl<Expr> ExpressionCounter for Identity<Expr> {
     fn expression_count(&self) -> usize {
         self.left.expression_count() + self.right.expression_count()
     }
@@ -50,20 +53,8 @@ impl ExpressionCounter for PublicDeclaration {
     }
 }
 
-impl<T> ExpressionCounter for SelectedExpressions<T> {
+impl<Expr> ExpressionCounter for SelectedExpressions<Expr> {
     fn expression_count(&self) -> usize {
-        self.selector.expression_count() + self.expressions.expression_count()
-    }
-}
-
-impl<T> ExpressionCounter for Vec<Expression<T>> {
-    fn expression_count(&self) -> usize {
-        self.len()
-    }
-}
-
-impl<T> ExpressionCounter for Option<Expression<T>> {
-    fn expression_count(&self) -> usize {
-        (self.is_some()).into()
+        self.selector.is_some() as usize + self.expressions.len()
     }
 }

executor/src/witgen/generator.rs

@@ -1,4 +1,5 @@
-use ast::analyzed::{Identity, IdentityKind, PolyID, PolynomialReference, SelectedExpressions};
+use ast::analyzed::{Expression, Identity, IdentityKind, PolyID, PolynomialReference};
+use ast::parsed::SelectedExpressions;
 use number::FieldElement;
 use std::collections::{HashMap, HashSet};
 
@@ -15,7 +16,7 @@ use super::{EvalResult, FixedData, MutableState};
 
 pub struct Generator<'a, T: FieldElement> {
     fixed_data: &'a FixedData<'a, T>,
-    identities: Vec<&'a Identity<T>>,
+    identities: Vec<&'a Identity<Expression<T>>>,
     witnesses: HashSet<PolyID>,
     global_range_constraints: WitnessColumnMap<Option<RangeConstraint<T>>>,
     data: Vec<Row<'a, T>>,
@@ -28,7 +29,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for Generator<'a, T> {
         _fixed_lookup: &mut FixedLookup<T>,
         _kind: IdentityKind,
         _left: &[AffineExpression<&'a PolynomialReference, T>],
-        _right: &'a SelectedExpressions<T>,
+        _right: &'a SelectedExpressions<Expression<T>>,
         _machines: Machines<'a, '_, T>,
     ) -> Option<EvalResult<'a, T>> {
         unimplemented!()
@@ -50,7 +51,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for Generator<'a, T> {
 impl<'a, T: FieldElement> Generator<'a, T> {
     pub fn new(
         fixed_data: &'a FixedData<'a, T>,
-        identities: &[&'a Identity<T>],
+        identities: &[&'a Identity<Expression<T>>],
         witnesses: HashSet<PolyID>,
         global_range_constraints: &WitnessColumnMap<Option<RangeConstraint<T>>>,
     ) -> Self {

executor/src/witgen/global_constraints.rs

@@ -35,7 +35,7 @@ impl<'a, T: FieldElement> RangeConstraintSet<&PolynomialReference, T>
 
 pub struct GlobalConstraints<'a, T: FieldElement> {
     pub known_witness_constraints: WitnessColumnMap<Option<RangeConstraint<T>>>,
-    pub retained_identities: Vec<&'a Identity<T>>,
+    pub retained_identities: Vec<&'a Identity<Expression<T>>>,
 }
 
 /// Determines global constraints on witness and fixed columns.
@@ -44,7 +44,7 @@ pub struct GlobalConstraints<'a, T: FieldElement> {
 /// TODO at some point, we should check that they still hold.
 pub fn determine_global_constraints<'a, T: FieldElement>(
     fixed_data: &'a FixedData<T>,
-    identities: Vec<&'a Identity<T>>,
+    identities: Vec<&'a Identity<Expression<T>>>,
 ) -> GlobalConstraints<'a, T> {
     let mut known_constraints = BTreeMap::new();
     // For these columns, we know that they are not only constrained to those bits
@@ -134,7 +134,7 @@ fn process_fixed_column<T: FieldElement>(fixed: &[T]) -> Option<(RangeConstraint
 /// no further information than the range constraint.
 fn propagate_constraints<T: FieldElement>(
     mut known_constraints: BTreeMap<PolyID, RangeConstraint<T>>,
-    identity: &Identity<T>,
+    identity: &Identity<Expression<T>>,
     full_span: &BTreeSet<PolyID>,
 ) -> (BTreeMap<PolyID, RangeConstraint<T>>, bool) {
     let mut remove = false;

executor/src/witgen/identity_processor.rs

@@ -1,6 +1,9 @@
 use std::{collections::HashMap, sync::Mutex};
 
-use ast::analyzed::{Expression, Identity, IdentityKind, PolynomialReference, SelectedExpressions};
+use ast::{
+    analyzed::{Expression, Identity, IdentityKind, PolynomialReference},
+    parsed::SelectedExpressions,
+};
 use itertools::{Either, Itertools};
 use lazy_static::lazy_static;
 use number::FieldElement;
@@ -86,7 +89,7 @@ impl<'a, 'b, T: FieldElement> IdentityProcessor<'a, 'b, T> {
     /// Returns the updates.
     pub fn process_identity(
         &mut self,
-        identity: &'a Identity<T>,
+        identity: &'a Identity<Expression<T>>,
         rows: &RowPair<'_, 'a, T>,
     ) -> EvalResult<'a, T> {
         let result = match identity.kind {
@@ -106,7 +109,7 @@ impl<'a, 'b, T: FieldElement> IdentityProcessor<'a, 'b, T> {
 
     fn process_polynomial_identity(
         &self,
-        identity: &'a Identity<T>,
+        identity: &'a Identity<Expression<T>>,
         rows: &RowPair<T>,
     ) -> EvalResult<'a, T> {
         match rows.evaluate(identity.expression_for_poly_id()) {
@@ -117,7 +120,7 @@ impl<'a, 'b, T: FieldElement> IdentityProcessor<'a, 'b, T> {
 
     fn process_plookup(
         &mut self,
-        identity: &'a Identity<T>,
+        identity: &'a Identity<Expression<T>>,
         rows: &RowPair<'_, 'a, T>,
     ) -> EvalResult<'a, T> {
         if let Some(left_selector) = &identity.left.selector {
@@ -187,7 +190,7 @@ impl<'a, 'b, T: FieldElement> IdentityProcessor<'a, 'b, T> {
     pub fn process_link(
         &mut self,
         left: &[AffineExpression<&'a PolynomialReference, T>],
-        right: &'a SelectedExpressions<T>,
+        right: &'a SelectedExpressions<Expression<T>>,
         current_rows: &RowPair<'_, 'a, T>,
     ) -> EvalResult<'a, T> {
         // sanity check that the right hand side selector is active
@@ -246,7 +249,10 @@ lazy_static! {
         Mutex::new(Default::default());
 }
 
-fn report_identity_solving<T: FieldElement, K>(identity: &Identity<T>, result: &EvalResult<T, K>) {
+fn report_identity_solving<T: FieldElement, K>(
+    identity: &Identity<Expression<T>>,
+    result: &EvalResult<T, K>,
+) {
     let success = result.as_ref().map(|r| r.is_complete()).unwrap_or_default() as u64;
     let mut stat = STATISTICS.lock().unwrap();
     stat.entry((identity.id, identity.kind))

executor/src/witgen/machines/block_machine.rs

@@ -12,9 +12,8 @@ use crate::witgen::Constraints;
 use crate::witgen::{
     machines::Machine, range_constraints::RangeConstraint, EvalError, EvalValue, IncompleteCause,
 };
-use ast::analyzed::{
-    Expression, Identity, IdentityKind, PolyID, PolynomialReference, SelectedExpressions,
-};
+use ast::analyzed::{Expression, Identity, IdentityKind, PolyID, PolynomialReference};
+use ast::parsed::SelectedExpressions;
 use number::{DegreeType, FieldElement};
 
 enum ProcessResult<'a, T: FieldElement> {
@@ -39,9 +38,9 @@ pub struct BlockMachine<'a, T: FieldElement> {
     block_size: usize,
     /// The right-hand side of the connecting identity, needed to identify
     /// when this machine is responsible.
-    selected_expressions: SelectedExpressions<T>,
+    selected_expressions: SelectedExpressions<Expression<T>>,
     /// The internal identities
-    identities: Vec<&'a Identity<T>>,
+    identities: Vec<&'a Identity<Expression<T>>>,
     /// The row factory
     row_factory: RowFactory<'a, T>,
     /// The data of the machine.
@@ -57,8 +56,8 @@ pub struct BlockMachine<'a, T: FieldElement> {
 impl<'a, T: FieldElement> BlockMachine<'a, T> {
     pub fn try_new(
         fixed_data: &'a FixedData<'a, T>,
-        connecting_identities: &[&'a Identity<T>],
-        identities: &[&'a Identity<T>],
+        connecting_identities: &[&'a Identity<Expression<T>>],
+        identities: &[&'a Identity<Expression<T>>],
         witness_cols: &HashSet<PolyID>,
         global_range_constraints: &WitnessColumnMap<Option<RangeConstraint<T>>>,
     ) -> Option<Self> {
@@ -140,7 +139,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
         fixed_lookup: &'b mut FixedLookup<T>,
         kind: IdentityKind,
         left: &[AffineExpression<&'a PolynomialReference, T>],
-        right: &'a SelectedExpressions<T>,
+        right: &'a SelectedExpressions<Expression<T>>,
         machines: Machines<'a, 'b, T>,
     ) -> Option<EvalResult<'a, T>> {
         if *right != self.selected_expressions || kind != IdentityKind::Plookup {
@@ -252,7 +251,7 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         &mut self,
         fixed_lookup: &'b mut FixedLookup<T>,
         left: &[AffineExpression<&'a PolynomialReference, T>],
-        right: &'a SelectedExpressions<T>,
+        right: &'a SelectedExpressions<Expression<T>>,
         machines: Machines<'a, 'b, T>,
     ) -> EvalResult<'a, T> {
         log::trace!("Start processing block machine '{}'", self.name());
@@ -351,7 +350,7 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         &self,
         identity_processor: &mut IdentityProcessor<'a, 'b, T>,
         left: &[AffineExpression<&'a PolynomialReference, T>],
-        right: &'a SelectedExpressions<T>,
+        right: &'a SelectedExpressions<Expression<T>>,
         sequence_iterator: &mut ProcessingSequenceIterator,
     ) -> Result<ProcessResult<'a, T>, EvalError<T>> {
         // Make the block two rows larger than the block size, it includes the last row of the previous block

executor/src/witgen/machines/double_sorted_witness_machine.rs

@@ -1,6 +1,7 @@
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::iter::once;
 
+use ast::parsed::SelectedExpressions;
 use itertools::Itertools;
 use num_traits::Zero;
 
@@ -12,9 +13,7 @@ use crate::witgen::{EvalResult, FixedData};
 use crate::witgen::{EvalValue, IncompleteCause};
 use number::{DegreeType, FieldElement};
 
-use ast::analyzed::{
-    Expression, Identity, IdentityKind, PolyID, PolynomialReference, Reference, SelectedExpressions,
-};
+use ast::analyzed::{Expression, Identity, IdentityKind, PolyID, PolynomialReference, Reference};
 
 /// TODO make this generic
 
@@ -43,7 +42,7 @@ impl<T: FieldElement> DoubleSortedWitnesses<T> {
 
     pub fn try_new(
         fixed_data: &FixedData<T>,
-        _identities: &[&Identity<T>],
+        _identities: &[&Identity<Expression<T>>],
         witness_cols: &HashSet<PolyID>,
     ) -> Option<Self> {
         // get the namespaces and column names
@@ -100,7 +99,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses<T> {
         _fixed_lookup: &mut FixedLookup<T>,
         kind: IdentityKind,
         left: &[AffineExpression<&'a PolynomialReference, T>],
-        right: &'a SelectedExpressions<T>,
+        right: &'a SelectedExpressions<Expression<T>>,
         _machines: Machines<'a, '_, T>,
     ) -> Option<EvalResult<'a, T>> {
         if kind != IdentityKind::Permutation
@@ -174,7 +173,7 @@ impl<T: FieldElement> DoubleSortedWitnesses<T> {
     fn process_plookup_internal<'a>(
         &mut self,
         left: &[AffineExpression<&'a PolynomialReference, T>],
-        right: &SelectedExpressions<T>,
+        right: &SelectedExpressions<Expression<T>>,
     ) -> EvalResult<'a, T> {
         // We blindly assume the lookup is of the form
         // OP { ADDR, STEP, X } is m_is_write { m_addr, m_step, m_value }

executor/src/witgen/machines/fixed_lookup_machine.rs

@@ -2,7 +2,8 @@ use std::collections::{BTreeMap, HashMap, HashSet};
 use std::mem;
 use std::num::NonZeroUsize;
 
-use ast::analyzed::{Identity, IdentityKind, PolyID, PolynomialReference, SelectedExpressions};
+use ast::analyzed::{Expression, Identity, IdentityKind, PolyID, PolynomialReference};
+use ast::parsed::SelectedExpressions;
 use itertools::Itertools;
 use number::FieldElement;
 
@@ -162,7 +163,7 @@ pub struct FixedLookup<T> {
 impl<T: FieldElement> FixedLookup<T> {
     pub fn try_new(
         _fixed_data: &FixedData<T>,
-        identities: &[&Identity<T>],
+        identities: &[&Identity<Expression<T>>],
         witness_names: &HashSet<&str>,
     ) -> Option<Self> {
         if identities.is_empty() && witness_names.is_empty() {
@@ -177,7 +178,7 @@ impl<T: FieldElement> FixedLookup<T> {
         fixed_data: &FixedData<T>,
         kind: IdentityKind,
         left: &[AffineExpression<&'b PolynomialReference, T>],
-        right: &'b SelectedExpressions<T>,
+        right: &'b SelectedExpressions<Expression<T>>,
     ) -> Option<EvalResult<'b, T>> {
         // This is a matching machine if it is a plookup and the RHS is fully constant.
         if kind != IdentityKind::Plookup

executor/src/witgen/machines/machine_extractor.rs

@@ -9,15 +9,16 @@ use super::KnownMachine;
 use crate::witgen::{
     column_map::WitnessColumnMap, generator::Generator, range_constraints::RangeConstraint,
 };
-use ast::analyzed::{Expression, Identity, IdentityKind, PolyID, Reference, SelectedExpressions};
+use ast::analyzed::{Expression, Identity, IdentityKind, PolyID, Reference};
 use ast::parsed::visitor::ExpressionVisitable;
+use ast::parsed::SelectedExpressions;
 use itertools::Itertools;
 use number::FieldElement;
 
 pub struct ExtractionOutput<'a, T: FieldElement> {
     pub fixed_lookup: FixedLookup<T>,
     pub machines: Vec<KnownMachine<'a, T>>,
-    pub base_identities: Vec<&'a Identity<T>>,
+    pub base_identities: Vec<&'a Identity<Expression<T>>>,
     pub base_witnesses: HashSet<PolyID>,
 }
 
@@ -26,7 +27,7 @@ pub struct ExtractionOutput<'a, T: FieldElement> {
 /// that are not "internal" to the machines.
 pub fn split_out_machines<'a, T: FieldElement>(
     fixed: &'a FixedData<'a, T>,
-    identities: Vec<&'a Identity<T>>,
+    identities: Vec<&'a Identity<Expression<T>>>,
     global_range_constraints: &WitnessColumnMap<Option<RangeConstraint<T>>>,
 ) -> ExtractionOutput<'a, T> {
     let fixed_lookup = FixedLookup::try_new(fixed, &[], &Default::default()).unwrap();
@@ -140,7 +141,7 @@ pub fn split_out_machines<'a, T: FieldElement>(
 fn all_row_connected_witnesses<T>(
     mut witnesses: HashSet<PolyID>,
     all_witnesses: &HashSet<PolyID>,
-    identities: &[&Identity<T>],
+    identities: &[&Identity<Expression<T>>],
 ) -> HashSet<PolyID> {
     loop {
         let count = witnesses.len();
@@ -173,7 +174,7 @@ fn all_row_connected_witnesses<T>(
 }
 
 /// Extracts all references to names from an identity.
-pub fn refs_in_identity<T>(identity: &Identity<T>) -> HashSet<PolyID> {
+pub fn refs_in_identity<T>(identity: &Identity<Expression<T>>) -> HashSet<PolyID> {
     let mut refs: HashSet<PolyID> = Default::default();
     identity.pre_visit_expressions(&mut |expr| {
         ref_of_expression(expr).map(|id| refs.insert(id));
@@ -182,7 +183,9 @@ pub fn refs_in_identity<T>(identity: &Identity<T>) -> HashSet<PolyID> {
 }
 
 /// Extracts all references to names from selected expressions.
-pub fn refs_in_selected_expressions<T>(selexpr: &SelectedExpressions<T>) -> HashSet<PolyID> {
+pub fn refs_in_selected_expressions<T>(
+    selexpr: &SelectedExpressions<Expression<T>>,
+) -> HashSet<PolyID> {
     let mut refs: HashSet<PolyID> = Default::default();
     selexpr.pre_visit_expressions(&mut |expr| {
         ref_of_expression(expr).map(|id| refs.insert(id));

executor/src/witgen/machines/mod.rs

@@ -1,14 +1,15 @@
 use std::collections::HashMap;
 
-use ast::analyzed::IdentityKind;
+use ast::analyzed::Expression;
 use ast::analyzed::PolynomialReference;
-use ast::analyzed::SelectedExpressions;
+use ast::parsed::SelectedExpressions;
 use number::FieldElement;
 
 use self::block_machine::BlockMachine;
 use self::double_sorted_witness_machine::DoubleSortedWitnesses;
 pub use self::fixed_lookup_machine::FixedLookup;
 use self::sorted_witness_machine::SortedWitnesses;
+use ast::analyzed::IdentityKind;
 
 use super::affine_expression::AffineExpression;
 use super::generator::Generator;
@@ -36,7 +37,7 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync {
         fixed_lookup: &'b mut FixedLookup<T>,
         kind: IdentityKind,
         left: &[AffineExpression<&'a PolynomialReference, T>],
-        right: &'a SelectedExpressions<T>,
+        right: &'a SelectedExpressions<Expression<T>>,
         machines: Machines<'a, 'b, T>,
     ) -> Option<EvalResult<'a, T>>;
 
@@ -72,7 +73,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
         fixed_lookup: &'b mut FixedLookup<T>,
         kind: IdentityKind,
         left: &[AffineExpression<&'a PolynomialReference, T>],
-        right: &'a SelectedExpressions<T>,
+        right: &'a SelectedExpressions<Expression<T>>,
         machines: Machines<'a, 'b, T>,
     ) -> Option<crate::witgen::EvalResult<'a, T>> {
         self.get()

executor/src/witgen/machines/sorted_witness_machine.rs

@@ -1,5 +1,6 @@
 use std::collections::{BTreeMap, HashMap, HashSet};
 
+use ast::parsed::SelectedExpressions;
 use itertools::Itertools;
 
 use super::super::affine_expression::AffineExpression;
@@ -12,9 +13,7 @@ use crate::witgen::{
     symbolic_evaluator::SymbolicEvaluator,
 };
 use crate::witgen::{EvalValue, IncompleteCause};
-use ast::analyzed::{
-    Expression, Identity, IdentityKind, PolyID, PolynomialReference, Reference, SelectedExpressions,
-};
+use ast::analyzed::{Expression, Identity, IdentityKind, PolyID, PolynomialReference, Reference};
 use number::FieldElement;
 
 /// A machine that can support a lookup in a set of columns that are sorted
@@ -34,7 +33,7 @@ pub struct SortedWitnesses<T> {
 impl<T: FieldElement> SortedWitnesses<T> {
     pub fn try_new(
         fixed_data: &FixedData<T>,
-        identities: &[&Identity<T>],
+        identities: &[&Identity<Expression<T>>],
         witnesses: &HashSet<PolyID>,
     ) -> Option<Self> {
         if identities.len() != 1 {
@@ -58,7 +57,10 @@ impl<T: FieldElement> SortedWitnesses<T> {
     }
 }
 
-fn check_identity<T: FieldElement>(fixed_data: &FixedData<T>, id: &Identity<T>) -> Option<PolyID> {
+fn check_identity<T: FieldElement>(
+    fixed_data: &FixedData<T>,
+    id: &Identity<Expression<T>>,
+) -> Option<PolyID> {
     // Looking for NOTLAST { A' - A } in { POSITIVE }
     if id.kind != IdentityKind::Plookup
         || id.right.selector.is_some()
@@ -125,7 +127,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for SortedWitnesses<T> {
         _fixed_lookup: &mut FixedLookup<T>,
         kind: IdentityKind,
         left: &[AffineExpression<&'a PolynomialReference, T>],
-        right: &'a SelectedExpressions<T>,
+        right: &'a SelectedExpressions<Expression<T>>,
         _machines: Machines<'a, '_, T>,
     ) -> Option<EvalResult<'a, T>> {
         if kind != IdentityKind::Plookup || right.selector.is_some() {
@@ -182,7 +184,7 @@ impl<T: FieldElement> SortedWitnesses<T> {
         &mut self,
         fixed_data: &FixedData<T>,
         left: &[AffineExpression<&'a PolynomialReference, T>],
-        right: &SelectedExpressions<T>,
+        right: &SelectedExpressions<Expression<T>>,
         rhs: Vec<&PolynomialReference>,
     ) -> EvalResult<'a, T> {
         let key_index = rhs

executor/src/witgen/processor.rs

@@ -1,6 +1,9 @@
 use std::{collections::HashSet, marker::PhantomData};
 
-use ast::analyzed::{Identity, PolyID, PolynomialReference, SelectedExpressions};
+use ast::{
+    analyzed::{Expression, Identity, PolyID, PolynomialReference},
+    parsed::SelectedExpressions,
+};
 use number::FieldElement;
 
 use crate::witgen::Constraint;
@@ -25,11 +28,11 @@ pub struct OuterQuery<'a, T: FieldElement> {
     /// This will be mutated while processing the block.
     left: Left<'a, T>,
     /// The right-hand side of the outer query.
-    right: &'a SelectedExpressions<T>,
+    right: &'a SelectedExpressions<Expression<T>>,
 }
 
 impl<'a, T: FieldElement> OuterQuery<'a, T> {
-    pub fn new(left: Left<'a, T>, right: &'a SelectedExpressions<T>) -> Self {
+    pub fn new(left: Left<'a, T>, right: &'a SelectedExpressions<Expression<T>>) -> Self {
         Self { left, right }
     }
 }
@@ -46,7 +49,7 @@ pub struct Processor<'a, 'b, 'c, T: FieldElement, CalldataAvailable> {
     /// The rows that are being processed.
     data: Vec<Row<'a, T>>,
     /// The list of identities
-    identities: &'c [&'a Identity<T>],
+    identities: &'c [&'a Identity<Expression<T>>],
     /// The identity processor
     identity_processor: &'c mut IdentityProcessor<'a, 'b, T>,
     /// The fixed data (containing information about all columns)
@@ -65,7 +68,7 @@ impl<'a, 'b, 'c, T: FieldElement> Processor<'a, 'b, 'c, T, WithoutCalldata> {
         row_offset: u64,
         data: Vec<Row<'a, T>>,
         identity_processor: &'c mut IdentityProcessor<'a, 'b, T>,
-        identities: &'c [&'a Identity<T>],
+        identities: &'c [&'a Identity<Expression<T>>],
         fixed_data: &'a FixedData<'a, T>,
         row_factory: RowFactory<'a, T>,
         witness_cols: &'c HashSet<PolyID>,

executor/src/witgen/vm_processor.rs

@@ -1,4 +1,4 @@
-use ast::analyzed::{Identity, IdentityKind, PolyID, PolynomialReference};
+use ast::analyzed::{Expression, Identity, IdentityKind, PolyID, PolynomialReference};
 use itertools::Itertools;
 use number::{DegreeType, FieldElement};
 use parser_util::lines::indent;
@@ -23,18 +23,18 @@ enum ProcessingPhase {
 
 /// A list of identities with a flag whether it is complete.
 struct CompletableIdentities<'a, T: FieldElement> {
-    identities_with_complete: Vec<(&'a Identity<T>, bool)>,
+    identities_with_complete: Vec<(&'a Identity<Expression<T>>, bool)>,
 }
 
 impl<'a, T: FieldElement> CompletableIdentities<'a, T> {
-    fn new(identities: impl Iterator<Item = &'a Identity<T>>) -> Self {
+    fn new(identities: impl Iterator<Item = &'a Identity<Expression<T>>>) -> Self {
         Self {
             identities_with_complete: identities.map(|identity| (identity, false)).collect(),
         }
     }
 
     /// Yields immutable references to the identity and mutable references to the complete flag.
-    fn iter_mut(&mut self) -> impl Iterator<Item = (&'a Identity<T>, &mut bool)> {
+    fn iter_mut(&mut self) -> impl Iterator<Item = (&'a Identity<Expression<T>>, &mut bool)> {
         self.identities_with_complete
             .iter_mut()
             .map(|(identity, complete)| (*identity, complete))
@@ -47,10 +47,10 @@ pub struct VmProcessor<'a, T: FieldElement> {
     fixed_data: &'a FixedData<'a, T>,
     /// The subset of identities that contains a reference to the next row
     /// (precomputed once for performance reasons)
-    identities_with_next_ref: Vec<&'a Identity<T>>,
+    identities_with_next_ref: Vec<&'a Identity<Expression<T>>>,
     /// The subset of identities that does not contain a reference to the next row
     /// (precomputed once for performance reasons)
-    identities_without_next_ref: Vec<&'a Identity<T>>,
+    identities_without_next_ref: Vec<&'a Identity<Expression<T>>>,
     data: Vec<Row<'a, T>>,
     last_report: DegreeType,
     last_report_time: Instant,
@@ -59,7 +59,7 @@ pub struct VmProcessor<'a, T: FieldElement> {
 impl<'a, T: FieldElement> VmProcessor<'a, T> {
     pub fn new(
         fixed_data: &'a FixedData<'a, T>,
-        identities: &[&'a Identity<T>],
+        identities: &[&'a Identity<Expression<T>>],
         witnesses: HashSet<PolyID>,
         data: Vec<Row<'a, T>>,
     ) -> Self {

halo2/src/circuit_builder.rs

@@ -1,4 +1,4 @@
-use ast::parsed::BinaryOperator;
+use ast::parsed::{BinaryOperator, SelectedExpressions};
 use num_bigint::BigUint;
 use polyexen::expr::{ColumnKind, ColumnQuery, Expr, PlonkVar};
 use polyexen::plaf::backends::halo2::PlafH2Circuit;
@@ -6,7 +6,7 @@ use polyexen::plaf::{
     ColumnFixed, ColumnWitness, Columns, Info, Lookup, Plaf, Poly, Shuffle, Witness,
 };
 
-use ast::analyzed::{Analyzed, Expression, IdentityKind, Reference, SelectedExpressions};
+use ast::analyzed::{Analyzed, Expression, IdentityKind, Reference};
 use num_traits::One;
 use number::{BigInt, FieldElement};
 
@@ -89,7 +89,7 @@ pub(crate) fn analyzed_to_circuit<T: FieldElement>(
 
     // build Plaf polys. -------------------------------------------------------------------------
 
-    let apply_selectors_to_set = |set: &SelectedExpressions<T>| {
+    let apply_selectors_to_set = |set: &SelectedExpressions<Expression<T>>| {
         let selector = set
             .selector
             .clone()

parser/src/powdr.lalrpop

@@ -140,7 +140,7 @@ PlookupIdentity: PilStatement<T> = {
     <@L> <SelectedExpressions> "in" <SelectedExpressions> => PilStatement::PlookupIdentity(<>)
 }
 
-SelectedExpressions: SelectedExpressions<T> = {
+SelectedExpressions: SelectedExpressions<Expression<T>> = {
     <selector:Expression?> "{" <expressions:ExpressionList> "}" => SelectedExpressions{<>},
     Expression => SelectedExpressions{selector: None, expressions: vec![<>]},
 }

pil_analyzer/src/pil_analyzer.rs

@@ -7,14 +7,14 @@ use analysis::MacroExpander;
 use ast::parsed::visitor::ExpressionVisitable;
 use ast::parsed::{
     self, ArrayExpression, ArrayLiteral, FunctionDefinition, LambdaExpression, MatchArm,
-    MatchPattern, PilStatement, PolynomialName,
+    MatchPattern, PilStatement, PolynomialName, SelectedExpressions,
 };
 use number::{DegreeType, FieldElement};
 
 use ast::analyzed::{
     Analyzed, Expression, FunctionValueDefinition, Identity, IdentityKind, PolynomialReference,
-    PolynomialType, PublicDeclaration, Reference, RepeatedArray, SelectedExpressions, SourceRef,
-    StatementIdentifier, Symbol, SymbolKind,
+    PolynomialType, PublicDeclaration, Reference, RepeatedArray, SourceRef, StatementIdentifier,
+    Symbol, SymbolKind,
 };
 
 use crate::evaluator::Evaluator;
@@ -39,7 +39,7 @@ struct PILAnalyzer<T> {
     constants: HashMap<String, T>,
     definitions: HashMap<String, (Symbol, Option<FunctionValueDefinition<T>>)>,
     public_declarations: HashMap<String, PublicDeclaration>,
-    identities: Vec<Identity<T>>,
+    identities: Vec<Identity<Expression<T>>>,
     /// The order in which definitions and identities
     /// appear in the source.
     source_order: Vec<StatementIdentifier>,
@@ -513,8 +513,8 @@ impl<'a, T: FieldElement> ExpressionProcessor<'a, T> {
 
     pub fn process_selected_expression(
         &mut self,
-        expr: ::ast::parsed::SelectedExpressions<T>,
-    ) -> SelectedExpressions<T> {
+        expr: SelectedExpressions<parsed::Expression<T>>,
+    ) -> SelectedExpressions<Expression<T>> {
         SelectedExpressions {
             selector: expr.selector.map(|e| self.process_expression(e)),
             expressions: self.process_expressions(expr.expressions),
@@ -550,18 +550,15 @@ impl<'a, T: FieldElement> ExpressionProcessor<'a, T> {
         }
     }
 
-    pub fn process_expressions(
-        &mut self,
-        exprs: Vec<::ast::parsed::Expression<T>>,
-    ) -> Vec<Expression<T>> {
+    pub fn process_expressions(&mut self, exprs: Vec<parsed::Expression<T>>) -> Vec<Expression<T>> {
         exprs
             .into_iter()
             .map(|e| self.process_expression(e))
             .collect()
     }
 
-    pub fn process_expression(&mut self, expr: ::ast::parsed::Expression<T>) -> Expression<T> {
-        use ast::parsed::Expression as PExpression;
+    pub fn process_expression(&mut self, expr: parsed::Expression<T>) -> Expression<T> {
+        use parsed::Expression as PExpression;
         match expr {
             PExpression::Constant(name) => Expression::Constant(name),
             PExpression::Reference(poly) => {
@@ -675,7 +672,9 @@ impl<'a, T: FieldElement> ExpressionProcessor<'a, T> {
     }
 }
 
-pub fn inline_intermediate_polynomials<T: Copy>(analyzed: &Analyzed<T>) -> Vec<Identity<T>> {
+pub fn inline_intermediate_polynomials<T: Copy>(
+    analyzed: &Analyzed<T>,
+) -> Vec<Identity<Expression<T>>> {
     substitute_intermediate(
         analyzed.identities.clone(),
         &analyzed
@@ -697,9 +696,9 @@ pub fn inline_intermediate_polynomials<T: Copy>(analyzed: &Analyzed<T>) -> Vec<I
 /// Takes identities as values and inlines intermediate polynomials everywhere, returning a vector of the updated identities
 /// TODO: this could return an iterator
 fn substitute_intermediate<T: Copy>(
-    identities: impl IntoIterator<Item = Identity<T>>,
+    identities: impl IntoIterator<Item = Identity<Expression<T>>>,
     intermediate_polynomials: &HashMap<u64, Expression<T>>,
-) -> Vec<Identity<T>> {
+) -> Vec<Identity<Expression<T>>> {
     identities
         .into_iter()
         .scan(HashMap::default(), |cache, mut identity| {