Return chips refactors

src/air/builder.rs

@@ -136,14 +136,17 @@ impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> LookupBuilder fo
 pub struct Record {
     pub nonce: u32,
     pub count: u32,
+    // Original query that did the lookup. `None` is for the root lookup
+    pub query_index: Option<usize>,
 }
 
 impl Record {
     /// Updates the provide record and returns the require record
-    pub fn new_lookup(&mut self, nonce: u32) -> Record {
+    pub fn new_lookup(&mut self, nonce: u32, query_index: usize) -> Record {
         let require = *self;
         self.nonce = nonce;
         self.count += 1;
+        self.query_index = Some(query_index);
         require
     }
 

src/core/big_num.rs

@@ -45,12 +45,13 @@ impl<F: PrimeField32> Chipset<F> for BigNum {
         &self,
         input: &[F],
         nonce: u32,
+        query_index: usize,
         queries: &mut QueryRecord<F>,
         requires: &mut Vec<Record>,
     ) -> Vec<F> {
         let in1: [F; 8] = input[0..8].try_into().unwrap();
         let in2: [F; 8] = input[8..16].try_into().unwrap();
-        let bytes = &mut queries.bytes.context(nonce, requires);
+        let bytes = &mut queries.bytes.context(nonce, query_index, requires);
         match self {
             BigNum::LessThan => {
                 let mut witness = BigNumCompareWitness::<F>::default();

src/core/chipset.rs

@@ -123,15 +123,22 @@ impl Chipset<BabyBear> for LurkChip {
         &self,
         input: &[BabyBear],
         nonce: u32,
+        query_index: usize,
         queries: &mut QueryRecord<BabyBear>,
         requires: &mut Vec<Record>,
     ) -> Vec<BabyBear> {
         match self {
-            LurkChip::Hasher3(hasher) => hasher.execute(input, nonce, queries, requires),
-            LurkChip::Hasher4(hasher) => hasher.execute(input, nonce, queries, requires),
-            LurkChip::Hasher5(hasher) => hasher.execute(input, nonce, queries, requires),
-            LurkChip::U64(op) => op.execute(input, nonce, queries, requires),
-            LurkChip::BigNum(op) => op.execute(input, nonce, queries, requires),
+            LurkChip::Hasher3(hasher) => {
+                hasher.execute(input, nonce, query_index, queries, requires)
+            }
+            LurkChip::Hasher4(hasher) => {
+                hasher.execute(input, nonce, query_index, queries, requires)
+            }
+            LurkChip::Hasher5(hasher) => {
+                hasher.execute(input, nonce, query_index, queries, requires)
+            }
+            LurkChip::U64(op) => op.execute(input, nonce, query_index, queries, requires),
+            LurkChip::BigNum(op) => op.execute(input, nonce, query_index, queries, requires),
         }
     }
 

src/core/u64.rs

@@ -85,6 +85,7 @@ impl<F: PrimeField32> Chipset<F> for U64 {
         &self,
         input: &[F],
         nonce: u32,
+        query_index: usize,
         queries: &mut QueryRecord<F>,
         requires: &mut Vec<Record>,
     ) -> Vec<F> {
@@ -93,7 +94,7 @@ impl<F: PrimeField32> Chipset<F> for U64 {
             U64::IsZero => 0, // unused
             _ => into_u64(&input[8..16]),
         };
-        let bytes = &mut queries.bytes.context(nonce, requires);
+        let bytes = &mut queries.bytes.context(nonce, query_index, requires);
         match self {
             U64::Add => {
                 let mut witness = Sum64::<F>::default();

src/gadgets/bytes/record.rs

@@ -13,7 +13,7 @@ use std::collections::BTreeMap;
 #[derive(Clone, Debug, Default, Eq, PartialEq)]
 pub struct BytesRecord {
     // Maps input bytes to a list of records for each potential operation that could've been
-    records: BTreeMap<ByteInput, BytesInputRecord>,
+    pub(crate) records: BTreeMap<ByteInput, BytesInputRecord>,
 }
 
 /// Temporary wrapper around `BytesRecord` which implements `ByteRecord`.
@@ -30,6 +30,7 @@ pub struct ByteRecordWithContext<'a> {
     pub nonce: u32,
     pub requires: &'a mut Vec<Record>,
     pub record: &'a mut BytesRecord,
+    pub query_index: usize,
 }
 
 /// For a given input byte pair, this structure records the nonce and count of the accesses to
@@ -68,6 +69,18 @@ impl BytesInputRecord {
             self.or,
         ]
     }
+
+    /// Returns a mutable iterator of all the records
+    pub fn iter_mut_records(&mut self) -> impl IntoIterator<Item = &mut Record> {
+        [
+            &mut self.range_u8,
+            &mut self.range_u16,
+            &mut self.less_than,
+            &mut self.and,
+            &mut self.xor,
+            &mut self.or,
+        ]
+    }
 }
 
 impl BytesRecord {
@@ -76,12 +89,14 @@ impl BytesRecord {
     pub fn context<'a>(
         &'a mut self,
         nonce: u32,
+        query_index: usize,
         requires: &'a mut Vec<Record>,
     ) -> ByteRecordWithContext<'a> {
         ByteRecordWithContext {
             nonce,
             record: self,
             requires,
+            query_index,
         }
     }
 
@@ -112,46 +127,52 @@ impl BytesRecord {
 impl ByteRecord for ByteRecordWithContext<'_> {
     fn range_check_u8_pair(&mut self, i1: u8, i2: u8) {
         let input = ByteInput::from_u8_pair(i1, i2);
+        let index = self.query_index;
         let range_u8 = &mut self.record.get_mut(input).range_u8;
-        let require = range_u8.new_lookup(self.nonce);
+        let require = range_u8.new_lookup(self.nonce, index);
         self.requires.push(require);
     }
 
     fn range_check_u16(&mut self, i: u16) {
         let input = ByteInput::from_u16(i);
+        let index = self.query_index;
         let range_u16 = &mut self.record.get_mut(input).range_u16;
-        let require = range_u16.new_lookup(self.nonce);
+        let require = range_u16.new_lookup(self.nonce, index);
         self.requires.push(require);
     }
 
     fn less_than(&mut self, i1: u8, i2: u8) -> bool {
         let input = ByteInput::from_u8_pair(i1, i2);
+        let index = self.query_index;
         let less_than = &mut self.record.get_mut(input).less_than;
-        let require = less_than.new_lookup(self.nonce);
+        let require = less_than.new_lookup(self.nonce, index);
         self.requires.push(require);
         input.less_than()
     }
 
     fn and(&mut self, i1: u8, i2: u8) -> u8 {
         let input = ByteInput::from_u8_pair(i1, i2);
+        let index = self.query_index;
         let and = &mut self.record.get_mut(input).and;
-        let require = and.new_lookup(self.nonce);
+        let require = and.new_lookup(self.nonce, index);
         self.requires.push(require);
         input.and()
     }
 
     fn xor(&mut self, i1: u8, i2: u8) -> u8 {
         let input = ByteInput::from_u8_pair(i1, i2);
+        let index = self.query_index;
         let xor = &mut self.record.get_mut(input).xor;
-        let require = xor.new_lookup(self.nonce);
+        let require = xor.new_lookup(self.nonce, index);
         self.requires.push(require);
         input.xor()
     }
 
     fn or(&mut self, i1: u8, i2: u8) -> u8 {
         let input = ByteInput::from_u8_pair(i1, i2);
+        let index = self.query_index;
         let or = &mut self.record.get_mut(input).or;
-        let require = or.new_lookup(self.nonce);
+        let require = or.new_lookup(self.nonce, index);
         self.requires.push(require);
         input.or()
     }

src/lair/air.rs

@@ -509,20 +509,16 @@ impl<F: Field> Ctrl<F> {
         <AB as AirBuilder>::Var: Debug,
     {
         match self {
-            Ctrl::Choose(_, cases, branches) => {
+            Ctrl::Choose(_, _, branches) => {
                 let map_len = map.len();
                 let init_state = index.save();
 
-                let mut process = |block: &Block<F>| {
+                branches.iter().for_each(|block| {
                     let sel = block.return_sel::<AB>(local);
                     block.eval(builder, local, &sel, index, map, toplevel, depth);
                     map.truncate(map_len);
                     index.restore(init_state);
-                };
-                branches.iter().for_each(&mut process);
-                if let Some(block) = cases.default.as_ref() {
-                    process(block)
-                };
+                });
             }
             Ctrl::ChooseMany(_, cases) => {
                 let map_len = map.len();

src/lair/bytecode.rs

@@ -71,10 +71,10 @@ pub struct Block<F> {
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum Ctrl<F> {
     /// `Choose(x, cases)` non-deterministically chooses which case to execute based on a value `x`
-    /// The third item is the list of unique branches, needed for `eval`
-    Choose(usize, Cases<F, F>, List<Block<F>>), // TODO use Arc or indices so that blocks are not duplicated
+    /// It works by mapping field elements into indices of the list of unique branches.
+    Choose(usize, Cases<F, usize>, List<Block<F>>),
     /// `ChooseMany(x, cases)` non-deterministically chooses which case to execute based on an array `x`
-    ChooseMany(List<usize>, Cases<List<F>, F>),
+    ChooseMany(List<usize>, Cases<List<F>, Block<F>>),
     /// Contains the variables whose bindings will construct the output of the
     /// block. The first `usize` is an unique identifier, representing the
     /// selector used for arithmetization
@@ -85,21 +85,21 @@ pub enum Ctrl<F> {
 /// matches and an optional default case in case there's no match. Each code path
 /// is encoded as its own block
 #[derive(Clone, Debug, Eq, PartialEq)]
-pub struct Cases<K, F> {
-    pub(crate) branches: Map<K, Block<F>>,
-    pub(crate) default: Option<Box<Block<F>>>,
+pub struct Cases<K, B> {
+    pub(crate) branches: Map<K, B>,
+    pub(crate) default: Option<Box<B>>,
 }
 
-impl<K: Ord, F> Cases<K, F> {
+impl<K: Ord, B> Cases<K, B> {
     /// Returns the block mapped from key `f`
     #[inline]
-    pub fn match_case(&self, k: &K) -> Option<&Block<F>> {
+    pub fn match_case(&self, k: &K) -> Option<&B> {
         self.branches.get(k).or(self.default.as_deref())
     }
 
     /// Returns the block at a given index
     #[inline]
-    pub fn get_index(&self, idx: usize) -> Option<&Block<F>> {
+    pub fn get_index(&self, idx: usize) -> Option<&B> {
         self.branches
             .get_index(idx)
             .map(|(_, b)| b)
@@ -117,7 +117,7 @@ impl<K: Ord, F> Cases<K, F> {
     }
 }
 
-impl<K, F> Cases<K, F> {
+impl<K, B> Cases<K, B> {
     /// Returns the size of `branches`, assuming that's the index of the default
     /// block
     #[inline]

src/lair/chipset.rs

@@ -23,6 +23,7 @@ pub trait Chipset<F>: Sync {
         &self,
         input: &[F],
         _nonce: u32,
+        _query_index: usize,
         _queries: &mut QueryRecord<F>,
         _requires: &mut Vec<Record>,
     ) -> Vec<F>
@@ -65,15 +66,16 @@ impl<F, C1: Chipset<F>, C2: Chipset<F>> Chipset<F> for &Either<C1, C2> {
         &self,
         input: &[F],
         nonce: u32,
+        query_index: usize,
         queries: &mut QueryRecord<F>,
         requires: &mut Vec<Record>,
     ) -> Vec<F>
     where
         F: PrimeField32,
     {
         match self {
-            Either::Left(c) => c.execute(input, nonce, queries, requires),
-            Either::Right(c) => c.execute(input, nonce, queries, requires),
+            Either::Left(c) => c.execute(input, nonce, query_index, queries, requires),
+            Either::Right(c) => c.execute(input, nonce, query_index, queries, requires),
         }
     }