From 3cfee55f23989147dbb18ec7245b4a6b06cd319c Mon Sep 17 00:00:00 2001
From: parazyd <parazyd@dyne.org>
Date: Thu, 20 Jul 2023 13:44:10 +0200
Subject: [PATCH] zk: Update to halo2 with dynamic circuit config and implement
 ProvingKey serialization.

---
 Cargo.toml      |  4 +--
 src/zk/proof.rs | 66 +++++++++++++++++++++++++++++++++++++++++++------
 src/zk/vm.rs    |  5 ++--
 3 files changed, 64 insertions(+), 11 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index d50a0f048..aa01e31dd 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -111,8 +111,8 @@ simplelog = {version = "0.12.1", optional = true}
 rand = {version = "0.8.5", optional = true}
 blake3 = {version = "1.4.1", features = ["rayon"], optional = true}
 crypto_api_chachapoly = {version = "0.5.0", optional = true}
-halo2_proofs = {version = "0.3.0", optional = true}
-halo2_gadgets = {version = "0.3.0", optional = true}
+halo2_proofs = {version = "0.3.0", features = ["circuit-params"], optional = true}
+halo2_gadgets = {version = "0.3.0", features = ["circuit-params"], optional = true}
 
 # Smart contract runtime
 darkfi-sdk = {path = "src/sdk", optional = true}
diff --git a/src/zk/proof.rs b/src/zk/proof.rs
index cbebc550b..46aa7081a 100644
--- a/src/zk/proof.rs
+++ b/src/zk/proof.rs
@@ -19,6 +19,7 @@ use std::{io, io::Cursor};
 
 use darkfi_serial::{SerialDecodable, SerialEncodable};
 use halo2_proofs::{
+    helpers::SerdeFormat,
     pasta::{pallas, vesta},
     plonk,
     plonk::{Circuit, SingleVerifier},
@@ -41,12 +42,11 @@ impl VerifyingKey {
     }
 
     pub fn write<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
-        // FIXME: This can be optimized.
         let mut params = vec![];
         self.params.write(&mut params)?;
 
         let mut vk = vec![];
-        self.vk.write(&mut vk)?;
+        self.vk.write(&mut vk, SerdeFormat::RawBytes)?;
 
         let _ = writer.write(&(params.len() as u32).to_le_bytes())?;
         let _ = writer.write(&params)?;
@@ -58,11 +58,8 @@ impl VerifyingKey {
 
     pub fn read<R: io::Read, ConcreteCircuit: Circuit<pallas::Base>>(
         reader: &mut R,
+        circuit: ConcreteCircuit,
     ) -> io::Result<Self> {
-        // FIXME: This can be optimized
-        // FIXME: Don't assert
-
-        // FIXME: Make sure that the size is legitimate.
         // The format chosen in write():
         // [params.len()<u32>, params..., vk.len()<u32>, vk...]
 
@@ -89,7 +86,11 @@ impl VerifyingKey {
 
         let mut vk_c = Cursor::new(vk_buf);
         let vk: plonk::VerifyingKey<vesta::Affine> =
-            plonk::VerifyingKey::read::<Cursor<Vec<u8>>, ConcreteCircuit>(&mut vk_c, &params)?;
+            plonk::VerifyingKey::read::<Cursor<Vec<u8>>, ConcreteCircuit>(
+                &mut vk_c,
+                SerdeFormat::RawBytes,
+                circuit.params(),
+            )?;
 
         Ok(Self { params, vk })
     }
@@ -108,6 +109,57 @@ impl ProvingKey {
         let pk = plonk::keygen_pk(&params, vk, c).unwrap();
         ProvingKey { params, pk }
     }
+
+    pub fn write<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
+        let mut params = vec![];
+        self.params.write(&mut params)?;
+
+        let mut pk = vec![];
+        self.pk.write(&mut pk, SerdeFormat::RawBytes)?;
+
+        let _ = writer.write(&(params.len() as u32).to_le_bytes())?;
+        let _ = writer.write(&params)?;
+        let _ = writer.write(&(pk.len() as u32).to_le_bytes())?;
+        let _ = writer.write(&pk)?;
+
+        Ok(())
+    }
+
+    pub fn read<R: io::Read, ConcreteCircuit: Circuit<pallas::Base>>(
+        reader: &mut R,
+        circuit: ConcreteCircuit,
+    ) -> io::Result<Self> {
+        let mut params_len = [0u8; 4];
+        reader.read_exact(&mut params_len)?;
+        let params_len = u32::from_le_bytes(params_len) as usize;
+
+        let mut params_buf = vec![0u8; params_len];
+        reader.read_exact(&mut params_buf)?;
+
+        assert!(params_buf.len() == params_len);
+
+        let mut pk_len = [0u8; 4];
+        reader.read_exact(&mut pk_len)?;
+        let pk_len = u32::from_le_bytes(pk_len) as usize;
+
+        let mut pk_buf = vec![0u8; pk_len];
+        reader.read_exact(&mut pk_buf)?;
+
+        assert!(pk_buf.len() == pk_len);
+
+        let mut params_c = Cursor::new(params_buf);
+        let params: Params<vesta::Affine> = Params::read(&mut params_c)?;
+
+        let mut pk_c = Cursor::new(pk_buf);
+        let pk: plonk::ProvingKey<vesta::Affine> =
+            plonk::ProvingKey::read::<Cursor<Vec<u8>>, ConcreteCircuit>(
+                &mut pk_c,
+                SerdeFormat::RawBytes,
+                circuit.params(),
+            )?;
+
+        Ok(Self { params, pk })
+    }
 }
 
 #[derive(Clone, Default, PartialEq, Eq, SerialEncodable, SerialDecodable)]
diff --git a/src/zk/vm.rs b/src/zk/vm.rs
index 5bafa46d8..2df93f2bf 100644
--- a/src/zk/vm.rs
+++ b/src/zk/vm.rs
@@ -129,17 +129,18 @@ pub struct ZkCircuit {
 }
 
 impl ZkCircuit {
-    pub fn new(witnesses: Vec<Witness>, circuit_code: ZkBinary) -> Self {
+    pub fn new(witnesses: Vec<Witness>, circuit_code: &ZkBinary) -> Self {
         let constants = circuit_code.constants.iter().map(|x| x.1.clone()).collect();
         #[allow(clippy::map_clone)]
         let literals = circuit_code.literals.iter().map(|x| x.clone()).collect();
-        Self { constants, witnesses, literals, opcodes: circuit_code.opcodes }
+        Self { constants, witnesses, literals, opcodes: circuit_code.opcodes.clone() }
     }
 }
 
 impl Circuit<pallas::Base> for ZkCircuit {
     type Config = VmConfig;
     type FloorPlanner = floor_planner::V1;
+    type Params = ();
 
     fn without_witnesses(&self) -> Self {
         Self {