From 8672419f2c488ecfd91c5405733a072a761fbffc Mon Sep 17 00:00:00 2001
From: parazyd <parazyd@dyne.org>
Date: Sat, 4 Mar 2023 22:20:10 +0100
Subject: [PATCH] Apply https://github.com/zcash/halo2/pull/661

Signed-off-by: parazyd <parazyd@dyne.org>
---
 halo2_proofs/src/plonk.rs             | 74 ++++++++++++++++++++++++++-
 halo2_proofs/src/plonk/keygen.rs      |  3 +-
 halo2_proofs/src/plonk/permutation.rs | 21 ++++++++
 3 files changed, 95 insertions(+), 3 deletions(-)

diff --git a/halo2_proofs/src/plonk.rs b/halo2_proofs/src/plonk.rs
index e96cf3fe97..50997299a4 100644
--- a/halo2_proofs/src/plonk.rs
+++ b/halo2_proofs/src/plonk.rs
@@ -9,9 +9,10 @@ use blake2b_simd::Params as Blake2bParams;
 use group::ff::{Field, FromUniformBytes, PrimeField};
 
 use crate::arithmetic::CurveAffine;
+use crate::helpers::CurveRead;
 use crate::poly::{
-    Coeff, EvaluationDomain, ExtendedLagrangeCoeff, LagrangeCoeff, PinnedEvaluationDomain,
-    Polynomial,
+    commitment::Params, Coeff, EvaluationDomain, ExtendedLagrangeCoeff, LagrangeCoeff,
+    PinnedEvaluationDomain, Polynomial,
 };
 use crate::transcript::{ChallengeScalar, EncodedChallenge, Transcript};
 
@@ -47,17 +48,85 @@ pub struct VerifyingKey<C: CurveAffine> {
     cs_degree: usize,
     /// The representative of this `VerifyingKey` in transcripts.
     transcript_repr: C::Scalar,
+    selectors: Vec<Vec<bool>>,
 }
 
 impl<C: CurveAffine> VerifyingKey<C>
 where
     C::Scalar: FromUniformBytes<64>,
 {
+    /// Writes a verifying key to a buffer.
+    pub fn write<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
+        writer.write_all(&(self.fixed_commitments.len() as u32).to_be_bytes())?;
+        for commitment in &self.fixed_commitments {
+            writer.write_all(commitment.to_bytes().as_ref())?;
+        }
+        self.permutation.write(writer)?;
+
+        // write self.selectors
+        for selector in &self.selectors {
+            let mut selector_bytes = vec![0u8; selector.len() / 8 + 1];
+            for (i, selector_idx) in selector.iter().enumerate() {
+                let byte_index = i / 8;
+                let bit_index = i % 8;
+                selector_bytes[byte_index] |= (*selector_idx as u8) << bit_index;
+            }
+            writer.write_all(&selector_bytes)?;
+        }
+
+        Ok(())
+    }
+
+    /// Reads a verification key from a buffer.
+    pub fn read<R: io::Read, ConcreteCircuit: Circuit<C::Scalar>>(
+        reader: &mut R,
+        params: &Params<C>,
+    ) -> io::Result<Self> {
+        let (domain, cs, _) = keygen::create_domain::<C, ConcreteCircuit>(params);
+        let mut num_fixed_columns_be_bytes = [0u8; 4];
+        reader.read_exact(&mut num_fixed_columns_be_bytes)?;
+        let num_fixed_columns = u32::from_be_bytes(num_fixed_columns_be_bytes);
+
+        let fixed_commitments: Vec<_> = (0..num_fixed_columns)
+            .map(|_| C::read(reader))
+            .collect::<Result<_, _>>()?;
+
+        let permutation = permutation::VerifyingKey::read(reader, &cs.permutation)?;
+
+        // read selectors
+        let selectors: Vec<Vec<bool>> = vec![vec![false; params.n as usize]; cs.num_selectors]
+            .into_iter()
+            .map(|mut selector| {
+                let mut selector_bytes = vec![0u8; selector.len() / 8 + 1];
+                reader
+                    .read_exact(&mut selector_bytes)
+                    .expect("unable to read selector bytes");
+                for (i, selector_idx) in selector.iter_mut().enumerate() {
+                    let byte_index = i / 8;
+                    let bit_index = i % 8;
+                    *selector_idx = (selector_bytes[byte_index] >> bit_index) & 1 == 1;
+                }
+                Ok(selector)
+            })
+            .collect::<Result<Vec<Vec<bool>>, &str>>()
+            .unwrap();
+        let (cs, _) = cs.compress_selectors(selectors.clone());
+
+        Ok(Self::from_parts(
+            domain,
+            fixed_commitments,
+            permutation,
+            cs,
+            selectors,
+        ))
+    }
+
     fn from_parts(
         domain: EvaluationDomain<C::Scalar>,
         fixed_commitments: Vec<C>,
         permutation: permutation::VerifyingKey<C>,
         cs: ConstraintSystem<C::Scalar>,
+        selectors: Vec<Vec<bool>>,
     ) -> Self {
         // Compute cached values.
         let cs_degree = cs.degree();
@@ -70,6 +139,7 @@ where
             cs_degree,
             // Temporary, this is not pinned.
             transcript_repr: C::Scalar::ZERO,
+            selectors,
         };
 
         let mut hasher = Blake2bParams::new()
diff --git a/halo2_proofs/src/plonk/keygen.rs b/halo2_proofs/src/plonk/keygen.rs
index bb26d5a6fb..1b2be6d0a0 100644
--- a/halo2_proofs/src/plonk/keygen.rs
+++ b/halo2_proofs/src/plonk/keygen.rs
@@ -219,7 +219,7 @@ where
     )?;
 
     let mut fixed = batch_invert_assigned(assembly.fixed);
-    let (cs, selector_polys) = cs.compress_selectors(assembly.selectors);
+    let (cs, selector_polys) = cs.compress_selectors(assembly.selectors.clone());
     fixed.extend(
         selector_polys
             .into_iter()
@@ -240,6 +240,7 @@ where
         fixed_commitments,
         permutation_vk,
         cs,
+        assembly.selectors,
     ))
 }
 
diff --git a/halo2_proofs/src/plonk/permutation.rs b/halo2_proofs/src/plonk/permutation.rs
index c81de27667..16f65f9da9 100644
--- a/halo2_proofs/src/plonk/permutation.rs
+++ b/halo2_proofs/src/plonk/permutation.rs
@@ -1,6 +1,7 @@
 use super::circuit::{Any, Column};
 use crate::{
     arithmetic::CurveAffine,
+    helpers::CurveRead,
     poly::{Coeff, ExtendedLagrangeCoeff, LagrangeCoeff, Polynomial},
 };
 
@@ -75,6 +76,26 @@ pub(crate) struct VerifyingKey<C: CurveAffine> {
     commitments: Vec<C>,
 }
 
+impl<C: CurveAffine> VerifyingKey<C> {
+    pub(crate) fn write<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
+        for commitment in &self.commitments {
+            writer.write_all(commitment.to_bytes().as_ref())?;
+        }
+
+        Ok(())
+    }
+
+    pub(crate) fn read<R: std::io::Read>(
+        reader: &mut R,
+        argument: &Argument,
+    ) -> std::io::Result<Self> {
+        let commitments = (0..argument.columns.len())
+            .map(|_| C::read(reader))
+            .collect::<Result<Vec<_>, _>>()?;
+        Ok(VerifyingKey { commitments })
+    }
+}
+
 /// The proving key for a single permutation argument.
 #[derive(Clone, Debug)]
 pub(crate) struct ProvingKey<C: CurveAffine> {