block statement fix

Author: ShuangWu121

autoprecompiles/src/lib.rs

@@ -434,15 +434,18 @@ pub fn build<A: Adapter>(
     let range_constraints = &execution_stats.column_ranges_by_pc;
     let equivalence_classes_by_block = &execution_stats.equivalence_classes_by_block;
 
-    // Mapping (instruction index, column index) -> AlgebraicReference
+    // Mapping (pc index, column index) -> AlgebraicReference
     let reverse_subs = column_allocator
         .subs
         .iter()
         .enumerate()
         .flat_map(|(instr_index, subs)| {
-            subs.iter()
-                .enumerate()
-                .map(move |(col_index, &poly_id)| (poly_id, (instr_index, col_index)))
+            subs.iter().enumerate().map(move |(col_index, &poly_id)| {
+                (
+                    poly_id,
+                    (block.start_pc as usize + (instr_index * 4), col_index),
+                )
+            })
         })
         .collect::<BTreeMap<_, _>>();
     let algebraic_references = machine
@@ -458,8 +461,12 @@ pub fn build<A: Adapter>(
         for (col_index, range) in range_constraints.iter().enumerate() {
             if range.0 == range.1 {
                 let value = A::PowdrField::from(range.0 as u64);
-                let Some(reference) = algebraic_references.get(&(i, col_index)).cloned() else {
-                    println!("Missing reference for (i: {}, col_index: {})", i, col_index);
+                let Some(reference) = algebraic_references.get(&(pc as usize, col_index)).cloned()
+                else {
+                    println!(
+                        "Missing reference for (pc: {}, col_index: {})",
+                        pc, col_index
+                    );
                     continue;
                 };
                 let constraint =
@@ -472,33 +479,32 @@ pub fn build<A: Adapter>(
         }
     }
 
-    // TODO: This can cause unsatisfiable constraints somehow.
-    // if let Some(equivalence_classes) = equivalence_classes_by_block.get(&block.start_pc) {
-    //     for equivalence_class in equivalence_classes {
-    //         let first = equivalence_class.first().unwrap();
-    //         let Some(first_ref) = algebraic_references.get(first).cloned() else {
-    //             println!(
-    //                 "Missing reference for (i: {}, col_index: {})",
-    //                 first.0, first.1
-    //             );
-    //             continue;
-    //         };
-    //         for other in equivalence_class.iter().skip(1) {
-    //             let Some(other_ref) = algebraic_references.get(other).cloned() else {
-    //                 println!(
-    //                     "Missing reference for (i: {}, col_index: {})",
-    //                     other.0, other.1
-    //                 );
-    //                 continue;
-    //             };
-    //             let constraint = AlgebraicExpression::Reference(first_ref.clone())
-    //                 - AlgebraicExpression::Reference(other_ref.clone());
-    //             machine
-    //                 .constraints
-    //                 .push(SymbolicConstraint { expr: constraint });
-    //         }
-    //     }
-    // }
+    if let Some(equivalence_classes) = equivalence_classes_by_block.get(&block.start_pc) {
+        for equivalence_class in equivalence_classes {
+            let first = equivalence_class.first().unwrap();
+            let Some(first_ref) = algebraic_references.get(first).cloned() else {
+                println!(
+                    "Missing reference for (i: {}, col_index: {})",
+                    first.0, first.1
+                );
+                continue;
+            };
+            for other in equivalence_class.iter().skip(1) {
+                let Some(other_ref) = algebraic_references.get(other).cloned() else {
+                    println!(
+                        "Missing reference for (i: {}, col_index: {})",
+                        other.0, other.1
+                    );
+                    continue;
+                };
+                let constraint = AlgebraicExpression::Reference(first_ref.clone())
+                    - AlgebraicExpression::Reference(other_ref.clone());
+                machine
+                    .constraints
+                    .push(SymbolicConstraint { expr: constraint });
+            }
+        }
+    }
 
     let labels = [("apc_start_pc", block.start_pc.to_string())];
     metrics::counter!("before_opt_cols", &labels)