Address comments.

valida-xyz · Jan 17, 2024 · a693dc9 · a693dc9
1 parent 9e48bf3
commit a693dc9
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/alu_u32/src/com/stark.rs b/alu_u32/src/com/stark.rs
@@ -30,7 +30,7 @@ where
                 .input_1
                 .into_iter()
                 .zip(local.input_2)
-                .map(|(a, b)| (a - b) * (a - b))
+                .map(|(a, b)| (a - b).square())
                 .sum::<AB::Expr>(),
         );
         builder.assert_bool(local.not_equal);
@@ -40,7 +40,9 @@ where
         builder.assert_bool(local.is_eq);
         builder.assert_bool(local.is_ne + local.is_eq);
 
-        builder.when(local.is_ne).assert_one(local.not_equal);
-        builder.when(local.is_eq).assert_zero(local.not_equal);
+        builder.assert_eq(
+            local.output,
+            local.is_ne * local.not_equal + local.is_eq * (AB::Expr::one() - local.not_equal),
+        )
     }
 }