Find and replace functions that compute "are all equal zero"

autoprecompiles/src/constraint_optimizer.rs

@@ -3,12 +3,23 @@ use std::{collections::HashSet, fmt::Display, hash::Hash};
 use inliner::DegreeBound;
 use num_traits::Zero;
 use powdr_constraint_solver::{
-    constraint_system::BusInteractionHandler, grouped_expression::GroupedExpression, inliner,
-    journaling_constraint_system::JournalingConstraintSystem, solver::solve_system,
+    constraint_system::{BusInteractionHandler, ConstraintSystem},
+    grouped_expression::GroupedExpression,
+    inliner,
+    journaling_constraint_system::JournalingConstraintSystem,
+    solver::solve_system,
 };
 use powdr_number::FieldElement;
 
-use crate::stats_logger::StatsLogger;
+use crate::{
+    constraint_optimizer::{
+        equal_zero_checks::replace_equal_zero_checks, reachability::reachable_variables,
+    },
+    stats_logger::StatsLogger,
+};
+
+mod equal_zero_checks;
+mod reachability;
 
 #[derive(Debug)]
 pub enum Error {
@@ -32,6 +43,7 @@ pub fn optimize_constraints<P: FieldElement, V: Ord + Clone + Eq + Hash + Displa
     bus_interaction_handler: impl BusInteractionHandler<P> + IsBusStateful<P> + Clone,
     degree_bound: DegreeBound,
     stats_logger: &mut StatsLogger,
+    new_var: &mut impl FnMut() -> V,
 ) -> Result<JournalingConstraintSystem<P, V>, Error> {
     let constraint_system =
         solver_based_optimization(constraint_system, bus_interaction_handler.clone())?;
@@ -41,6 +53,9 @@ pub fn optimize_constraints<P: FieldElement, V: Ord + Clone + Eq + Hash + Displa
         remove_disconnected_columns(constraint_system, bus_interaction_handler.clone());
     stats_logger.log("removing disconnected columns", &constraint_system);
 
+    let constraint_system =
+        replace_equal_zero_checks(constraint_system, bus_interaction_handler.clone(), new_var);
+
     let constraint_system =
         inliner::replace_constrained_witness_columns(constraint_system, degree_bound);
     stats_logger.log("in-lining witness columns", &constraint_system);
@@ -84,36 +99,18 @@ fn solver_based_optimization<T: FieldElement, V: Clone + Ord + Hash + Display>(
 /// change the statement being proven.
 fn remove_disconnected_columns<T: FieldElement, V: Clone + Ord + Hash + Display>(
     mut constraint_system: JournalingConstraintSystem<T, V>,
-    bus_interaction_handler: impl IsBusStateful<T>,
+    bus_interaction_handler: impl IsBusStateful<T> + Clone,
 ) -> JournalingConstraintSystem<T, V> {
-    // Initialize variables_to_keep with any variables that appear in stateful bus interactions.
-    let mut variables_to_keep = constraint_system
-        .system()
-        .bus_interactions
-        .iter()
-        .filter(|bus_interaction| {
-            let bus_id = bus_interaction.bus_id.try_to_number().unwrap();
-            bus_interaction_handler.is_stateful(bus_id)
-        })
-        .flat_map(|bus_interaction| bus_interaction.referenced_variables())
-        .cloned()
-        .collect::<HashSet<_>>();
-
-    // Any variable that is connected to a variable in variables_to_keep must also be kept.
-    loop {
-        let size_before = variables_to_keep.len();
-        for expr in constraint_system.system().iter() {
-            if expr
-                .referenced_variables()
-                .any(|var| variables_to_keep.contains(var))
-            {
-                variables_to_keep.extend(expr.referenced_variables().cloned());
-            }
-        }
-        if variables_to_keep.len() == size_before {
-            break;
-        }
-    }
+    let initial_variables = variables_in_stateful_bus_interactions(
+        constraint_system.system(),
+        bus_interaction_handler.clone(),
+    )
+    .cloned();
+    let variables_to_keep = reachable_variables(
+        initial_variables,
+        std::iter::empty(),
+        constraint_system.system(),
+    );
 
     constraint_system.retain_algebraic_constraints(|constraint| {
         constraint
@@ -134,6 +131,21 @@ fn remove_disconnected_columns<T: FieldElement, V: Clone + Ord + Hash + Display>
     constraint_system
 }
 
+/// Returns an iterator over all variables that are referenced in stateful bus interactions.
+fn variables_in_stateful_bus_interactions<'a, P: FieldElement, V: Ord + Clone + Eq + Hash>(
+    constraint_system: &'a ConstraintSystem<P, V>,
+    bus_interaction_handler: impl IsBusStateful<P> + 'a,
+) -> impl Iterator<Item = &'a V> {
+    constraint_system
+        .bus_interactions
+        .iter()
+        .filter(move |bus_interaction| {
+            let bus_id = bus_interaction.bus_id.try_to_number().unwrap();
+            bus_interaction_handler.is_stateful(bus_id)
+        })
+        .flat_map(|bus_interaction| bus_interaction.referenced_variables())
+}
+
 fn remove_trivial_constraints<P: FieldElement, V: PartialEq + Clone + Hash + Ord>(
     mut constraint_system: JournalingConstraintSystem<P, V>,
 ) -> JournalingConstraintSystem<P, V> {