added serialization support for gates and gadgets

Cargo.toml

@@ -10,8 +10,8 @@ edition = "2021"
 anyhow = { version = "1.0.40", default-features = false }
 itertools = { version = "0.10.0", default-features = false }
 num = { version = "0.4", default-features = false }
-plonky2 = { version = "0.1.2", default-features = false }
+plonky2 = { git = "https://github.com/mir-protocol/plonky2.git", default-features = false }
 
 [dev-dependencies]
-plonky2 = { version = "0.1.2", default-features = false, features = ["gate_testing"] }
+plonky2 = { git = "https://github.com/mir-protocol/plonky2.git", features = ["gate_testing"] }
 rand = { version = "0.8.4", default-features = false, features = ["getrandom"] }

src/gadgets/arithmetic_u32.rs

@@ -1,5 +1,7 @@
+use alloc::string::{String, ToString};
 use alloc::vec;
 use alloc::vec::Vec;
+use plonky2::util::serialization::{IoResult, Buffer, Read, Write};
 use core::marker::PhantomData;
 
 use plonky2::field::extension::Extendable;
@@ -12,6 +14,7 @@ use plonky2::plonk::circuit_builder::CircuitBuilder;
 use crate::gates::add_many_u32::U32AddManyGate;
 use crate::gates::arithmetic_u32::U32ArithmeticGate;
 use crate::gates::subtraction_u32::U32SubtractionGate;
+use crate::serialization::{WriteU32, ReadU32};
 use crate::witness::GeneratedValuesU32;
 
 #[derive(Clone, Copy, Debug)]
@@ -241,6 +244,23 @@ struct SplitToU32Generator, const D: usize> {
 impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
     for SplitToU32Generator<F, D>
 {
+    fn id(&self) -> String {
+        "SplitToU32Generator".to_string()
+    }
+
+    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
+        dst.write_target(self.x)?;
+        dst.write_target_u32(self.low)?;
+        dst.write_target_u32(self.high)
+    }
+
+    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
+        let x = src.read_target()?;
+        let low = src.read_target_u32()?;
+        let high = src.read_target_u32()?;
+        Ok(Self { x, low, high, _phantom: PhantomData })
+    }
+
     fn dependencies(&self) -> Vec<Target> {
         vec![self.x]
     }
@@ -254,6 +274,7 @@ impl, const D: usize> SimpleGenerator
         out_buffer.set_u32_target(self.low, low);
         out_buffer.set_u32_target(self.high, high);
     }
+
 }
 
 #[cfg(test)]

src/gates/add_many_u32.rs

@@ -1,7 +1,7 @@
-use alloc::boxed::Box;
 use alloc::format;
-use alloc::string::String;
+use alloc::string::{String, ToString};
 use alloc::vec::Vec;
+use plonky2::util::serialization::{IoResult, Buffer, Write, Read};
 use core::marker::PhantomData;
 
 use itertools::unfold;
@@ -11,7 +11,7 @@ use plonky2::gates::gate::Gate;
 use plonky2::gates::util::StridedConstraintConsumer;
 use plonky2::hash::hash_types::RichField;
 use plonky2::iop::ext_target::ExtensionTarget;
-use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGenerator};
+use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGeneratorRef};
 use plonky2::iop::target::Target;
 use plonky2::iop::wire::Wire;
 use plonky2::iop::witness::{PartitionWitness, Witness, WitnessWrite};
@@ -91,6 +91,17 @@ impl, const D: usize> Gate for U32AddManyGate
         format!("{self:?}")
     }
 
+    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
+        dst.write_usize(self.num_addends);
+        dst.write_usize(self.num_ops)
+    }
+
+    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
+        let num_addends = src.read_usize()?;
+        let num_ops = src.read_usize()?;
+        Ok(Self { num_addends, num_ops, _phantom: PhantomData })
+    }
+
     fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<F::Extension> {
         let mut constraints = Vec::with_capacity(self.num_constraints());
         for i in 0..self.num_ops {
@@ -236,10 +247,10 @@ impl, const D: usize> Gate for U32AddManyGate
         constraints
     }
 
-    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<Box<dyn WitnessGenerator<F>>> {
+    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<WitnessGeneratorRef<F>> {
         (0..self.num_ops)
             .map(|i| {
-                let g: Box<dyn WitnessGenerator<F>> = Box::new(
+                WitnessGeneratorRef::new(
                     U32AddManyGenerator {
                         gate: *self,
                         row,
@@ -247,8 +258,7 @@ impl, const D: usize> Gate for U32AddManyGate
                         _phantom: PhantomData,
                     }
                     .adapter(),
-                );
-                g
+                )
             })
             .collect()
     }
@@ -281,6 +291,23 @@ struct U32AddManyGenerator, const D: usize> {
 impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
     for U32AddManyGenerator<F, D>
 {
+    fn id(&self) -> String {
+        "U32AddManyGenerator".to_string()
+    }
+
+    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
+        self.gate.serialize(dst)?;
+        dst.write_usize(self.row)?;
+        dst.write_usize(self.i)
+    }
+
+    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
+        let gate = U32AddManyGate::deserialize(src)?;
+        let row = src.read_usize()?;
+        let i = src.read_usize()?;
+        Ok(Self { gate, row, i, _phantom: PhantomData })
+    }
+
     fn dependencies(&self) -> Vec<Target> {
         let local_target = |column| Target::wire(self.row, column);
 

src/gates/arithmetic_u32.rs

@@ -1,7 +1,7 @@
-use alloc::boxed::Box;
-use alloc::string::String;
+use alloc::string::{String, ToString};
 use alloc::vec::Vec;
 use alloc::{format, vec};
+use plonky2::util::serialization::{Buffer, IoResult, Read, Write};
 use core::marker::PhantomData;
 
 use itertools::unfold;
@@ -13,7 +13,7 @@ use plonky2::gates::packed_util::PackedEvaluableBase;
 use plonky2::gates::util::StridedConstraintConsumer;
 use plonky2::hash::hash_types::RichField;
 use plonky2::iop::ext_target::ExtensionTarget;
-use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGenerator};
+use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGeneratorRef};
 use plonky2::iop::target::Target;
 use plonky2::iop::wire::Wire;
 use plonky2::iop::witness::{PartitionWitness, Witness, WitnessWrite};
@@ -93,6 +93,15 @@ impl, const D: usize> Gate for U32ArithmeticG
         format!("{self:?}")
     }
 
+    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
+        dst.write_usize(self.num_ops)
+    }
+
+    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
+        let num_ops = src.read_usize()?;
+        Ok(Self { num_ops, _phantom: PhantomData })
+    }
+
     fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<F::Extension> {
         let mut constraints = Vec::with_capacity(self.num_constraints());
         for i in 0..self.num_ops {
@@ -240,10 +249,10 @@ impl, const D: usize> Gate for U32ArithmeticG
         constraints
     }
 
-    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<Box<dyn WitnessGenerator<F>>> {
+    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<WitnessGeneratorRef<F>> {
         (0..self.num_ops)
             .map(|i| {
-                let g: Box<dyn WitnessGenerator<F>> = Box::new(
+                let g: WitnessGeneratorRef<F> = WitnessGeneratorRef::new(
                     U32ArithmeticGenerator {
                         gate: *self,
                         row,
@@ -350,6 +359,10 @@ struct U32ArithmeticGenerator, const D: usize> {
 impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
     for U32ArithmeticGenerator<F, D>
 {
+    fn id(&self) -> String {
+        "U32ArithmeticGenerator".to_string()
+    }
+
     fn dependencies(&self) -> Vec<Target> {
         let local_target = |column| Target::wire(self.row, column);
 
@@ -411,6 +424,19 @@ impl, const D: usize> SimpleGenerator
             out_buffer.set_wire(wire, output_limb);
         }
     }
+
+    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
+        self.gate.serialize(dst)?;
+        dst.write_usize(self.row)?;
+        dst.write_usize(self.i)
+    }
+
+    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
+        let gate = U32ArithmeticGate::deserialize(src)?;
+        let row = src.read_usize()?;
+        let i = src.read_usize()?;
+        Ok(Self { gate, row, i, _phantom: PhantomData })
+    }
 }
 
 #[cfg(test)]

src/gates/comparison.rs

@@ -1,7 +1,7 @@
-use alloc::boxed::Box;
-use alloc::string::String;
+use alloc::string::{String, ToString};
 use alloc::vec::Vec;
 use alloc::{format, vec};
+use plonky2::util::serialization::{Buffer, IoResult, Read, Write};
 use core::marker::PhantomData;
 
 use plonky2::field::extension::Extendable;
@@ -12,7 +12,7 @@ use plonky2::gates::packed_util::PackedEvaluableBase;
 use plonky2::gates::util::StridedConstraintConsumer;
 use plonky2::hash::hash_types::RichField;
 use plonky2::iop::ext_target::ExtensionTarget;
-use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGenerator};
+use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGeneratorRef};
 use plonky2::iop::target::Target;
 use plonky2::iop::wire::Wire;
 use plonky2::iop::witness::{PartitionWitness, Witness, WitnessWrite};
@@ -98,6 +98,18 @@ impl, const D: usize> Gate for ComparisonGate
         format!("{self:?}<D={D}>")
     }
 
+    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
+        dst.write_usize(self.num_bits)?;
+        dst.write_usize(self.num_chunks)?;
+        Ok(())
+    }
+
+    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
+        let num_bits = src.read_usize()?;
+        let num_chunks = src.read_usize()?;
+        Ok(Self { num_bits, num_chunks, _phantom: PhantomData })
+    }
+
     fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<F::Extension> {
         let mut constraints = Vec::with_capacity(self.num_constraints());
 
@@ -287,12 +299,12 @@ impl, const D: usize> Gate for ComparisonGate
         constraints
     }
 
-    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<Box<dyn WitnessGenerator<F>>> {
+    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<WitnessGeneratorRef<F>> {
         let gen = ComparisonGenerator::<F, D> {
             row,
             gate: self.clone(),
         };
-        vec![Box::new(gen.adapter())]
+        vec![WitnessGeneratorRef::new(gen.adapter())]
     }
 
     fn num_wires(&self) -> usize {
@@ -404,6 +416,10 @@ struct ComparisonGenerator, const D: usize> {
 impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
     for ComparisonGenerator<F, D>
 {
+    fn id(&self) -> String {
+        "ComparisonGenerator".to_string()
+    }
+
     fn dependencies(&self) -> Vec<Target> {
         let local_target = |column| Target::wire(self.row, column);
 
@@ -512,6 +528,18 @@ impl, const D: usize> SimpleGenerator
             );
         }
     }
+
+    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
+        dst.write_usize(self.row)?;
+        self.gate.serialize(dst)
+    }
+
+    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
+        let row = src.read_usize()?;
+        let gate = ComparisonGate::deserialize(src)?;
+        Ok(Self { row, gate })
+    }
+
 }
 
 #[cfg(test)]

src/gates/range_check_u32.rs

@@ -1,7 +1,7 @@
-use alloc::boxed::Box;
-use alloc::string::String;
+use alloc::string::{String, ToString};
 use alloc::vec::Vec;
 use alloc::{format, vec};
+use plonky2::util::serialization::{Buffer, IoResult, Read, Write};
 use core::marker::PhantomData;
 
 use plonky2::field::extension::Extendable;
@@ -10,7 +10,7 @@ use plonky2::gates::gate::Gate;
 use plonky2::gates::util::StridedConstraintConsumer;
 use plonky2::hash::hash_types::RichField;
 use plonky2::iop::ext_target::ExtensionTarget;
-use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGenerator};
+use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGeneratorRef};
 use plonky2::iop::target::Target;
 use plonky2::iop::witness::{PartitionWitness, Witness, WitnessWrite};
 use plonky2::plonk::circuit_builder::CircuitBuilder;
@@ -55,6 +55,15 @@ impl, const D: usize> Gate for U32RangeCheckG
         format!("{self:?}")
     }
 
+    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
+        dst.write_usize(self.num_input_limbs)
+    }
+
+    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
+        let num_input_limbs = src.read_usize()?;
+        Ok(Self { num_input_limbs, _phantom: PhantomData })
+    }
+
     fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<F::Extension> {
         let mut constraints = Vec::with_capacity(self.num_constraints());
 
@@ -138,9 +147,9 @@ impl, const D: usize> Gate for U32RangeCheckG
         constraints
     }
 
-    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<Box<dyn WitnessGenerator<F>>> {
+    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<WitnessGeneratorRef<F>> {
         let gen = U32RangeCheckGenerator { gate: *self, row };
-        vec![Box::new(gen.adapter())]
+        vec![WitnessGeneratorRef::new(gen.adapter())]
     }
 
     fn num_wires(&self) -> usize {
@@ -171,6 +180,10 @@ pub struct U32RangeCheckGenerator, const D: usize>
 impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
     for U32RangeCheckGenerator<F, D>
 {
+    fn id(&self) -> String {
+        "U32RangeCheckGenerator".to_string()
+    }
+
     fn dependencies(&self) -> Vec<Target> {
         let num_input_limbs = self.gate.num_input_limbs;
         (0..num_input_limbs)
@@ -201,6 +214,18 @@ impl, const D: usize> SimpleGenerator
             }
         }
     }
+
+    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
+        dst.write_usize(self.row)?;
+        self.gate.serialize(dst)
+    }
+
+    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
+        let row = src.read_usize()?;
+        let gate = U32RangeCheckGate::deserialize(src)?;
+        Ok(Self { row, gate })
+    }
+
 }
 
 #[cfg(test)]

src/gates/subtraction_u32.rs

@@ -1,7 +1,7 @@
-use alloc::boxed::Box;
-use alloc::string::String;
+use alloc::string::{String, ToString};
 use alloc::vec::Vec;
 use alloc::{format, vec};
+use plonky2::util::serialization::{Buffer, IoResult, Read, Write};
 use core::marker::PhantomData;
 
 use plonky2::field::extension::Extendable;
@@ -12,7 +12,7 @@ use plonky2::gates::packed_util::PackedEvaluableBase;
 use plonky2::gates::util::StridedConstraintConsumer;
 use plonky2::hash::hash_types::RichField;
 use plonky2::iop::ext_target::ExtensionTarget;
-use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGenerator};
+use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGeneratorRef};
 use plonky2::iop::target::Target;
 use plonky2::iop::wire::Wire;
 use plonky2::iop::witness::{PartitionWitness, Witness, WitnessWrite};
@@ -87,6 +87,15 @@ impl, const D: usize> Gate for U32Subtraction
         format!("{self:?}")
     }
 
+    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
+        dst.write_usize(self.num_ops)
+    }
+
+    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
+        let num_ops = src.read_usize()?;
+        Ok(Self { num_ops, _phantom: PhantomData })
+    }
+
     fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<F::Extension> {
         let mut constraints = Vec::with_capacity(self.num_constraints());
         for i in 0..self.num_ops {
@@ -186,10 +195,10 @@ impl, const D: usize> Gate for U32Subtraction
         constraints
     }
 
-    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<Box<dyn WitnessGenerator<F>>> {
+    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<WitnessGeneratorRef<F>> {
         (0..self.num_ops)
             .map(|i| {
-                let g: Box<dyn WitnessGenerator<F>> = Box::new(
+                let g = WitnessGeneratorRef::new(
                     U32SubtractionGenerator {
                         gate: *self,
                         row,
@@ -273,6 +282,10 @@ struct U32SubtractionGenerator, const D: usize> {
 impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
     for U32SubtractionGenerator<F, D>
 {
+    fn id(&self) -> String {
+        "U32SubtractionGenerator".to_string()
+    }
+
     fn dependencies(&self) -> Vec<Target> {
         let local_target = |column| Target::wire(self.row, column);
 
@@ -329,6 +342,20 @@ impl, const D: usize> SimpleGenerator
             out_buffer.set_wire(wire, output_limbs[j]);
         }
     }
+
+    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
+        self.gate.serialize(dst)?;
+        dst.write_usize(self.row)?;
+        dst.write_usize(self.i)
+    }
+
+    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
+        let gate = U32SubtractionGate::deserialize(src)?;
+        let row = src.read_usize()?;
+        let i = src.read_usize()?;
+        Ok(Self { gate, row, i, _phantom: PhantomData })
+    }
+
 }
 
 #[cfg(test)]

src/lib.rs

@@ -5,4 +5,5 @@ extern crate alloc;
 
 pub mod gadgets;
 pub mod gates;
+pub mod serialization;
 pub mod witness;

src/serialization.rs

@@ -0,0 +1,26 @@
+use alloc::vec::Vec;
+use plonky2::util::serialization::{Buffer, IoResult, Read, Write};
+
+use crate::gadgets::arithmetic_u32::U32Target;
+
+pub trait WriteU32 {
+    fn write_target_u32(&mut self, x: U32Target) -> IoResult<()>;
+}
+
+impl WriteU32 for Vec<u8> {
+    #[inline]
+    fn write_target_u32(&mut self, x: U32Target) -> IoResult<()> {
+        self.write_target(x.0)
+    }
+}
+
+pub trait ReadU32 {
+    fn read_target_u32(&mut self) -> IoResult<U32Target>;
+}
+
+impl ReadU32 for Buffer {
+    #[inline]
+    fn read_target_u32(&mut self) -> IoResult<U32Target> {
+        Ok(U32Target(self.read_target()?))
+    }
+}