Remove constant intermediates (#2163)

See [added
test](https://github.com/powdr-labs/powdr/blob/48de0e832e14da0e5dd5377bc75c6eea55c1bcc6/pilopt/tests/optimizer.rs#L34)

This propagates knowledge about constant intermediates forward in source
order. We could run this backwards as well but I think in practice
intermediates are defined before usage.

---------

Co-authored-by: chriseth <[email protected]>

ast/src/analyzed/mod.rs

@@ -139,7 +139,7 @@ impl<T> Analyzed<T> {
         self.definitions_in_source_order(PolynomialType::Committed)
     }
 
-    fn intermediate_polys_in_source_order(
+    pub fn intermediate_polys_in_source_order(
         &self,
     ) -> impl Iterator<Item = &(Symbol, Vec<AlgebraicExpression<T>>)> {
         self.source_order.iter().filter_map(move |statement| {

pilopt/src/lib.rs

@@ -28,6 +28,7 @@ pub fn optimize<T: FieldElement>(mut pil_file: Analyzed<T>) -> Analyzed<T> {
     simplify_identities(&mut pil_file);
     extract_constant_lookups(&mut pil_file);
     remove_constant_witness_columns(&mut pil_file);
+    remove_constant_intermediate_columns(&mut pil_file);
     simplify_identities(&mut pil_file);
     remove_trivial_identities(&mut pil_file);
     remove_duplicate_identities(&mut pil_file);
@@ -466,6 +467,34 @@ fn remove_constant_witness_columns<T: FieldElement>(pil_file: &mut Analyzed<T>)
     substitute_polynomial_references(pil_file, constant_polys);
 }
 
+/// Identifies intermediate columns that are constrained to a single value, replaces every
+/// reference to this column by the value and deletes the column.
+fn remove_constant_intermediate_columns<T: FieldElement>(pil_file: &mut Analyzed<T>) {
+    let intermediate_polys = pil_file
+        .intermediate_polys_in_source_order()
+        .filter_map(|(symbol, definitions)| {
+            let mut symbols_and_definitions = symbol.array_elements().zip_eq(definitions);
+            match symbol.is_array() {
+                true => None,
+                false => {
+                    let ((name, poly_id), definition) = symbols_and_definitions.next().unwrap();
+                    match definition {
+                        AlgebraicExpression::Number(value) => {
+                            log::debug!(
+                                "Determined intermediate column {name} to be constant {value}. Removing.",
+                            );
+                            Some(((name.clone(), poly_id), value.to_arbitrary_integer()))
+                        }
+                        _ => None,
+                    }
+                }
+            }
+        })
+        .collect::<Vec<((String, PolyID), _)>>();
+
+    substitute_polynomial_references(pil_file, intermediate_polys);
+}
+
 /// Substitutes all references to certain polynomials by the given field elements.
 fn substitute_polynomial_references<T: FieldElement>(
     pil_file: &mut Analyzed<T>,

pilopt/tests/optimizer.rs

@@ -30,6 +30,23 @@ fn replace_fixed() {
     assert_eq!(optimized, expectation);
 }
 
+#[test]
+fn replace_intermediate() {
+    let input = r#"namespace N(65536);
+    col witness X;
+    col intermediate = 1;
+    col other_intermediate = (intermediate - 1) * X;
+    X' = X + intermediate + other_intermediate;
+"#;
+    let expectation = r#"namespace N(65536);
+    col witness X;
+    col other_intermediate = 0;
+    N::X' = N::X + 1 + N::other_intermediate;
+"#;
+    let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
+    assert_eq!(optimized, expectation);
+}
+
 #[test]
 fn deduplicate_fixed() {
     let input = r#"namespace N(65536);