From 091047f7f10cae082716f3738ad59a583835f7b6 Mon Sep 17 00:00:00 2001
From: Linda Guiga <101227802+LindaGuiga@users.noreply.github.com>
Date: Tue, 28 Jan 2025 17:39:13 +0100
Subject: [PATCH] Fix padding in LookupTableGate (#1656)

* Fix soundness in Lookups

* FIx recursion
---
 plonky2/src/gates/lookup.rs         |  6 ++++--
 plonky2/src/gates/lookup_table.rs   |  8 +++++---
 plonky2/src/plonk/vanishing_poly.rs | 21 +++++++++++++++++++--
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/plonky2/src/gates/lookup.rs b/plonky2/src/gates/lookup.rs
index 6a28745c..c0ece3f6 100644
--- a/plonky2/src/gates/lookup.rs
+++ b/plonky2/src/gates/lookup.rs
@@ -197,8 +197,10 @@ impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F, D> for Loo
         let get_wire = |wire: usize| -> F { witness.get_target(Target::wire(self.row, wire)) };
 
         let input_val = get_wire(LookupGate::wire_ith_looking_inp(self.slot_nb));
-        let (input, output) = self.lut[input_val.to_canonical_u64() as usize];
-        if input_val == F::from_canonical_u16(input) {
+        if (input_val.to_canonical_u64() as usize) < self.lut.len()
+            && input_val == F::from_canonical_u16(self.lut[input_val.to_canonical_u64() as usize].0)
+        {
+            let (_, output) = self.lut[input_val.to_canonical_u64() as usize];
             let output_val = F::from_canonical_u16(output);
 
             let out_wire = Target::wire(self.row, LookupGate::wire_ith_looking_out(self.slot_nb));
diff --git a/plonky2/src/gates/lookup_table.rs b/plonky2/src/gates/lookup_table.rs
index 97b49595..dde0ed6e 100644
--- a/plonky2/src/gates/lookup_table.rs
+++ b/plonky2/src/gates/lookup_table.rs
@@ -224,9 +224,11 @@ impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F, D> for Loo
             out_buffer.set_target(slot_input_target, F::from_canonical_usize(input as usize))?;
             out_buffer.set_target(slot_output_target, F::from_canonical_usize(output as usize))
         } else {
-            // Pad with zeros.
-            out_buffer.set_target(slot_input_target, F::ZERO)?;
-            out_buffer.set_target(slot_output_target, F::ZERO)
+            // Pad with first element in the LUT.
+            assert!(!self.lut.is_empty(), "Empty LUTs are not supported.");
+            let (input, output) = self.lut[0];
+            out_buffer.set_target(slot_input_target, F::from_canonical_usize(input as usize))?;
+            out_buffer.set_target(slot_output_target, F::from_canonical_usize(output as usize))
         }
     }
 
diff --git a/plonky2/src/plonk/vanishing_poly.rs b/plonky2/src/plonk/vanishing_poly.rs
index a5fac4ab..5d26f68e 100644
--- a/plonky2/src/plonk/vanishing_poly.rs
+++ b/plonky2/src/plonk/vanishing_poly.rs
@@ -36,10 +36,17 @@ pub(crate) fn get_lut_poly<F: RichField + Extendable<D>, const D: usize>(
     let b = deltas[LookupChallenges::ChallengeB as usize];
     let mut coeffs = Vec::with_capacity(common_data.luts[lut_index].len());
     let n = common_data.luts[lut_index].len();
+    let nb_padded_elts = LookupTableGate::num_slots(&common_data.config)
+        - n % LookupTableGate::num_slots(&common_data.config);
+    let (padding_inp, padding_out) = common_data.luts[lut_index][0];
     for (input, output) in common_data.luts[lut_index].iter() {
         coeffs.push(F::from_canonical_u16(*input) + b * F::from_canonical_u16(*output));
     }
-    coeffs.append(&mut vec![F::ZERO; degree - n]);
+    // Padding with the first element of the LUT.
+    for _ in 0..nb_padded_elts {
+        coeffs.push(F::from_canonical_u16(padding_inp) + b * F::from_canonical_u16(padding_out));
+    }
+    coeffs.append(&mut vec![F::ZERO; degree - (n + nb_padded_elts)]);
     coeffs.reverse();
     PolynomialCoeffs::new(coeffs)
 }
@@ -756,6 +763,9 @@ pub(crate) fn get_lut_poly_circuit<F: RichField + Extendable<D>, const D: usize>
     let b = deltas[LookupChallenges::ChallengeB as usize];
     let delta = deltas[LookupChallenges::ChallengeDelta as usize];
     let n = common_data.luts[lut_index].len();
+    let nb_padded_elts = LookupTableGate::num_slots(&common_data.config)
+        - n % LookupTableGate::num_slots(&common_data.config);
+    let (padding_inp, padding_out) = common_data.luts[lut_index][0];
     let mut coeffs: Vec<Target> = common_data.luts[lut_index]
         .iter()
         .map(|(input, output)| {
@@ -763,7 +773,14 @@ pub(crate) fn get_lut_poly_circuit<F: RichField + Extendable<D>, const D: usize>
             builder.add_const(temp, F::from_canonical_u16(*input))
         })
         .collect();
-    for _ in n..degree {
+
+    // Padding with the first element of the LUT.
+    for _ in 0..nb_padded_elts {
+        let temp = builder.mul_const(F::from_canonical_u16(padding_out), b);
+        let temp = builder.add_const(temp, F::from_canonical_u16(padding_inp));
+        coeffs.push(temp);
+    }
+    for _ in (n + nb_padded_elts)..degree {
         coeffs.push(builder.zero());
     }
     coeffs.reverse();