refactor: various AIR tweaks (#311)

core/src/air/builder.rs

@@ -32,9 +32,9 @@ pub trait EmptyMessageBuilder: AirBuilder {}
 
 /// A trait which contains basic methods for building an AIR.
 pub trait BaseAirBuilder: AirBuilder + MessageBuilder<AirInteraction<Self::Expr>> {
-    /// Returns a sub-builder whose constraints are enforced only when condition is one.
+    /// Returns a sub-builder whose constraints are enforced only when `condition` is not one.
     fn when_not<I: Into<Self::Expr>>(&mut self, condition: I) -> FilteredAirBuilder<Self> {
-        self.when(Self::Expr::from(Self::F::one()) - condition.into())
+        self.when_ne(condition, Self::F::one())
     }
 
     /// Asserts that an iterator of expressions are all equal.

core/src/alu/divrem/mod.rs

@@ -546,13 +546,13 @@ where
                         );
                     builder
                         .when(not_overflow.clone())
-                        .when(one.clone() - local.b_neg)
-                        .assert_eq(c_times_quotient_plus_remainder[i].clone(), zero.clone());
+                        .when_ne(one.clone(), local.b_neg)
+                        .assert_zero(c_times_quotient_plus_remainder[i].clone());
 
                     // The only exception to the upper-4-byte check is the overflow case.
                     builder
                         .when(local.is_overflow)
-                        .assert_eq(c_times_quotient_plus_remainder[i].clone(), zero.clone());
+                        .assert_zero(c_times_quotient_plus_remainder[i].clone());
                 }
             }
         }

core/src/alu/divrem/utils.rs

@@ -72,7 +72,7 @@ pub fn eval_abs_value<AB>(
             .assert_eq(value[i] + abs_value[i], exp_sum_if_negative.clone());
 
         builder
-            .when(AB::Expr::one() - *is_negative)
+            .when_not(*is_negative)
             .assert_eq(value[i], abs_value[i]);
     }
 }

core/src/cpu/air/memory.rs

@@ -245,40 +245,33 @@ impl CpuChip {
             - memory_columns.offset_is_two
             - memory_columns.offset_is_three;
 
-        // Assert that the value flags are boolean
-        builder
-            .when(is_mem_op.clone())
-            .assert_bool(memory_columns.offset_is_one);
+        let mut filtered_builder = builder.when(is_mem_op);
 
-        builder
-            .when(is_mem_op.clone())
-            .assert_bool(memory_columns.offset_is_two);
-
-        builder
-            .when(is_mem_op.clone())
-            .assert_bool(memory_columns.offset_is_three);
+        // Assert that the value flags are boolean
+        filtered_builder.assert_bool(memory_columns.offset_is_one);
+        filtered_builder.assert_bool(memory_columns.offset_is_two);
+        filtered_builder.assert_bool(memory_columns.offset_is_three);
 
         // Assert that only one of the value flags is true
-        builder.when(is_mem_op.clone()).assert_eq(
+        filtered_builder.assert_one(
             offset_is_zero.clone()
                 + memory_columns.offset_is_one
                 + memory_columns.offset_is_two
                 + memory_columns.offset_is_three,
-            AB::Expr::one(),
         );
 
         // Assert that the correct value flag is set
-        builder
-            .when(is_mem_op.clone() * offset_is_zero)
-            .assert_eq(memory_columns.addr_offset, AB::Expr::zero());
-        builder
-            .when(is_mem_op.clone() * memory_columns.offset_is_one)
-            .assert_eq(memory_columns.addr_offset, AB::Expr::one());
-        builder
-            .when(is_mem_op.clone() * memory_columns.offset_is_two)
+        filtered_builder
+            .when(offset_is_zero)
+            .assert_zero(memory_columns.addr_offset);
+        filtered_builder
+            .when(memory_columns.offset_is_one)
+            .assert_one(memory_columns.addr_offset);
+        filtered_builder
+            .when(memory_columns.offset_is_two)
             .assert_eq(memory_columns.addr_offset, AB::Expr::two());
-        builder
-            .when(is_mem_op * memory_columns.offset_is_three)
+        filtered_builder
+            .when(memory_columns.offset_is_three)
             .assert_eq(memory_columns.addr_offset, AB::Expr::from_canonical_u8(3));
     }
 }

core/src/cpu/air/mod.rs

@@ -57,7 +57,7 @@ where
             AB::Expr::one() - local.selectors.imm_b,
         );
         builder
-            .when(AB::Expr::one() - local.selectors.imm_b)
+            .when_not(local.selectors.imm_b)
             .assert_word_eq(local.op_b_val(), *local.op_b_access.prev_value());
 
         builder.constraint_memory_access(
@@ -68,7 +68,7 @@ where
             AB::Expr::one() - local.selectors.imm_c,
         );
         builder
-            .when(AB::Expr::one() - local.selectors.imm_c)
+            .when_not(local.selectors.imm_c)
             .assert_word_eq(local.op_c_val(), *local.op_c_access.prev_value());
 
         // Write the `a` or the result to the first register described in the instruction unless

core/src/syscall/precompiles/edwards/ed_decompress.rs

@@ -185,7 +185,7 @@ impl<V: Copy> EdDecompressCols<V> {
             .assert_all_eq(self.neg_x.result, x_limbs);
         builder
             .when(self.is_real)
-            .when(AB::Expr::one() - sign.clone())
+            .when_not(sign.clone())
             .assert_all_eq(self.x.multiplication.result, x_limbs);
     }
 }