From e3492b32881fdd82209b3d40619ce9e7af55a5d8 Mon Sep 17 00:00:00 2001
From: Yu Jiang Tham <yoojbruin@gmail.com>
Date: Thu, 14 Dec 2023 16:31:12 -0800
Subject: [PATCH] Update structs

---
 halo2-wasm/src/lib.rs      | 42 ++++++++++++++++++++++++++++----------
 halo2-wasm/src/metadata.rs | 11 +++++++---
 2 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/halo2-wasm/src/lib.rs b/halo2-wasm/src/lib.rs
index 8e2732f..3f11b1d 100644
--- a/halo2-wasm/src/lib.rs
+++ b/halo2-wasm/src/lib.rs
@@ -4,6 +4,7 @@ use std::cell::RefCell;
 use std::io::BufReader;
 use std::rc::Rc;
 
+use ethers_core::types::H256;
 use halo2_base::{
     gates::circuit::{builder::BaseCircuitBuilder, BaseCircuitParams},
     AssignedValue,
@@ -48,6 +49,8 @@ pub mod tests;
 use vkey::{write_onchain_vkey, OnchainVerifyingKey};
 use metadata::AxiomV2CircuitMetadata;
 
+use crate::metadata::decode_axiom_v2_circuit_metadata;
+
 #[wasm_bindgen]
 extern "C" {
     #[wasm_bindgen(js_namespace = console)]
@@ -83,20 +86,22 @@ pub struct CircuitConfig {
     num_virtual_instance: usize,
 }
 
-#[derive(Tsify, Serialize, Deserialize)]
+#[derive(Tsify, Serialize, Deserialize, Clone, Debug)]
 #[serde(rename_all = "camelCase")]
 #[tsify(into_wasm_abi, from_wasm_abi)]
 pub struct CircuitLimits {
-    max_user_subqueries: usize,
-    max_user_outputs: usize,
+    user_max_outputs: usize,
+    user_max_subqueries: usize,
 }
 
-#[derive(Tsify, Serialize, Deserialize)]
-#[serde(rename_all = "camelCase")]
-#[tsify(into_wasm_abi, from_wasm_abi)]
-pub struct CircuitComposition {
-    circuit_config: CircuitConfig,
-    circuit_limits: CircuitLimits,
+impl Default for CircuitLimits {
+    fn default() -> Self {
+        CircuitLimits {
+            // Note: ideally should grab from axiom-tools
+            user_max_outputs: 128,
+            user_max_subqueries: 128,
+        }
+    }
 }
 
 #[wasm_bindgen]
@@ -107,6 +112,8 @@ pub struct Halo2Wasm {
     pub public: Vec<Vec<AssignedValue<Fr>>>,
     #[wasm_bindgen(skip)]
     pub circuit_params: Option<BaseCircuitParams>,
+    #[wasm_bindgen(skip)]
+    pub circuit_limits: Option<CircuitLimits>,
     params: Option<ParamsKZG<Bn256>>,
     pk: Option<ProvingKey<G1Affine>>,
     vk: Option<VerifyingKey<G1Affine>>,
@@ -119,6 +126,7 @@ impl Default for Halo2Wasm {
             circuit: Rc::new(RefCell::new(circuit)),
             public: vec![],
             circuit_params: None,
+            circuit_limits: Some(CircuitLimits::default()),
             params: None,
             pk: None,
             vk: None,
@@ -249,10 +257,19 @@ impl Halo2Wasm {
 
     #[wasm_bindgen(js_name = getEncodedCircuitMetadata)]
     pub fn get_encoded_circuit_metadata(&mut self) -> Vec<u8> {
-        let circuit_metadata = AxiomV2CircuitMetadata::from_circuit_params(&self.circuit_params.clone().unwrap());
+        let circuit_metadata = AxiomV2CircuitMetadata::from_circuit_params(
+            &self.circuit_params.clone().unwrap(),
+            &self.circuit_limits.clone().unwrap()
+        );
         circuit_metadata.encode().unwrap().0.to_vec()
     }
 
+    #[wasm_bindgen(js_name = decodeCircuitMetadata)]
+    pub fn decode_circuit_metadata(bytes: &[u8]) -> AxiomV2CircuitMetadata {
+        let encoded = H256::from_slice(bytes);
+        decode_axiom_v2_circuit_metadata(encoded).unwrap()
+    }
+
     #[wasm_bindgen(js_name = getVk)]
     pub fn get_vk(&self) -> Vec<u8> {
         let file = self
@@ -266,7 +283,10 @@ impl Halo2Wasm {
     #[wasm_bindgen(js_name = getOnchainVk)]
     pub fn get_onchain_vk(&self) -> Vec<u8> {
         let vk = self.vk.as_ref().unwrap();
-        let circuit_metadata = AxiomV2CircuitMetadata::from_circuit_params(&self.circuit_params.clone().unwrap());
+        let circuit_metadata = AxiomV2CircuitMetadata::from_circuit_params(
+            &self.circuit_params.clone().unwrap(),
+            &self.circuit_limits.clone().unwrap()
+        );
         let preprocessed = vk
             .fixed_commitments()
             .iter()
diff --git a/halo2-wasm/src/metadata.rs b/halo2-wasm/src/metadata.rs
index 816fb9a..b37ec67 100644
--- a/halo2-wasm/src/metadata.rs
+++ b/halo2-wasm/src/metadata.rs
@@ -3,14 +3,18 @@ use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 use ethers_core::types::H256;
 use halo2_base::gates::circuit::BaseCircuitParams;
 use serde::{Deserialize, Serialize};
+use tsify::Tsify;
+
+use crate::CircuitLimits;
 
 /// All configuration parameters of an Axiom Client Circuit that are
 /// hard-coded into the Verify Compute Circuit (which is an Aggregation Circuit with Universality::Full).
 ///
 /// This metadata is only for a circuit built using `RlcCircuitBuilder`
 /// or `BaseCircuitBuilder`, where the circuit _may_ be an aggregation circuit.
-#[derive(Clone, Debug, Default, Serialize, Deserialize, Hash)]
+#[derive(Tsify, Clone, Debug, Default, Serialize, Deserialize, Hash)]
 #[serde(rename_all = "camelCase")]
+#[tsify(into_wasm_abi, from_wasm_abi)]
 pub struct AxiomV2CircuitMetadata {
     /// Version byte for domain separation on version of Axiom client, halo2-lib, snark-verifier (for example if we switch to mv_lookup)
     /// If `version = x`, this should be thought of as Axiom Query v2.x
@@ -104,7 +108,8 @@ impl AxiomV2CircuitMetadata {
     }
 
     pub fn from_circuit_params(
-        circuit_params: &BaseCircuitParams
+        circuit_params: &BaseCircuitParams,
+        circuit_limits: &CircuitLimits
     ) -> Self {
         Self {
             version: 0,
@@ -114,7 +119,7 @@ impl AxiomV2CircuitMetadata {
             num_advice_per_phase: circuit_params.num_advice_per_phase.clone().iter().map(|x| *x as u16).collect(),
             num_lookup_advice_per_phase: circuit_params.num_lookup_advice_per_phase.clone().iter().map(|x| *x as u8).collect(),
             num_fixed: circuit_params.num_fixed.clone() as u8,
-            max_outputs: 1,
+            max_outputs: circuit_limits.user_max_outputs as u16,
             ..Default::default()
         }
     }