optimizer: deduplicate fixed columns (#2102)

With bus lookups, we insert `is_first` once per lookup. Only keep a single column per syntactic value. Depends on #2103
powdr-labs · Nov 22, 2024 · c45a4d5 · c45a4d5
1 parent 51cff4c
commit c45a4d5
Show file tree

Hide file tree

Showing 7 changed files with 178 additions and 32 deletions.
diff --git a/ast/src/analyzed/mod.rs b/ast/src/analyzed/mod.rs
@@ -796,7 +796,7 @@ pub enum SymbolKind {
     Other(),
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq, Hash)]
 pub enum FunctionValueDefinition {
     Array(ArrayExpression<Reference>),
     Expression(TypedExpression),
@@ -1152,7 +1152,9 @@ impl<T> SelectedExpressions<T> {
 pub type Expression = parsed::Expression<Reference>;
 pub type TypedExpression = crate::parsed::TypedExpression<Reference, u64>;
 
-#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq, Hash, PartialOrd, Ord,
+)]
 pub enum Reference {
     LocalVar(u64, String),
     Poly(PolynomialReference),
@@ -1577,7 +1579,9 @@ impl<T> From<T> for AlgebraicExpression<T> {
 /// Reference to a symbol with optional type arguments.
 /// Named `PolynomialReference` for historical reasons, it can reference
 /// any symbol.
-#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq, PartialOrd, Ord, Hash,
+)]
 pub struct PolynomialReference {
     /// Absolute name of the symbol.
     pub name: String,

diff --git a/ast/src/parsed/mod.rs b/ast/src/parsed/mod.rs
@@ -251,7 +251,9 @@ impl Children<Expression> for PilStatement {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub enum TypeDeclaration<E = u64> {
     Enum(EnumDeclaration<E>),
     Struct(StructDeclaration<E>),
@@ -289,7 +291,9 @@ impl<R> Children<Expression<R>> for TypeDeclaration<u64> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct StructDeclaration<E = u64> {
     pub name: String,
     pub type_vars: TypeBounds,
@@ -326,7 +330,9 @@ impl<R> Children<Expression<R>> for StructDeclaration<u64> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct EnumDeclaration<E = u64> {
     pub name: String,
     pub type_vars: TypeBounds,
@@ -351,7 +357,9 @@ impl<R> Children<Expression<R>> for EnumDeclaration<Expression<R>> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct EnumVariant<E = u64> {
     pub name: String,
     pub fields: Option<Vec<Type<E>>>,
@@ -452,7 +460,9 @@ impl<R> Children<Expression<R>> for TraitImplementation<Expression<R>> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct NamedExpression<Expr> {
     pub name: String,
     pub body: Expr,
@@ -476,7 +486,9 @@ impl<R> Children<Expression<R>> for NamedExpression<Arc<Expression<R>>> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct TraitDeclaration<E = u64> {
     pub name: String,
     pub type_vars: Vec<String>,
@@ -502,7 +514,9 @@ impl<R> Children<Expression<R>> for TraitDeclaration<Expression<R>> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct NamedType<E = u64> {
     pub name: String,
     pub ty: Type<E>,
@@ -548,7 +562,9 @@ impl<T> Children<Expression<T>> for SelectedExpressions<Expression<T>> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub enum Expression<Ref = NamespacedPolynomialReference> {
     Reference(SourceRef, Ref),
     PublicReference(SourceRef, String),
@@ -641,7 +657,9 @@ impl_source_reference!(
     Expression
 );
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct UnaryOperation<E = Expression<NamespacedPolynomialReference>> {
     pub op: UnaryOperator,
     pub expr: Box<E>,
@@ -663,7 +681,9 @@ impl<E> Children<E> for UnaryOperation<E> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct BinaryOperation<E = Expression<NamespacedPolynomialReference>> {
     pub left: Box<E>,
     pub op: BinaryOperator,
@@ -686,7 +706,9 @@ impl<E> Children<E> for BinaryOperation<E> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct Number {
     #[schemars(skip)]
     pub value: BigUint,
@@ -739,7 +761,9 @@ impl<Ref> Expression<Ref> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct MatchExpression<E = Expression<NamespacedPolynomialReference>> {
     pub scrutinee: Box<E>,
     pub arms: Vec<MatchArm<E>>,
@@ -766,7 +790,9 @@ impl<E> Children<E> for MatchExpression<E> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct BlockExpression<E> {
     pub statements: Vec<StatementInsideBlock<E>>,
     pub expr: Option<Box<E>>,
@@ -959,7 +985,9 @@ impl NamespacedPolynomialReference {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct LambdaExpression<E = Expression<NamespacedPolynomialReference>> {
     pub kind: FunctionKind,
     pub params: Vec<Pattern>,
@@ -985,15 +1013,17 @@ impl<E> Children<E> for LambdaExpression<E> {
 }
 
 #[derive(
-    Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, JsonSchema,
+    Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, JsonSchema, Hash,
 )]
 pub enum FunctionKind {
     Pure,
     Constr,
     Query,
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct ArrayLiteral<E = Expression<NamespacedPolynomialReference>> {
     pub items: Vec<E>,
 }
@@ -1156,7 +1186,9 @@ impl BinaryOperator {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct IndexAccess<E = Expression<NamespacedPolynomialReference>> {
     pub array: Box<E>,
     pub index: Box<E>,
@@ -1178,7 +1210,9 @@ impl<E> Children<E> for IndexAccess<E> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct FunctionCall<E = Expression<NamespacedPolynomialReference>> {
     pub function: Box<E>,
     pub arguments: Vec<E>,
@@ -1200,7 +1234,9 @@ impl<E> Children<E> for FunctionCall<E> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct MatchArm<E = Expression<NamespacedPolynomialReference>> {
     pub pattern: Pattern,
     pub value: E,
@@ -1216,7 +1252,9 @@ impl<E> Children<E> for MatchArm<E> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct IfExpression<E = Expression<NamespacedPolynomialReference>> {
     pub condition: Box<E>,
     pub body: Box<E>,
@@ -1249,7 +1287,9 @@ impl<E> Children<E> for IfExpression<E> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct StructExpression<Ref = NamespacedPolynomialReference> {
     pub name: Ref,
     pub fields: Vec<NamedExpression<Box<Expression<Ref>>>>,
@@ -1271,7 +1311,9 @@ impl<R> Children<Expression<R>> for StructExpression<R> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub enum StatementInsideBlock<E = Expression<NamespacedPolynomialReference>> {
     // TODO add a source ref here.
     LetStatement(LetStatementInsideBlock<E>),
@@ -1294,7 +1336,9 @@ impl<E> Children<E> for StatementInsideBlock<E> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct LetStatementInsideBlock<E = Expression<NamespacedPolynomialReference>> {
     pub pattern: Pattern,
     pub ty: Option<Type<u64>>,
@@ -1352,7 +1396,9 @@ impl Children<Expression> for FunctionDefinition {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub enum ArrayExpression<Ref = NamespacedPolynomialReference> {
     Value(Vec<Expression<Ref>>),
     RepeatedValue(Vec<Expression<Ref>>),
@@ -1511,7 +1557,9 @@ impl<Ref> Children<Expression<Ref>> for ArrayExpression<Ref> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub enum Pattern {
     CatchAll(SourceRef), // "_", matches a single value
     Ellipsis(SourceRef), // "..", matches a series of values, only valid inside array patterns
@@ -1604,7 +1652,9 @@ impl SourceReference for Pattern {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct TypedExpression<Ref = NamespacedPolynomialReference, E = Expression<Ref>> {
     pub e: Expression<Ref>,
     pub type_scheme: Option<TypeScheme<E>>,

diff --git a/ast/src/parsed/types.rs b/ast/src/parsed/types.rs
@@ -422,7 +422,9 @@ impl From<FunctionType> for Type {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+#[derive(
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
+)]
 pub struct TypeScheme<E = u64> {
     /// Type variables and their trait bounds.
     pub vars: TypeBounds,
@@ -481,7 +483,7 @@ impl From<Type> for TypeScheme {
 }
 
 #[derive(
-    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Default, Serialize, Deserialize, JsonSchema,
+    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Default, Serialize, Deserialize, JsonSchema, Hash,
 )]
 // TODO bounds should be SymbolPaths in the future.
 pub struct TypeBounds(Vec<(String, BTreeSet<String>)>);

diff --git a/parser-util/src/lib.rs b/parser-util/src/lib.rs
@@ -4,6 +4,7 @@
 
 use std::{
     fmt::{self, Debug, Formatter},
+    hash::Hash,
     sync::Arc,
 };
 
@@ -42,6 +43,10 @@ impl PartialEq for SourceRef {
 
 impl Eq for SourceRef {}
 
+impl Hash for SourceRef {
+    fn hash<H: std::hash::Hasher>(&self, _: &mut H) {}
+}
+
 impl SourceRef {
     pub fn unknown() -> Self {
         Default::default()

diff --git a/pilopt/Cargo.toml b/pilopt/Cargo.toml
@@ -13,6 +13,7 @@ powdr-number.workspace = true
 
 log = "0.4.17"
 pretty_assertions = "1.4.0"
+itertools = "0.13.0"
 
 [dev-dependencies]
 powdr-pil-analyzer.workspace = true