improving comments

src/compiler/cse.rs

@@ -20,6 +20,14 @@ use crate::{
     wit_gen::NullTraceGenerator,
 };
 
+
+/// Common Subexpression Elimination (CSE) optimization.
+/// This optimization replaces common subexpressions with new internal signals for the step type.
+/// This is done by each time finding the optimal subexpression to replace and creating a new signal
+/// for it and replacing it in all constraints.
+/// The process is repeated until no more common subexpressions are found.
+/// Equivalent expressions are found by hashing the expressions with random assignments to the queriables. Using
+/// the Schwartz-Zippel lemma, we can determine if two expressions are equivalent with high probability.
 pub(super) fn cse<F: Field + Hash>(
     mut circuit: SBPIR<F, NullTraceGenerator>,
 ) -> SBPIR<F, NullTraceGenerator> {
@@ -67,6 +75,7 @@ pub(super) fn cse<F: Field + Hash>(
                 let (common_se, decomp) =
                     create_common_ses_signal(&common_expr, &mut signal_factory);
 
+                // Add the new signal to the step type and a constraint for it
                 decomp.auto_signals.iter().for_each(|(q, expr)| {
                     if let Queriable::Internal(signal) = q {
                         step_type_with_hash.add_internal(signal.clone());
@@ -196,16 +205,6 @@ impl<F> poly::SignalFactory<Queriable<F>> for SignalFactory<F> {
     }
 }
 
-// cse
-// 0. collects a set of expressions Expr<F, V, ()>
-// 1. turns all Expr<F, V, ()> into Expr<F, V, HashResult>
-// 2. traverse all the expressions and find common subexpressions counting the number of times they
-//    appear
-// 3. Sort the common subexpressions by the degree and number of times they appear
-// 4. Replace the best common subexpression with a signal
-// 5. Repeat until no common subexpressions are found
-//
-
 #[cfg(test)]
 mod test {
     use std::collections::HashSet;

src/poly/cse.rs

@@ -3,6 +3,7 @@ use crate::field::Field;
 use super::{ConstrDecomp, Expr, HashResult, SignalFactory};
 use std::{fmt::Debug, hash::Hash};
 
+/// This function replaces a common subexpression in an expression with a new signal.
 pub fn replace_expr<F: Field + Hash, V: Clone + Eq + Hash + Debug, SF: SignalFactory<V>>(
     expr: &Expr<F, V, HashResult>,
     common_se: &Expr<F, V, HashResult>,
@@ -15,6 +16,7 @@ pub fn replace_expr<F: Field + Hash, V: Clone + Eq + Hash + Debug, SF: SignalFac
     (new_expr, ConstrDecomp::default())
 }
 
+/// This function creates a new signal for a common subexpression.
 pub fn create_common_ses_signal<
     F: Field,
     V: Clone + PartialEq + Eq + Hash + Debug,
@@ -30,6 +32,7 @@ pub fn create_common_ses_signal<
     (Expr::Query(signal, common_se.meta().clone()), decomp)
 }
 
+/// This function replaces a common subexpression in an expression with a new signal.
 fn replace_subexpr<F: Field + Hash, V: Clone + Eq + Hash + Debug, SF: SignalFactory<V>>(
     expr: &Expr<F, V, HashResult>,
     common_se: &Expr<F, V, HashResult>,
@@ -44,13 +47,11 @@ fn replace_subexpr<F: Field + Hash, V: Clone + Eq + Hash + Debug, SF: SignalFact
         return expr.clone();
     }
 
-    // If the expression is the same as the common subexpression, create a new signal and return it
+    // If the expression is the same as the common subexpression return the signal
     if expr.meta().hash == common_expr_hash {
-        // Find the signal or create a new signal for the expression
         return common_se.clone();
     } else {
-        // Only recurse if we haven't found a match and the expression could potentially contain the
-        // common subexpression
+        // Recursively apply the function to the subexpressions
         expr.apply_subexpressions(|se| replace_subexpr(se, common_se, signal_factory, decomp))
     }
 }