Merge pull request #686 from powdr-labs/prepare_expression_split

Prepare expression split
This commit is contained in:
chriseth
2023-10-11 09:15:58 +00:00
committed by GitHub
4 changed files with 55 additions and 61 deletions

View File

@@ -163,14 +163,16 @@ impl<T> Analyzed<T> {
poly.id = replacements[&poly_id].id;
}
});
self.pre_visit_expressions_mut(&mut |expr| {
let visitor = &mut |expr: &mut Expression<_>| {
if let Expression::Reference(Reference::Poly(poly)) = expr {
poly.poly_id = poly.poly_id.map(|poly_id| {
assert!(!to_remove.contains(&poly_id));
replacements[&poly_id]
});
}
});
};
self.post_visit_expressions_in_definitions_mut(visitor);
self.post_visit_expressions_in_identities_mut(visitor);
}
/// Adds a polynomial identity and returns the ID.
@@ -225,6 +227,34 @@ impl<T> Analyzed<T> {
retain
})
}
pub fn post_visit_expressions_in_identities_mut<F>(&mut self, f: &mut F)
where
F: FnMut(&mut Expression<T>),
{
self.identities
.iter_mut()
.for_each(|i| i.post_visit_expressions_mut(f));
}
pub fn post_visit_expressions_in_definitions_mut<F>(&mut self, f: &mut F)
where
F: FnMut(&mut Expression<T>),
{
// TODO add public inputs if we change them to expressions at some point.
self.definitions
.values_mut()
.for_each(|(_poly, definition)| match definition {
Some(FunctionValueDefinition::Mapping(e))
| Some(FunctionValueDefinition::Query(e)) => e.post_visit_expressions_mut(f),
Some(FunctionValueDefinition::Array(elements)) => elements
.iter_mut()
.flat_map(|e| e.pattern.iter_mut())
.for_each(|e| e.post_visit_expressions_mut(f)),
Some(FunctionValueDefinition::Expression(e)) => e.post_visit_expressions_mut(f),
None => {}
});
}
}
#[derive(Debug, Clone)]

View File

@@ -2,54 +2,6 @@ use crate::parsed::visitor::VisitOrder;
use super::*;
impl<T> ExpressionVisitable<parsed::Expression<T, Reference>> for Analyzed<T> {
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>,
{
// TODO add constants if we change them to expressions at some point.
self.definitions
.values_mut()
.try_for_each(|(_poly, definition)| match definition {
Some(FunctionValueDefinition::Mapping(e))
| Some(FunctionValueDefinition::Query(e)) => e.visit_expressions_mut(f, o),
Some(FunctionValueDefinition::Array(elements)) => elements
.iter_mut()
.flat_map(|e| e.pattern.iter_mut())
.try_for_each(|e| e.visit_expressions_mut(f, o)),
Some(FunctionValueDefinition::Expression(e)) => e.visit_expressions_mut(f, o),
None => ControlFlow::Continue(()),
})?;
self.identities
.iter_mut()
.try_for_each(|i| i.visit_expressions_mut(f, o))
}
fn visit_expressions<F, B>(&self, f: &mut F, o: VisitOrder) -> ControlFlow<B>
where
F: FnMut(&parsed::Expression<T, Reference>) -> ControlFlow<B>,
{
// TODO add constants if we change them to expressions at some point.
self.definitions
.values()
.try_for_each(|(_poly, definition)| match definition {
Some(FunctionValueDefinition::Mapping(e))
| Some(FunctionValueDefinition::Query(e)) => e.visit_expressions(f, o),
Some(FunctionValueDefinition::Array(elements)) => elements
.iter()
.flat_map(|e| e.pattern.iter())
.try_for_each(|e| e.visit_expressions(f, o)),
Some(FunctionValueDefinition::Expression(e)) => e.visit_expressions(f, o),
None => ControlFlow::Continue(()),
})?;
self.identities
.iter()
.try_for_each(|i| i.visit_expressions(f, o))
}
}
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

View File

@@ -77,11 +77,14 @@ impl<T: Copy> From<PILAnalyzer<T>> for Analyzed<T> {
reference.poly_id = Some(poly.into());
}
};
result.pre_visit_expressions_mut(&mut |e| {
let expr_visitor = &mut |e: &mut Expression<_>| {
if let Expression::Reference(Reference::Poly(reference)) = e {
assign_id(reference);
}
});
};
result.post_visit_expressions_in_definitions_mut(expr_visitor);
result.post_visit_expressions_in_identities_mut(expr_visitor);
// TODO at some point, merge public declarations with definitions as well.
result
.public_declarations
.values_mut()

View File

@@ -47,7 +47,7 @@ pub fn optimize_constants<T: FieldElement>(mut pil_file: Analyzed<T>) -> Analyze
/// Inlines references to symbols with a single constant value.
fn inline_constant_values<T: FieldElement>(pil_file: &mut Analyzed<T>) {
let constants = compute_constants(pil_file);
pil_file.post_visit_expressions_mut(&mut |e| {
let visitor = &mut |e: &mut Expression<_>| {
if let Expression::Reference(Reference::Poly(poly)) = e {
if !poly.next && poly.index.is_none() {
if let Some(value) = constants.get(&poly.name) {
@@ -55,12 +55,14 @@ fn inline_constant_values<T: FieldElement>(pil_file: &mut Analyzed<T>) {
}
}
}
});
};
pil_file.post_visit_expressions_in_definitions_mut(visitor);
pil_file.post_visit_expressions_in_identities_mut(visitor);
}
/// Substitutes expression that evaluate to a constant value.
fn evaluate_constant_subtrees<T: FieldElement>(pil_file: &mut Analyzed<T>) {
pil_file.post_visit_expressions_mut(&mut |e| match e {
let visitor = &mut |e: &mut Expression<_>| match e {
Expression::BinaryOperation(left, op, right) => {
if let (Expression::Number(l), Expression::Number(r)) = (left.as_ref(), right.as_ref())
{
@@ -73,7 +75,9 @@ fn evaluate_constant_subtrees<T: FieldElement>(pil_file: &mut Analyzed<T>) {
}
}
_ => {}
});
};
pil_file.post_visit_expressions_in_definitions_mut(visitor);
pil_file.post_visit_expressions_in_identities_mut(visitor);
}
/// Identifies fixed columns that only have a single value, replaces every
@@ -95,7 +99,7 @@ fn remove_constant_fixed_columns<T: FieldElement>(pil_file: &mut Analyzed<T>) {
})
.collect::<BTreeMap<PolyID, _>>();
pil_file.pre_visit_expressions_mut(&mut |e| {
let visitor = &mut |e: &mut Expression<_>| {
if let Expression::Reference(Reference::Poly(PolynomialReference {
name: _,
index,
@@ -108,7 +112,9 @@ fn remove_constant_fixed_columns<T: FieldElement>(pil_file: &mut Analyzed<T>) {
*e = Expression::Number(*value);
}
}
});
};
pil_file.post_visit_expressions_in_definitions_mut(visitor);
pil_file.post_visit_expressions_in_identities_mut(visitor);
pil_file.remove_polynomials(&constant_polys.keys().cloned().collect());
}
@@ -143,7 +149,8 @@ fn constant_value<T: FieldElement>(function: &FunctionValueDefinition<T>) -> Opt
/// Simplifies multiplications by zero and one.
fn simplify_expressions<T: FieldElement>(pil_file: &mut Analyzed<T>) {
pil_file.post_visit_expressions_mut(&mut simplify_expression_single);
pil_file.post_visit_expressions_in_definitions_mut(&mut simplify_expression_single);
pil_file.post_visit_expressions_in_identities_mut(&mut simplify_expression_single);
}
fn simplify_expression<T: FieldElement>(mut e: Expression<T>) -> Expression<T> {
@@ -315,7 +322,7 @@ fn remove_constant_witness_columns<T: FieldElement>(pil_file: &mut Analyzed<T>)
.filter_map(|expr| constrained_to_constant(expr))
.collect::<BTreeMap<PolyID, _>>();
pil_file.pre_visit_expressions_mut(&mut |e| {
let visitor = &mut |e: &mut Expression<_>| {
if let Expression::Reference(Reference::Poly(PolynomialReference {
name: _,
index,
@@ -328,7 +335,9 @@ fn remove_constant_witness_columns<T: FieldElement>(pil_file: &mut Analyzed<T>)
*e = Expression::Number(*value);
}
}
});
};
pil_file.post_visit_expressions_in_definitions_mut(visitor);
pil_file.post_visit_expressions_in_identities_mut(visitor);
pil_file.remove_polynomials(&constant_polys.keys().cloned().collect());
}