github/zerokit
1
1
Fork 0
mirror of https://github.com/vacp2p/zerokit.git synced 2026-04-04 03:00:25 -04:00
Code Issues Packages Projects Releases Wiki Activity

feat(rln): improve graph initialization by deserialize only once (#368)

Browse Source
This commit is contained in:
Vinh Trịnh
2026-01-21 11:23:04 +07:00
committed by GitHub
parent 832bdfe8d9
commit c35e62a635
8 changed files with 107 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
rln/src/circuit/error.rs
View File

@@ -7,11 +7,20 @@ pub enum ZKeyReadError {
SerializationError(#[from] ark_serialize::SerializationError),
}
/// Errors that can occur during witness calculation graph reading operations
#[derive(Debug, thiserror::Error)]
pub enum GraphReadError {
#[error("Empty graph bytes provided")]
EmptyBytes,
#[error("Failed to deserialize witness calculation graph: {0}")]
GraphDeserialization(#[from] std::io::Error),
#[error("Tree depth mismatch: circuit expects depth {expected}, but {actual} was provided")]
TreeDepthMismatch { expected: usize, actual: usize },
}
/// Errors that can occur during witness calculation
#[derive(Debug, thiserror::Error)]
pub enum WitnessCalcError {
#[error("Failed to deserialize witness calculation graph: {0}")]
GraphDeserialization(#[from] std::io::Error),
#[error("Failed to evaluate witness calculation graph: {0}")]
GraphEvaluation(String),
#[error("Invalid input length for '{name}': expected {expected}, got {actual}")]

16
rln/src/circuit/iden3calc.rs
View File

@@ -1,15 +1,14 @@
// This crate is based on the code by iden3. Its preimage can be found here:
// https://github.com/iden3/circom-witnesscalc/blob/5cb365b6e4d9052ecc69d4567fcf5bc061c20e94/src/lib.rs
mod graph;
pub(crate) mod graph;
mod proto;
mod storage;
pub(crate) mod storage;
use std::collections::HashMap;
use graph::Node;
use ruint::aliases::U256;
use storage::deserialize_witnesscalc_graph;
use zeroize::zeroize_flat_type;
use self::graph::fr_to_u256;
@@ -20,7 +19,7 @@ pub(crate) type InputSignalsInfo = HashMap<String, (usize, usize)>;
pub(crate) fn calc_witness<I: IntoIterator<Item = (String, Vec<FrOrSecret>)>>(
inputs: I,
graph_data: &[u8],
graph: &super::Graph,
) -> Result<Vec<Fr>, WitnessCalcError> {
let mut inputs: HashMap<String, Vec<U256>> = inputs
.into_iter()
@@ -38,12 +37,9 @@ pub(crate) fn calc_witness<I: IntoIterator<Item = (String, Vec<FrOrSecret>)>>(
})
.collect();
let (nodes, signals, input_mapping): (Vec<Node>, Vec<usize>, InputSignalsInfo) =
deserialize_witnesscalc_graph(std::io::Cursor::new(graph_data))?;
let mut inputs_buffer = get_inputs_buffer(get_inputs_size(&graph.nodes));
let mut inputs_buffer = get_inputs_buffer(get_inputs_size(&nodes));
populate_inputs(&inputs, &input_mapping, &mut inputs_buffer)?;
populate_inputs(&inputs, &graph.input_mapping, &mut inputs_buffer)?;
if let Some(v) = inputs.get_mut("identitySecret") {
// DO NOT USE: unsafe { zeroize_flat_type(v) } only clears the Vec pointer, not the data—can cause memory leaks
@@ -53,7 +49,7 @@ pub(crate) fn calc_witness<I: IntoIterator<Item = (String, Vec<FrOrSecret>)>>(
}
}
let res = graph::evaluate(&nodes, inputs_buffer.as_slice(), &signals)
let res = graph::evaluate(&graph.nodes, inputs_buffer.as_slice(), &graph.signals)
.map_err(WitnessCalcError::GraphEvaluation)?;
for val in inputs_buffer.iter_mut() {

59
rln/src/circuit/mod.rs
View File

@@ -18,7 +18,11 @@ use ark_groth16::{
use ark_relations::r1cs::ConstraintMatrices;
use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
use self::error::ZKeyReadError;
use self::{
error::{GraphReadError, ZKeyReadError},
iden3calc::InputSignalsInfo,
};
use crate::circuit::iden3calc::{graph::Node, storage::deserialize_witnesscalc_graph};
#[cfg(not(target_arch = "wasm32"))]
const GRAPH_BYTES: &[u8] = include_bytes!("../../resources/tree_depth_20/graph.bin");
@@ -31,6 +35,11 @@ static ARKZKEY: LazyLock<Zkey> = LazyLock::new(|| {
read_arkzkey_from_bytes_uncompressed(ARKZKEY_BYTES).expect("Default zkey must be valid")
});
#[cfg(not(target_arch = "wasm32"))]
static GRAPH: LazyLock<Graph> = LazyLock::new(|| {
graph_from_raw(GRAPH_BYTES, Some(DEFAULT_TREE_DEPTH)).expect("Default graph must be valid")
});
pub const DEFAULT_TREE_DEPTH: usize = 20;
pub const COMPRESS_PROOF_SIZE: usize = 128;
@@ -73,6 +82,17 @@ pub type Zkey = (ArkProvingKey<Curve>, ConstraintMatrices<Fr>);
/// Verifying key for the Groth16 proof system.
pub type VerifyingKey = ArkVerifyingKey<Curve>;
/// Parsed witness calculator graph.
///
/// Contains the deserialized computation graph used for witness calculation.
/// Parsing this once and reusing it avoids repeated deserialization overhead.
#[derive(Clone)]
pub struct Graph {
pub(crate) nodes: Vec<Node>,
pub(crate) signals: Vec<usize>,
pub(crate) input_mapping: InputSignalsInfo,
}
/// Loads the zkey from raw bytes
pub fn zkey_from_raw(zkey_data: &[u8]) -> Result<Zkey, ZKeyReadError> {
if zkey_data.is_empty() {
@@ -84,16 +104,47 @@ pub fn zkey_from_raw(zkey_data: &[u8]) -> Result<Zkey, ZKeyReadError> {
Ok(proving_key_and_matrices)
}
/// Parses the witness calculator graph from raw bytes
pub fn graph_from_raw(
graph_data: &[u8],
expected_tree_depth: Option<usize>,
) -> Result<Graph, GraphReadError> {
if graph_data.is_empty() {
return Err(GraphReadError::EmptyBytes);
}
let (nodes, signals, input_mapping) =
deserialize_witnesscalc_graph(std::io::Cursor::new(graph_data))
.map_err(GraphReadError::GraphDeserialization)?;
if let Some(expected) = expected_tree_depth {
let actual = input_mapping
.get("pathElements")
.map(|(_, len)| *len)
.unwrap_or(0);
if expected != actual {
return Err(GraphReadError::TreeDepthMismatch { expected, actual });
}
}
Ok(Graph {
nodes,
signals,
input_mapping,
})
}
// Loads default zkey from folder
#[cfg(not(target_arch = "wasm32"))]
pub fn zkey_from_folder() -> &'static Zkey {
&ARKZKEY
}
// Loads default graph from folder
// Loads default parsed graph from folder
#[cfg(not(target_arch = "wasm32"))]
pub fn graph_from_folder() -> &'static [u8] {
GRAPH_BYTES
pub fn graph_from_folder() -> &'static Graph {
&GRAPH
}
// The following functions and structs are based on code from ark-zkey:

4
rln/src/error.rs
View File

@@ -6,7 +6,7 @@ use thiserror::Error;
use zerokit_utils::error::{FromConfigError, HashError, ZerokitMerkleTreeError};
use crate::circuit::{
error::{WitnessCalcError, ZKeyReadError},
error::{GraphReadError, WitnessCalcError, ZKeyReadError},
Fr,
};
@@ -76,6 +76,8 @@ pub enum RLNError {
Hash(#[from] HashError),
#[error("ZKey error: {0}")]
ZKey(#[from] ZKeyReadError),
#[error("Graph error: {0}")]
Graph(#[from] GraphReadError),
#[error("Protocol error: {0}")]
Protocol(#[from] ProtocolError),
#[error("Verification error: {0}")]

4
rln/src/prelude.rs
View File

@@ -12,8 +12,8 @@ pub use crate::protocol::compute_tree_root;
pub use crate::protocol::{generate_zk_proof, verify_zk_proof};
pub use crate::{
circuit::{
zkey_from_raw, Curve, Fq, Fq2, Fr, G1Affine, G1Projective, G2Affine, G2Projective, Proof,
VerifyingKey, Zkey, COMPRESS_PROOF_SIZE, DEFAULT_TREE_DEPTH,
graph_from_raw, zkey_from_raw, Curve, Fq, Fq2, Fr, G1Affine, G1Projective, G2Affine,
G2Projective, Graph, Proof, VerifyingKey, Zkey, COMPRESS_PROOF_SIZE, DEFAULT_TREE_DEPTH,
},
error::{ProtocolError, RLNError, UtilsError, VerifyError},
hashers::{hash_to_field_be, hash_to_field_le, poseidon_hash, PoseidonHash},

6
rln/src/protocol/proof.rs
View File

@@ -8,7 +8,7 @@ use num_traits::Signed;
use super::witness::{inputs_for_witness_calculation, RLNWitnessInput};
use crate::{
circuit::{
iden3calc::calc_witness, qap::CircomReduction, Curve, Fr, Proof, VerifyingKey, Zkey,
iden3calc::calc_witness, qap::CircomReduction, Curve, Fr, Graph, Proof, VerifyingKey, Zkey,
COMPRESS_PROOF_SIZE,
},
error::ProtocolError,
@@ -292,13 +292,13 @@ pub fn generate_zk_proof_with_witness(
pub fn generate_zk_proof(
zkey: &Zkey,
witness: &RLNWitnessInput,
graph_data: &[u8],
graph: &Graph,
) -> Result<Proof, ProtocolError> {
let inputs = inputs_for_witness_calculation(witness)?
.into_iter()
.map(|(name, values)| (name.to_string(), values));
let full_assignment = calc_witness(inputs, graph_data)?;
let full_assignment = calc_witness(inputs, graph)?;
// Random Values
let mut rng = thread_rng();

23
rln/src/public.rs
View File

@@ -14,7 +14,7 @@ use {
#[cfg(not(target_arch = "wasm32"))]
use crate::{
circuit::{graph_from_folder, zkey_from_folder},
circuit::{graph_from_folder, graph_from_raw, zkey_from_folder, Graph},
protocol::generate_zk_proof,
};
use crate::{
@@ -59,7 +59,7 @@ impl TreeConfigInput for <PoseidonTree as ZerokitMerkleTree>::Config {
pub struct RLN {
pub(crate) zkey: Zkey,
#[cfg(not(target_arch = "wasm32"))]
pub(crate) graph_data: Vec<u8>,
pub(crate) graph: Graph,
#[cfg(not(feature = "stateless"))]
pub(crate) tree: PoseidonTree,
}
@@ -101,7 +101,7 @@ impl RLN {
#[cfg(all(not(target_arch = "wasm32"), not(feature = "stateless")))]
pub fn new<T: TreeConfigInput>(tree_depth: usize, tree_config: T) -> Result<RLN, RLNError> {
let zkey = zkey_from_folder().to_owned();
let graph_data = graph_from_folder().to_owned();
let graph = graph_from_folder().to_owned();
let config = tree_config.into_tree_config()?;
// We compute a default empty tree
@@ -113,7 +113,7 @@ impl RLN {
Ok(RLN {
zkey,
graph_data,
graph,
#[cfg(not(feature = "stateless"))]
tree,
})
@@ -129,9 +129,9 @@ impl RLN {
#[cfg(all(not(target_arch = "wasm32"), feature = "stateless"))]
pub fn new() -> Result<RLN, RLNError> {
let zkey = zkey_from_folder().to_owned();
let graph_data = graph_from_folder().to_owned();
let graph = graph_from_folder().clone();
Ok(RLN { zkey, graph_data })
Ok(RLN { zkey, graph })
}
/// Creates a new RLN object by passing circuit resources as byte vectors.
@@ -176,6 +176,8 @@ impl RLN {
tree_config: T,
) -> Result<RLN, RLNError> {
let zkey = zkey_from_raw(&zkey_data)?;
let graph = graph_from_raw(&graph_data, Some(tree_depth))?;
let config = tree_config.into_tree_config()?;
// We compute a default empty tree
@@ -187,7 +189,7 @@ impl RLN {
Ok(RLN {
zkey,
graph_data,
graph,
#[cfg(not(feature = "stateless"))]
tree,
})
@@ -221,8 +223,9 @@ impl RLN {
#[cfg(all(not(target_arch = "wasm32"), feature = "stateless"))]
pub fn new_with_params(zkey_data: Vec<u8>, graph_data: Vec<u8>) -> Result<RLN, RLNError> {
let zkey = zkey_from_raw(&zkey_data)?;
let graph = graph_from_raw(&graph_data, None)?;
Ok(RLN { zkey, graph_data })
Ok(RLN { zkey, graph })
}
/// Creates a new stateless RLN object by passing circuit resources as a byte vector.
@@ -566,7 +569,7 @@ impl RLN {
/// ```
#[cfg(not(target_arch = "wasm32"))]
pub fn generate_zk_proof(&self, witness: &RLNWitnessInput) -> Result<Proof, RLNError> {
let proof = generate_zk_proof(&self.zkey, witness, &self.graph_data)?;
let proof = generate_zk_proof(&self.zkey, witness, &self.graph)?;
Ok(proof)
}
@@ -585,7 +588,7 @@ impl RLN {
witness: &RLNWitnessInput,
) -> Result<(Proof, RLNProofValues), RLNError> {
let proof_values = proof_values_from_witness(witness)?;
let proof = generate_zk_proof(&self.zkey, witness, &self.graph_data)?;
let proof = generate_zk_proof(&self.zkey, witness, &self.graph)?;
Ok((proof, proof_values))
}

27
rln/tests/public.rs
View File

@@ -179,17 +179,23 @@ mod test {
assert!(verified);
}
#[test]
fn test_initialization_with_params() {
let zkey_data = include_bytes!("../resources/tree_depth_20/rln_final.arkzkey").to_vec();
let graph_data = include_bytes!("../resources/tree_depth_20/graph.bin").to_vec();
#[cfg(all(not(target_arch = "wasm32"), not(feature = "stateless")))]
assert!(RLN::new_with_params(DEFAULT_TREE_DEPTH, zkey_data, graph_data, "").is_ok());
#[cfg(all(not(target_arch = "wasm32"), feature = "stateless"))]
assert!(RLN::new_with_params(zkey_data, graph_data).is_ok());
}
#[cfg(not(feature = "stateless"))]
mod tree_test {
use ark_std::{rand::thread_rng, UniformRand};
use rand::{rngs::ThreadRng, Rng};
use rln::{
circuit::{Fq, Fr, Proof, DEFAULT_TREE_DEPTH},
hashers::{hash_to_field_le, poseidon_hash},
pm_tree_adapter::PmtreeConfig,
protocol::*,
public::RLN,
};
use rln::prelude::*;
use serde_json::json;
const NO_OF_LEAVES: usize = 256;
@@ -1126,12 +1132,7 @@ mod test {
mod stateless_test {
use ark_std::{rand::thread_rng, UniformRand};
use rand::Rng;
use rln::{
circuit::Fr,
hashers::{hash_to_field_le, poseidon_hash, PoseidonHash},
protocol::*,
public::RLN,
};
use rln::prelude::*;
use zerokit_utils::merkle_tree::{
OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree,
};