Require that row is enabled for non-poseidon lookups

src/constraint_builder.rs

@@ -1,7 +1,9 @@
+use crate::gadgets::poseidon::PoseidonLookup;
 use halo2_proofs::{
     arithmetic::FieldExt,
     plonk::{ConstraintSystem, SecondPhase},
 };
+use itertools::Itertools;
 
 mod binary_column;
 mod binary_query;
@@ -73,35 +75,58 @@ impl<F: FieldExt> ConstraintBuilder<F> {
         let condition = self
             .conditions
             .iter()
-            .skip(1) // Save a degree by skipping every row selector
             .fold(BinaryQuery::one(), |a, b| a.and(b.clone()));
-        let lookup = left
+        let mut lookup: Vec<_> = left
             .into_iter()
             .map(|q| q * condition.clone())
             .zip(right.into_iter())
             .collect();
+        // If condition is true, every_row_selector must be enabled.
+        lookup.push((condition.into(), self.every_row_selector().into()));
         self.lookups.push((name, lookup))
     }
 
-    pub fn add_lookup_with_default<const N: usize>(
+    pub fn poseidon_lookup(
         &mut self,
         name: &'static str,
-        left: [Query<F>; N],
-        right: [Query<F>; N],
-        default: [Query<F>; N],
+        [left, right, domain, hash]: [Query<F>; 4],
+        poseidon: &impl PoseidonLookup,
     ) {
         let condition = self
             .conditions
             .iter()
             .skip(1) // Save a degree by skipping every row selector
             .fold(BinaryQuery::one(), |a, b| a.and(b.clone()));
-        let lookup = left
-            .into_iter()
-            .zip(default.into_iter())
-            .map(|(a, b)| condition.select(a, b))
-            .zip(right.into_iter())
-            .collect();
-        self.lookups.push((name, lookup))
+        let extended_queries = [
+            Query::one(),
+            hash,
+            left,
+            right,
+            Query::zero(),
+            domain,
+            Query::one(),
+        ]
+        .map(|q| q * condition.clone());
+
+        let (q_enable, [hash, left, right, control, domain_spec, head_mark]) =
+            poseidon.lookup_columns();
+        let poseidon_lookup_queries = [
+            q_enable.current(),
+            hash.current(),
+            left.current(),
+            right.current(),
+            control.current(),
+            domain_spec.current(),
+            head_mark.current(),
+        ];
+
+        self.lookups.push((
+            name,
+            extended_queries
+                .into_iter()
+                .zip_eq(poseidon_lookup_queries)
+                .collect(),
+        ))
     }
 
     pub fn build_columns<const A: usize, const B: usize, const C: usize>(

src/gadgets/byte_representation.rs

@@ -219,7 +219,7 @@ mod test {
             layouter.assign_region(
                 || "",
                 |mut region| {
-                    for offset in 0..1024 {
+                    for offset in 0..(8 * 256) {
                         selector.enable(&mut region, offset);
                     }
                     byte_bit.assign(&mut region);

src/gadgets/canonical_representation.rs

@@ -250,7 +250,7 @@ mod test {
             layouter.assign_region(
                 || "",
                 |mut region| {
-                    for offset in 0..256 {
+                    for offset in 0..(8 * 256) {
                         selector.enable(&mut region, offset);
                     }
                     byte_bit.assign(&mut region);

src/gadgets/key_bit.rs

@@ -191,7 +191,7 @@ mod test {
             layouter.assign_region(
                 || "",
                 |mut region| {
-                    for offset in 0..32 {
+                    for offset in 0..(8 * 256) {
                         selector.enable(&mut region, offset);
                     }
 

src/gadgets/poseidon.rs

@@ -1,11 +1,10 @@
-use crate::constraint_builder::{AdviceColumn, ConstraintBuilder, FixedColumn, Query};
+use crate::constraint_builder::{AdviceColumn, FixedColumn};
+use halo2_proofs::plonk::{Advice, Column, Fixed};
+#[cfg(test)]
 use halo2_proofs::{
-    arithmetic::FieldExt,
-    plonk::{Advice, Column, Fixed},
+    arithmetic::FieldExt, circuit::Region, halo2curves::bn256::Fr, plonk::ConstraintSystem,
 };
 #[cfg(test)]
-use halo2_proofs::{circuit::Region, halo2curves::bn256::Fr, plonk::ConstraintSystem};
-#[cfg(test)]
 use hash_circuit::hash::Hashable;
 
 /// Lookup  represent the poseidon table in zkevm circuit
@@ -20,42 +19,6 @@ pub trait PoseidonLookup {
     }
 }
 
-impl<F: FieldExt> ConstraintBuilder<F> {
-    pub fn poseidon_lookup(
-        &mut self,
-        name: &'static str,
-        [left, right, domain, hash]: [Query<F>; 4],
-        poseidon: &impl PoseidonLookup,
-    ) {
-        let extended_queries = [
-            Query::one(),
-            hash,
-            left,
-            right,
-            Query::zero(),
-            domain,
-            Query::one(),
-        ];
-
-        let (q_enable, [hash, left, right, control, domain_spec, head_mark]) =
-            poseidon.lookup_columns();
-
-        self.add_lookup(
-            name,
-            extended_queries,
-            [
-                q_enable.current(),
-                hash.current(),
-                left.current(),
-                right.current(),
-                control.current(),
-                domain_spec.current(),
-                head_mark.current(),
-            ],
-        )
-    }
-}
-
 #[cfg(test)]
 #[derive(Clone, Copy)]
 pub struct PoseidonTable {

src/tests.rs

@@ -11,7 +11,7 @@ use halo2_proofs::{
 };
 use mpt_zktrie::state::{builder::HASH_SCHEME_DONE, witness::WitnessGenerator, ZktrieState};
 
-const N_ROWS: usize = 1024;
+const N_ROWS: usize = 8 * 256 + 1;
 const STORAGE_ADDRESS: Address = Address::repeat_byte(1);
 
 fn initial_generator() -> WitnessGenerator {