Fix lookup serialization and update with latest serialization changes

plonky2/src/plonk/circuit_builder.rs

@@ -66,7 +66,7 @@ pub enum LookupChallenges {
 }
 
 /// Structure containing, for each lookup table, the indices of the last lookup row, the last lookup table row and the first lookup table row. Since the rows are in reverse order in the trace, they actually correspond, respectively, to: the indices of the first `LookupGate`, the first `LookupTableGate` and the last `LookupTableGate`.
-#[derive(Eq, PartialEq, Clone, Debug)]
+#[derive(Clone, Debug, Eq, PartialEq)]
 pub struct LookupWire {
     /// Index of the last lookup row (i.e. the first `LookupGate`).
     pub last_lu_gate: usize,

plonky2/src/util/serialization/mod.rs

@@ -5,6 +5,7 @@ pub mod generator_serialization;
 pub mod gate_serialization;
 
 use alloc::collections::BTreeMap;
+use alloc::sync::Arc;
 use alloc::vec;
 use alloc::vec::Vec;
 use core::convert::Infallible;
@@ -38,6 +39,7 @@ use crate::iop::ext_target::ExtensionTarget;
 use crate::iop::generator::WitnessGeneratorRef;
 use crate::iop::target::{BoolTarget, Target};
 use crate::iop::wire::Wire;
+use crate::plonk::circuit_builder::LookupWire;
 use crate::plonk::circuit_data::{
     CircuitConfig, CircuitData, CommonCircuitData, ProverCircuitData, ProverOnlyCircuitData,
     VerifierCircuitData, VerifierCircuitTarget, VerifierOnlyCircuitData,
@@ -112,6 +114,14 @@ pub trait Read {
         Ok(buf[0])
     }
 
+    /// Reads a `u16` value from `self`.
+    #[inline]
+    fn read_u16(&mut self) -> IoResult<u16> {
+        let mut buf = [0; size_of::<u16>()];
+        self.read_exact(&mut buf)?;
+        Ok(u16::from_le_bytes(buf))
+    }
+
     /// Reads a `u32` value from `self`.
     #[inline]
     fn read_u32(&mut self) -> IoResult<u32> {
@@ -334,8 +344,8 @@ pub trait Read {
         let wires = self.read_field_ext_vec::<F, D>(config.num_wires)?;
         let plonk_zs = self.read_field_ext_vec::<F, D>(config.num_challenges)?;
         let plonk_zs_next = self.read_field_ext_vec::<F, D>(config.num_challenges)?;
-        let lookup_zs = self.read_field_ext_vec::<F, D>(common_data.num_lookup_polys)?;
-        let lookup_zs_next = self.read_field_ext_vec::<F, D>(common_data.num_lookup_polys)?;
+        let lookup_zs = self.read_field_ext_vec::<F, D>(common_data.num_all_lookup_polys())?;
+        let lookup_zs_next = self.read_field_ext_vec::<F, D>(common_data.num_all_lookup_polys())?;
         let partial_products = self
             .read_field_ext_vec::<F, D>(common_data.num_partial_products * config.num_challenges)?;
         let quotient_polys = self.read_field_ext_vec::<F, D>(
@@ -362,6 +372,8 @@ pub trait Read {
         let wires = self.read_target_ext_vec::<D>()?;
         let plonk_zs = self.read_target_ext_vec::<D>()?;
         let plonk_zs_next = self.read_target_ext_vec::<D>()?;
+        let lookup_zs = self.read_target_ext_vec::<D>()?;
+        let next_lookup_zs = self.read_target_ext_vec::<D>()?;
         let partial_products = self.read_target_ext_vec::<D>()?;
         let quotient_polys = self.read_target_ext_vec::<D>()?;
 
@@ -371,6 +383,8 @@ pub trait Read {
             wires,
             plonk_zs,
             plonk_zs_next,
+            lookup_zs,
+            next_lookup_zs,
             partial_products,
             quotient_polys,
         })
@@ -746,6 +760,15 @@ pub trait Read {
 
         let num_partial_products = self.read_usize()?;
 
+        let num_lookup_polys = self.read_usize()?;
+        let num_lookup_selectors = self.read_usize()?;
+        let length = self.read_usize()?;
+        let mut luts = Vec::with_capacity(length);
+
+        for _ in 0..length {
+            luts.push(Arc::new(self.read_lut()?));
+        }
+
         Ok(CommonCircuitData {
             config,
             fri_params,
@@ -757,6 +780,9 @@ pub trait Read {
             num_public_inputs,
             k_is,
             num_partial_products,
+            num_lookup_polys,
+            num_lookup_selectors,
+            luts,
         })
     }
 
@@ -831,6 +857,22 @@ pub trait Read {
 
         let circuit_digest = self.read_hash::<F, <C as GenericConfig<D>>::Hasher>()?;
 
+        let length = self.read_usize()?;
+        let mut lookup_rows = Vec::with_capacity(length);
+        for _ in 0..length {
+            lookup_rows.push(LookupWire {
+                last_lu_gate: self.read_usize()?,
+                last_lut_gate: self.read_usize()?,
+                first_lut_gate: self.read_usize()?,
+            });
+        }
+
+        let length = self.read_usize()?;
+        let mut lut_to_lookups = Vec::with_capacity(length);
+        for _ in 0..length {
+            lut_to_lookups.push(self.read_target_lut()?);
+        }
+
         Ok(ProverOnlyCircuitData {
             generators,
             generator_indices_by_watches,
@@ -841,6 +883,8 @@ pub trait Read {
             representative_map,
             fft_root_table,
             circuit_digest,
+            lookup_rows,
+            lut_to_lookups,
         })
     }
 
@@ -1095,6 +1139,30 @@ pub trait Read {
             public_inputs,
         })
     }
+
+    /// Reads a lookup table stored as `Vec<(u16, u16)>` from `self`.
+    #[inline]
+    fn read_lut(&mut self) -> IoResult<Vec<(u16, u16)>> {
+        let length = self.read_usize()?;
+        let mut lut = Vec::with_capacity(length);
+        for _ in 0..length {
+            lut.push((self.read_u16()?, self.read_u16()?));
+        }
+
+        Ok(lut)
+    }
+
+    /// Reads a target lookup table stored as `Vec<(Target, Target)>` from `self`.
+    #[inline]
+    fn read_target_lut(&mut self) -> IoResult<Vec<(Target, Target)>> {
+        let length = self.read_usize()?;
+        let mut lut = Vec::with_capacity(length);
+        for _ in 0..length {
+            lut.push((self.read_target()?, self.read_target()?));
+        }
+
+        Ok(lut)
+    }
 }
 
 /// Writing
@@ -1134,6 +1202,12 @@ pub trait Write {
         self.write_all(&[x])
     }
 
+    /// Writes a word `x` to `self`.
+    #[inline]
+    fn write_u16(&mut self, x: u16) -> IoResult<()> {
+        self.write_all(&x.to_le_bytes())
+    }
+
     /// Writes a word `x` to `self.`
     #[inline]
     fn write_u32(&mut self, x: u32) -> IoResult<()> {
@@ -1357,6 +1431,8 @@ pub trait Write {
         self.write_target_ext_vec::<D>(&os.wires)?;
         self.write_target_ext_vec::<D>(&os.plonk_zs)?;
         self.write_target_ext_vec::<D>(&os.plonk_zs_next)?;
+        self.write_target_ext_vec::<D>(&os.lookup_zs)?;
+        self.write_target_ext_vec::<D>(&os.next_lookup_zs)?;
         self.write_target_ext_vec::<D>(&os.partial_products)?;
         self.write_target_ext_vec::<D>(&os.quotient_polys)
     }
@@ -1672,6 +1748,9 @@ pub trait Write {
             num_public_inputs,
             k_is,
             num_partial_products,
+            num_lookup_polys,
+            num_lookup_selectors,
+            luts,
         } = common_data;
 
         self.write_circuit_config(config)?;
@@ -1693,6 +1772,13 @@ pub trait Write {
 
         self.write_usize(*num_partial_products)?;
 
+        self.write_usize(*num_lookup_polys)?;
+        self.write_usize(*num_lookup_selectors)?;
+        self.write_usize(luts.len())?;
+        for lut in luts.iter() {
+            self.write_lut(lut)?;
+        }
+
         Ok(())
     }
 
@@ -1730,6 +1816,8 @@ pub trait Write {
             representative_map,
             fft_root_table,
             circuit_digest,
+            lookup_rows,
+            lut_to_lookups,
         } = prover_only_circuit_data;
 
         self.write_usize(generators.len())?;
@@ -1768,6 +1856,18 @@ pub trait Write {
 
         self.write_hash::<F, <C as GenericConfig<D>>::Hasher>(*circuit_digest)?;
 
+        self.write_usize(lookup_rows.len())?;
+        for wire in lookup_rows.iter() {
+            self.write_usize(wire.last_lu_gate)?;
+            self.write_usize(wire.last_lut_gate)?;
+            self.write_usize(wire.first_lut_gate)?;
+        }
+
+        self.write_usize(lut_to_lookups.len())?;
+        for tlut in lut_to_lookups.iter() {
+            self.write_target_lut(tlut)?;
+        }
+
         Ok(())
     }
 
@@ -1970,6 +2070,30 @@ pub trait Write {
         self.write_compressed_proof(proof)?;
         self.write_field_vec(public_inputs)
     }
+
+    /// Writes a lookup table to `self`.
+    #[inline]
+    fn write_lut(&mut self, lut: &[(u16, u16)]) -> IoResult<()> {
+        self.write_usize(lut.len())?;
+        for (a, b) in lut.iter() {
+            self.write_u16(*a)?;
+            self.write_u16(*b)?;
+        }
+
+        Ok(())
+    }
+
+    /// Writes a target lookup table to `self`.
+    #[inline]
+    fn write_target_lut(&mut self, lut: &[(Target, Target)]) -> IoResult<()> {
+        self.write_usize(lut.len())?;
+        for (a, b) in lut.iter() {
+            self.write_target(*a)?;
+            self.write_target(*b)?;
+        }
+
+        Ok(())
+    }
 }
 
 impl Write for Vec<u8> {