From 7a6d813f36bf28ba90e76342e0a76358a2e317d0 Mon Sep 17 00:00:00 2001
From: Dan Dore <dorebell@gmail.com>
Date: Fri, 3 May 2024 16:18:26 -0700
Subject: [PATCH] fixes soundness gap in LT32 chip

---
 alu_u32/src/lt/columns.rs | 3 +++
 alu_u32/src/lt/mod.rs     | 9 ++++++---
 alu_u32/src/lt/stark.rs   | 5 +++++
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/alu_u32/src/lt/columns.rs b/alu_u32/src/lt/columns.rs
index 55fa10c..a87a249 100644
--- a/alu_u32/src/lt/columns.rs
+++ b/alu_u32/src/lt/columns.rs
@@ -21,6 +21,9 @@ pub struct Lt32Cols<T> {
 
     pub is_lt: T,
     pub is_lte: T,
+
+    // inverse of input_1[i] - input_2[i] where i is the first byte that differs
+    pub diff_inv: T,
 }
 
 pub const NUM_LT_COLS: usize = size_of::<Lt32Cols<u8>>();
diff --git a/alu_u32/src/lt/mod.rs b/alu_u32/src/lt/mod.rs
index 1177c22..c1348d0 100644
--- a/alu_u32/src/lt/mod.rs
+++ b/alu_u32/src/lt/mod.rs
@@ -119,6 +119,9 @@ impl Lt32Chip {
         F: PrimeField,
     {
         // Set the input columns
+        debug_assert_eq!(a.0.len(), 4);
+        debug_assert_eq!(b.0.len(), 4);
+        debug_assert_eq!(c.0.len(), 4);
         cols.input_1 = b.transform(F::from_canonical_u8);
         cols.input_2 = c.transform(F::from_canonical_u8);
         cols.output = F::from_canonical_u8(a[3]);
@@ -133,9 +136,9 @@ impl Lt32Chip {
             for i in 0..10 {
                 cols.bits[i] = F::from_canonical_u16(z >> i & 1);
             }
-            if n < 4 {
-                cols.byte_flag[n] = F::one();
-            }
+            cols.byte_flag[n] = F::one();
+            // b[n] != c[n] always here, so the difference is never zero.
+            cols.diff_inv = (cols.input_1[n] - cols.input_2[n]).inverse();
         }
         cols.multiplicity = F::one();
     }
diff --git a/alu_u32/src/lt/stark.rs b/alu_u32/src/lt/stark.rs
index 0402b89..48684df 100644
--- a/alu_u32/src/lt/stark.rs
+++ b/alu_u32/src/lt/stark.rs
@@ -68,6 +68,11 @@ where
                 AB::Expr::from_canonical_u32(256) + local.input_1[i] - local.input_2[i],
                 bit_comp.clone(),
             );
+            // ensure that when the n-th byte flag is set, the n-th bytes are actually different
+            builder.when(local.byte_flag[i]).assert_eq(
+                (local.input_1[i] - local.input_2[i]) * local.diff_inv,
+                AB::Expr::one(),
+            );
             builder.assert_bool(local.byte_flag[i]);
         }