Merge pull request #795 from powdr-labs/move_util_functions

Move inlining function.

ast/src/analyzed/mod.rs

@@ -6,7 +6,7 @@ use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use std::fmt::Display;
 use std::ops::{self, ControlFlow};
 
-use number::DegreeType;
+use number::{DegreeType, FieldElement};
 
 use crate::parsed::utils::expr_any;
 use crate::parsed::visitor::ExpressionVisitable;
@@ -294,6 +294,84 @@ impl<T> Analyzed<T> {
     }
 }
 
+impl<T: FieldElement> Analyzed<T> {
+    /// @returns all identities with intermediate polynomials inlined.
+    pub fn identities_with_inlined_intermediate_polynomials(
+        &self,
+    ) -> Vec<Identity<AlgebraicExpression<T>>> {
+        let intermediates = &self
+            .intermediate_polys_in_source_order()
+            .iter()
+            .map(|(symbol, def)| (symbol.id, def))
+            .collect();
+
+        substitute_intermediate(self.identities.clone(), intermediates)
+    }
+}
+
+/// 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<AlgebraicExpression<T>>>,
+    intermediate_polynomials: &HashMap<u64, &AlgebraicExpression<T>>,
+) -> Vec<Identity<AlgebraicExpression<T>>> {
+    identities
+        .into_iter()
+        .scan(HashMap::default(), |cache, mut identity| {
+            identity.post_visit_expressions_mut(&mut |e| {
+                if let AlgebraicExpression::Reference(poly) = e {
+                    match poly.poly_id.ptype {
+                        PolynomialType::Committed => {}
+                        PolynomialType::Constant => {}
+                        PolynomialType::Intermediate => {
+                            // recursively inline intermediate polynomials, updating the cache
+                            *e = inlined_expression_from_intermediate_poly_id(
+                                poly.poly_id.id,
+                                intermediate_polynomials,
+                                cache,
+                            );
+                        }
+                    }
+                }
+            });
+            Some(identity)
+        })
+        .collect()
+}
+
+/// Recursively inlines intermediate polynomials inside an expression and returns the new expression
+/// This uses a cache to avoid resolving an intermediate polynomial twice
+fn inlined_expression_from_intermediate_poly_id<T: Copy>(
+    poly_id: u64,
+    intermediate_polynomials: &HashMap<u64, &AlgebraicExpression<T>>,
+    cache: &mut HashMap<u64, AlgebraicExpression<T>>,
+) -> AlgebraicExpression<T> {
+    if let Some(e) = cache.get(&poly_id) {
+        return e.clone();
+    }
+    let mut expr = intermediate_polynomials[&poly_id].clone();
+    expr.post_visit_expressions_mut(&mut |e| {
+        if let AlgebraicExpression::Reference(r) = e {
+            match r.poly_id.ptype {
+                PolynomialType::Committed => {}
+                PolynomialType::Constant => {}
+                PolynomialType::Intermediate => {
+                    // read from the cache, if no cache hit, compute the inlined expression
+                    *e = cache.get(&r.poly_id.id).cloned().unwrap_or_else(|| {
+                        inlined_expression_from_intermediate_poly_id(
+                            r.poly_id.id,
+                            intermediate_polynomials,
+                            cache,
+                        )
+                    });
+                }
+            }
+        }
+    });
+    cache.insert(poly_id, expr.clone());
+    expr
+}
+
 #[derive(Debug, Clone)]
 pub struct Symbol {
     pub id: u64,

executor/src/witgen/mod.rs

@@ -16,8 +16,6 @@ use self::identity_processor::Machines;
 use self::machines::machine_extractor::ExtractionOutput;
 use self::machines::{FixedLookup, Machine};
 
-use pil_analyzer::pil_analyzer::inline_intermediate_polynomials;
-
 mod affine_expression;
 mod block_processor;
 mod data_structures;
@@ -92,7 +90,9 @@ impl<'a, 'b, T: FieldElement, Q: QueryCallback<T>> WitnessGenerator<'a, 'b, T, Q
             self.fixed_col_values,
             self.external_witness_values,
         );
-        let identities = inline_intermediate_polynomials(self.analyzed);
+        let identities = self
+            .analyzed
+            .identities_with_inlined_intermediate_polynomials();
 
         let (
             constraints,

halo2/src/circuit_builder.rs

@@ -14,8 +14,6 @@ use number::{BigInt, FieldElement};
 
 use super::circuit_data::CircuitData;
 
-use pil_analyzer::pil_analyzer::inline_intermediate_polynomials;
-
 pub(crate) fn analyzed_to_circuit<T: FieldElement>(
     analyzed: &Analyzed<T>,
     fixed: &[(String, Vec<T>)],
@@ -116,7 +114,7 @@ pub(crate) fn analyzed_to_circuit<T: FieldElement>(
             .collect()
     };
 
-    let identities = inline_intermediate_polynomials(analyzed);
+    let identities = analyzed.identities_with_inlined_intermediate_polynomials();
     for id in &identities {
         match id.kind {
             IdentityKind::Polynomial => {

pil_analyzer/src/pil_analyzer.rs

@@ -4,15 +4,14 @@ use std::path::{Path, PathBuf};
 
 use analysis::MacroExpander;
 
-use ast::parsed::visitor::ExpressionVisitable;
 use ast::parsed::{
     self, FunctionDefinition, LambdaExpression, PilStatement, PolynomialName, SelectedExpressions,
 };
 use number::{DegreeType, FieldElement};
 
 use ast::analyzed::{
-    AlgebraicExpression, Analyzed, Expression, FunctionValueDefinition, Identity, IdentityKind,
-    PolynomialType, PublicDeclaration, SourceRef, StatementIdentifier, Symbol, SymbolKind,
+    Analyzed, Expression, FunctionValueDefinition, Identity, IdentityKind, PolynomialType,
+    PublicDeclaration, SourceRef, StatementIdentifier, Symbol, SymbolKind,
 };
 
 use crate::evaluator::EvalError;
@@ -477,81 +476,6 @@ impl<'a, T: FieldElement> ReferenceResolver for PILResolver<'a, T> {
     }
 }
 
-pub fn inline_intermediate_polynomials<T: FieldElement>(
-    analyzed: &Analyzed<T>,
-) -> Vec<Identity<AlgebraicExpression<T>>> {
-    let intermediates = &analyzed
-        .intermediate_polys_in_source_order()
-        .iter()
-        .map(|(symbol, def)| (symbol.id, def))
-        .collect();
-
-    substitute_intermediate(analyzed.identities.clone(), intermediates)
-}
-
-/// 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<AlgebraicExpression<T>>>,
-    intermediate_polynomials: &HashMap<u64, &AlgebraicExpression<T>>,
-) -> Vec<Identity<AlgebraicExpression<T>>> {
-    identities
-        .into_iter()
-        .scan(HashMap::default(), |cache, mut identity| {
-            identity.post_visit_expressions_mut(&mut |e| {
-                if let AlgebraicExpression::Reference(poly) = e {
-                    match poly.poly_id.ptype {
-                        PolynomialType::Committed => {}
-                        PolynomialType::Constant => {}
-                        PolynomialType::Intermediate => {
-                            // recursively inline intermediate polynomials, updating the cache
-                            *e = inlined_expression_from_intermediate_poly_id(
-                                poly.poly_id.id,
-                                intermediate_polynomials,
-                                cache,
-                            );
-                        }
-                    }
-                }
-            });
-            Some(identity)
-        })
-        .collect()
-}
-
-/// Recursively inlines intermediate polynomials inside an expression and returns the new expression
-/// This uses a cache to avoid resolving an intermediate polynomial twice
-fn inlined_expression_from_intermediate_poly_id<T: Copy>(
-    poly_id: u64,
-    intermediate_polynomials: &HashMap<u64, &AlgebraicExpression<T>>,
-    cache: &mut HashMap<u64, AlgebraicExpression<T>>,
-) -> AlgebraicExpression<T> {
-    if let Some(e) = cache.get(&poly_id) {
-        return e.clone();
-    }
-    let mut expr = intermediate_polynomials[&poly_id].clone();
-    expr.post_visit_expressions_mut(&mut |e| {
-        if let AlgebraicExpression::Reference(r) = e {
-            match r.poly_id.ptype {
-                PolynomialType::Committed => {}
-                PolynomialType::Constant => {}
-                PolynomialType::Intermediate => {
-                    // read from the cache, if no cache hit, compute the inlined expression
-                    *e = cache.get(&r.poly_id.id).cloned().unwrap_or_else(|| {
-                        inlined_expression_from_intermediate_poly_id(
-                            r.poly_id.id,
-                            intermediate_polynomials,
-                            cache,
-                        )
-                    });
-                }
-            }
-        }
-    });
-    cache.insert(poly_id, expr.clone());
-    expr
-}
-
 #[cfg(test)]
 mod test {
     use number::GoldilocksField;