From 6528914366db21351b5a25b159869625767064a7 Mon Sep 17 00:00:00 2001 From: Ahmad Afuni Date: Tue, 8 Apr 2025 02:15:46 +1000 Subject: [PATCH] chore(backend): implement more circuit op logic (#173) * Add backend MerkleProof type * Add eval_not_contains * Remove print statement * Handle some edge cases * Add test * Add missing ? * Optimisation and stylistic changes * Code review --- src/backends/plonky2/circuits/common.rs | 59 ++++- src/backends/plonky2/circuits/mainpod.rs | 226 +++++++++++++++++- src/backends/plonky2/mainpod.rs | 28 ++- src/backends/plonky2/mock/mainpod/mod.rs | 75 ++++-- .../plonky2/mock/mainpod/operation.rs | 128 +++++++++- src/backends/plonky2/primitives/merkletree.rs | 14 +- .../plonky2/primitives/merkletree_circuit.rs | 11 +- src/middleware/mod.rs | 5 + src/middleware/operation.rs | 3 + src/middleware/statement.rs | 1 + 10 files changed, 502 insertions(+), 48 deletions(-) diff --git a/src/backends/plonky2/circuits/common.rs b/src/backends/plonky2/circuits/common.rs index 0510258..1f4cfac 100644 --- a/src/backends/plonky2/circuits/common.rs +++ b/src/backends/plonky2/circuits/common.rs @@ -3,14 +3,16 @@ use crate::backends::plonky2::basetypes::D; use crate::backends::plonky2::mock::mainpod::Statement; use crate::backends::plonky2::mock::mainpod::{Operation, OperationArg}; +use crate::backends::plonky2::primitives::merkletree::MerkleClaimAndProofTarget; use crate::middleware::{ NativeOperation, NativePredicate, Params, Predicate, StatementArg, ToFields, Value, - EMPTY_VALUE, F, HASH_SIZE, OPERATION_ARG_F_LEN, STATEMENT_ARG_F_LEN, VALUE_SIZE, + EMPTY_VALUE, F, HASH_SIZE, OPERATION_ARG_F_LEN, OPERATION_AUX_F_LEN, STATEMENT_ARG_F_LEN, + VALUE_SIZE, }; use anyhow::Result; use plonky2::field::extension::Extendable; use plonky2::field::types::{Field, PrimeField64}; -use plonky2::hash::hash_types::RichField; +use plonky2::hash::hash_types::{HashOutTarget, RichField, NUM_HASH_OUT_ELTS}; use plonky2::iop::target::{BoolTarget, Target}; use plonky2::iop::witness::{PartialWitness, WitnessWrite}; use plonky2::plonk::circuit_builder::CircuitBuilder; @@ -155,6 +157,7 @@ impl StatementTarget { pub struct OperationTarget { pub op_type: [Target; Params::operation_type_size()], pub args: Vec<[Target; OPERATION_ARG_F_LEN]>, + pub aux: [Target; OPERATION_AUX_F_LEN], } impl OperationTarget { @@ -174,6 +177,7 @@ impl OperationTarget { { pw.set_target_arr(&self.args[i], &arg.to_fields(params))?; } + pw.set_target_arr(&self.aux, &op.aux().to_fields(params))?; Ok(()) } @@ -197,6 +201,56 @@ pub trait Flattenable { fn from_flattened(vs: &[Target]) -> Self; } +/// For the purpose of op verification, we need only look up the +/// Merkle claim rather than the Merkle proof since it is verified +/// elsewhere. +pub struct MerkleClaimTarget { + pub(crate) enabled: BoolTarget, + pub(crate) root: HashOutTarget, + pub(crate) key: ValueTarget, + pub(crate) value: ValueTarget, + pub(crate) existence: BoolTarget, +} + +impl From for MerkleClaimTarget { + fn from(pf: MerkleClaimAndProofTarget) -> Self { + Self { + enabled: pf.enabled, + root: pf.root, + key: pf.key, + value: pf.value, + existence: pf.existence, + } + } +} + +impl Flattenable for MerkleClaimTarget { + fn flatten(&self) -> Vec { + [ + vec![self.enabled.target], + self.root.elements.to_vec(), + self.key.elements.to_vec(), + self.value.elements.to_vec(), + vec![self.existence.target], + ] + .concat() + } + + fn from_flattened(vs: &[Target]) -> Self { + Self { + enabled: BoolTarget::new_unsafe(vs[0]), + root: HashOutTarget::from_vec((&vs[1..1 + NUM_HASH_OUT_ELTS]).to_vec()), + key: ValueTarget::from_slice( + &vs[1 + NUM_HASH_OUT_ELTS..1 + NUM_HASH_OUT_ELTS + VALUE_SIZE], + ), + value: ValueTarget::from_slice( + &vs[1 + NUM_HASH_OUT_ELTS + VALUE_SIZE..1 + NUM_HASH_OUT_ELTS + 2 * VALUE_SIZE], + ), + existence: BoolTarget::new_unsafe(vs[1 + NUM_HASH_OUT_ELTS + 2 * VALUE_SIZE]), + } + } +} + impl Flattenable for StatementTarget { fn flatten(&self) -> Vec { self.predicate @@ -290,6 +344,7 @@ impl CircuitBuilderPod for CircuitBuilder { args: (0..params.max_operation_args) .map(|_| self.add_virtual_target_arr()) .collect(), + aux: self.add_virtual_target_arr(), } } diff --git a/src/backends/plonky2/circuits/mainpod.rs b/src/backends/plonky2/circuits/mainpod.rs index 2a4ec14..4be750b 100644 --- a/src/backends/plonky2/circuits/mainpod.rs +++ b/src/backends/plonky2/circuits/mainpod.rs @@ -1,4 +1,4 @@ -use anyhow::Result; +use anyhow::{anyhow, Result}; use itertools::zip_eq; use plonky2::{ hash::{hash_types::HashOutTarget, poseidon::PoseidonHash}, @@ -6,19 +6,27 @@ use plonky2::{ plonk::circuit_builder::CircuitBuilder, }; -use crate::backends::plonky2::basetypes::{Value, D, EMPTY_HASH, F, VALUE_SIZE}; -use crate::backends::plonky2::circuits::common::{ - CircuitBuilderPod, OperationTarget, StatementTarget, ValueTarget, -}; use crate::backends::plonky2::mock::mainpod; use crate::backends::plonky2::signedpod::SignedPod; +use crate::backends::plonky2::{ + basetypes::{Value, D, EMPTY_HASH, F, VALUE_SIZE}, + mock::mainpod::MerkleClaimAndProof, + primitives::merkletree::{MerkleClaimAndProofTarget, MerkleProofGadget}, +}; use crate::middleware::{ hash_str, AnchoredKey, NativeOperation, NativePredicate, Params, PodType, Statement, StatementArg, ToFields, KEY_TYPE, SELF, }; +use crate::{ + backends::plonky2::{ + circuits::common::{CircuitBuilderPod, OperationTarget, StatementTarget, ValueTarget}, + primitives::merkletree, + }, + middleware, +}; use super::{ - common::Flattenable, + common::{Flattenable, MerkleClaimTarget}, signedpod::{SignedPodVerifyGadget, SignedPodVerifyTarget}, }; @@ -37,6 +45,7 @@ impl OperationVerifyGadget { st: &StatementTarget, op: &OperationTarget, prev_statements: &[StatementTarget], + merkle_claims: &[MerkleClaimTarget], ) -> Result { let _true = builder._true(); let _false = builder._false(); @@ -54,6 +63,12 @@ impl OperationVerifyGadget { .collect::>() }; + // Certain operations (Contains/NotContains) will refer to one + // of the provided Merkle proofs (if any). These proofs have already + // been verified, so we need only look up the claim. + let resolved_merkle_claim = + (merkle_claims.len() > 0).then(|| builder.vec_ref(merkle_claims, op.aux[0])); + // The verification may require aux data which needs to be stored in the // `OperationVerifyTarget` so that we can set during witness generation. @@ -77,6 +92,18 @@ impl OperationVerifyGadget { self.eval_lt_from_entries(builder, st, op, &resolved_op_args), ] }, + // Skip these if there are no resolved Merkle claims + if let Some(resolved_merkle_claim) = resolved_merkle_claim { + vec![self.eval_not_contains_from_entries( + builder, + st, + op, + resolved_merkle_claim, + &resolved_op_args, + )] + } else { + vec![] + }, ] .concat(); @@ -87,6 +114,74 @@ impl OperationVerifyGadget { Ok(OperationVerifyTarget {}) } + fn eval_not_contains_from_entries( + &self, + builder: &mut CircuitBuilder, + st: &StatementTarget, + op: &OperationTarget, + resolved_merkle_claim: MerkleClaimTarget, + resolved_op_args: &[StatementTarget], + ) -> BoolTarget { + let op_code_ok = op.has_native_type(builder, NativeOperation::NotContainsFromEntries); + + // Expect 2 op args of type `ValueOf`. + let op_arg_type_checks = resolved_op_args + .iter() + .take(2) + .map(|op_arg| op_arg.has_native_type(builder, &self.params, NativePredicate::ValueOf)) + .collect::>(); + let op_arg_types_ok = builder.all(op_arg_type_checks); + + // The values embedded in the op args must be values, i.e. the + // last `STATEMENT_ARG_F_LEN - VALUE_SIZE` slots of each being + // 0. + let merkle_root_arg = &resolved_op_args[0].args[1]; + let key_arg = &resolved_op_args[1].args[1]; + let op_arg_range_checks = [ + builder.statement_arg_is_value(merkle_root_arg), + builder.statement_arg_is_value(key_arg), + ]; + let op_arg_range_ok = builder.all(op_arg_range_checks); + + // Check Merkle proof (verified elsewhere) against op args. + let merkle_proof_checks = [ + /* The supplied Merkle proof must be enabled. */ + resolved_merkle_claim.enabled, + /* ...and it must be a nonexistence proof. */ + builder.not(resolved_merkle_claim.existence), + /* ...for the root-key pair in the resolved op args. */ + builder.is_equal_slice( + &merkle_root_arg.elements[..VALUE_SIZE], + &resolved_merkle_claim.root.elements, + ), + builder.is_equal_slice( + &key_arg.elements[..VALUE_SIZE], + &resolved_merkle_claim.key.elements, + ), + ]; + + let merkle_proof_ok = builder.all(merkle_proof_checks); + + // Check output statement + let arg1_key = resolved_op_args[0].args[0].clone(); + let arg2_key = resolved_op_args[1].args[0].clone(); + let expected_statement = StatementTarget::new_native( + builder, + &self.params, + NativePredicate::NotContains, + &[arg1_key, arg2_key], + ); + let st_ok = builder.is_equal_flattenable(st, &expected_statement); + + builder.all([ + op_code_ok, + op_arg_types_ok, + op_arg_range_ok, + merkle_proof_ok, + st_ok, + ]) + } + fn eval_eq_from_entries( &self, builder: &mut CircuitBuilder, @@ -315,6 +410,19 @@ impl MainPodVerifyGadget { let pub_statements = &input_statements[input_statements.len() - params.max_public_statements..]; + // Add Merkle claim/proof targets + let mp_gadget = MerkleProofGadget { + max_depth: params.max_depth_mt_gadget, + }; + let merkle_proofs: Vec<_> = (0..params.max_merkle_proofs) + .map(|_| mp_gadget.eval(builder)) + .collect::>()?; + let merkle_claims: Vec<_> = merkle_proofs + .clone() + .into_iter() + .map(|pf| pf.into()) + .collect(); + // 2. Calculate the Pod Id from the public statements let pub_statements_flattened = pub_statements .iter() @@ -350,7 +458,7 @@ impl MainPodVerifyGadget { let op_verification = OperationVerifyGadget { params: params.clone(), } - .eval(builder, st, op, prev_statements)?; + .eval(builder, st, op, prev_statements, &merkle_claims)?; op_verifications.push(op_verification); } @@ -360,6 +468,7 @@ impl MainPodVerifyGadget { signed_pods, statements: input_statements.to_vec(), operations, + merkle_proofs, op_verifications, }) } @@ -372,6 +481,7 @@ pub struct MainPodVerifyTarget { // The KEY_TYPE statement must be the first public one statements: Vec, operations: Vec, + merkle_proofs: Vec, op_verifications: Vec, } @@ -379,6 +489,7 @@ pub struct MainPodVerifyInput { pub signed_pods: Vec, pub statements: Vec, pub operations: Vec, + pub merkle_proofs: Vec, } impl MainPodVerifyTarget { @@ -402,6 +513,23 @@ impl MainPodVerifyTarget { self.statements[i].set_targets(pw, &self.params, st)?; self.operations[i].set_targets(pw, &self.params, op)?; } + assert_eq!(input.merkle_proofs.len(), self.params.max_merkle_proofs); + for (i, mp) in input.merkle_proofs.iter().enumerate() { + assert_eq!(mp.siblings.len(), self.params.max_depth_mt_gadget); + self.merkle_proofs[i].set_targets( + pw, + mp.enabled, + mp.existence, + mp.root, + mp.clone().try_into().unwrap_or(merkletree::MerkleProof { + existence: mp.existence, + siblings: mp.siblings.clone(), + other_leaf: None, + }), + mp.key, + mp.value, + )?; + } Ok(()) } } @@ -423,6 +551,7 @@ impl MainPodVerifyCircuit { #[cfg(test)] mod tests { + use merkletree::MerkleTree; use plonky2::plonk::{circuit_builder::CircuitBuilder, circuit_data::CircuitConfig}; use super::*; @@ -437,8 +566,12 @@ mod tests { st: mainpod::Statement, op: mainpod::Operation, prev_statements: Vec, + merkle_proofs: Vec, ) -> Result<()> { let params = Params::default(); + let mp_gadget = MerkleProofGadget { + max_depth: params.max_depth_mt_gadget, + }; let config = CircuitConfig::standard_recursion_config(); let mut builder = CircuitBuilder::::new(config); @@ -448,6 +581,15 @@ mod tests { let prev_statements_target: Vec<_> = (0..prev_statements.len()) .map(|_| builder.add_virtual_statement(¶ms)) .collect(); + let merkle_proofs_target: Vec<_> = merkle_proofs + .iter() + .map(|_| mp_gadget.eval(&mut builder)) + .collect::>()?; + let merkle_claims_target: Vec<_> = merkle_proofs_target + .clone() + .into_iter() + .map(|pf| pf.into()) + .collect(); let operation_verify = OperationVerifyGadget { params: params.clone(), @@ -457,6 +599,7 @@ mod tests { &st_target, &op_target, &prev_statements_target, + &merkle_claims_target, )?; let mut pw = PartialWitness::::new(); @@ -465,6 +608,19 @@ mod tests { for (prev_st_target, prev_st) in prev_statements_target.iter().zip(prev_statements.iter()) { prev_st_target.set_targets(&mut pw, ¶ms, prev_st)?; } + for (merkle_proof_target, merkle_proof) in + merkle_proofs_target.iter().zip(merkle_proofs.iter()) + { + merkle_proof_target.set_targets( + &mut pw, + merkle_proof.enabled, + merkle_proof.existence, + merkle_proof.root, + merkle_proof.clone().try_into()?, + merkle_proof.key, + merkle_proof.value, + )? + } let input = OperationVerifyInput {}; operation_verify.set_targets(&mut pw, &input)?; @@ -478,6 +634,8 @@ mod tests { #[test] fn test_operation_verify() -> Result<()> { + let params = Params::default(); + // None let st: mainpod::Statement = Statement::None.into(); let op = mainpod::Operation( @@ -486,7 +644,13 @@ mod tests { OperationAux::None, ); let prev_statements = vec![Statement::None.into()]; - operation_verify(st.clone(), op, prev_statements.clone())?; + let merkle_proofs = vec![]; + operation_verify( + st.clone(), + op, + prev_statements.clone(), + merkle_proofs.clone(), + )?; // NewEntry let st1: mainpod::Statement = @@ -502,7 +666,12 @@ mod tests { vec![], OperationAux::None, ); - operation_verify(st1.clone(), op, prev_statements.clone())?; + operation_verify( + st1.clone(), + op, + prev_statements.clone(), + merkle_proofs.clone(), + )?; // Copy let st: mainpod::Statement = Statement::None.into(); @@ -512,7 +681,7 @@ mod tests { OperationAux::None, ); let prev_statements = vec![Statement::None.into()]; - operation_verify(st, op, prev_statements)?; + operation_verify(st, op, prev_statements, merkle_proofs.clone())?; // Eq let st2: mainpod::Statement = Statement::ValueOf( @@ -531,7 +700,7 @@ mod tests { OperationAux::None, ); let prev_statements = vec![st1.clone(), st2]; - operation_verify(st, op, prev_statements)?; + operation_verify(st, op, prev_statements, merkle_proofs.clone())?; // Lt let st2: mainpod::Statement = Statement::ValueOf( @@ -550,7 +719,40 @@ mod tests { OperationAux::None, ); let prev_statements = vec![st1.clone(), st2]; - operation_verify(st, op, prev_statements)?; + operation_verify(st, op, prev_statements, merkle_proofs.clone())?; + + // NotContainsFromEntries + let kvs = [ + (1.into(), 55.into()), + (2.into(), 88.into()), + (175.into(), 0.into()), + ] + .into_iter() + .collect(); + let mt = MerkleTree::new(params.max_depth_mt_gadget, &kvs)?; + + let root = mt.root().into(); + let root_ak = AnchoredKey(PodId(Value::from(88).into()), "merkle root".into()); + + let key = 5.into(); + let key_ak = AnchoredKey(PodId(Value::from(88).into()), "key".into()); + + let no_key_pf = mt.prove_nonexistence(&key)?; + + let root_st: mainpod::Statement = Statement::ValueOf(root_ak, root).into(); + let key_st: mainpod::Statement = Statement::ValueOf(key_ak, key).into(); + let st: mainpod::Statement = Statement::NotContains(root_ak, key_ak).into(); + let op = mainpod::Operation( + OperationType::Native(NativeOperation::NotContainsFromEntries), + vec![OperationArg::Index(0), OperationArg::Index(1)], + OperationAux::MerkleProofIndex(0), + ); + + let merkle_proofs = vec![mainpod::MerkleClaimAndProof::try_from_middleware( + ¶ms, &root, &key, None, &no_key_pf, + )?]; + let prev_statements = vec![root_st, key_st]; + operation_verify(st, op, prev_statements, merkle_proofs.clone())?; Ok(()) } diff --git a/src/backends/plonky2/mainpod.rs b/src/backends/plonky2/mainpod.rs index dfaf89b..3cc0954 100644 --- a/src/backends/plonky2/mainpod.rs +++ b/src/backends/plonky2/mainpod.rs @@ -23,6 +23,8 @@ use crate::middleware::{ SELF, }; +use super::mock::mainpod::MerkleClaimAndProof; + pub struct Prover {} impl PodProver for Prover { @@ -47,12 +49,14 @@ impl PodProver for Prover { }) .collect_vec(); + let merkle_proofs = MockMainPod::extract_merkle_proofs(params, &inputs.operations)?; + // TODO: Move these methods from the mock main pod to a common place let statements = MockMainPod::layout_statements(params, &inputs); let operations = MockMainPod::process_private_statements_operations( params, &statements, - &[], // TODO: fill in the merkle proofs for Contains/NotContains ops + &merkle_proofs, inputs.operations, )?; let operations = @@ -67,6 +71,7 @@ impl PodProver for Prover { signed_pods: signed_pods_input, statements: statements[statements.len() - params.max_statements..].to_vec(), operations, + merkle_proofs, }; main_pod.set_targets(&mut pw, &input)?; @@ -173,19 +178,28 @@ pub mod tests { pay_stub: &frontend::SignedPod, sanction_list: &frontend::SignedPod, ) -> Result { + let sanction_set = match sanction_list.kvs.get("sanctionList") { + Some(frontend::Value::Set(s)) => Ok(s), + _ => Err(anyhow!("Missing sanction list!")), + }?; let now_minus_18y: i64 = 1169909388; let now_minus_1y: i64 = 1706367566; + let gov_id_kvs = gov_id.kvs(); + let id_number_value = gov_id_kvs.get(&"idNumber".into()).unwrap(); + let mut kyc = frontend::MainPodBuilder::new(params); kyc.add_signed_pod(gov_id); kyc.add_signed_pod(pay_stub); kyc.add_signed_pod(sanction_list); - // NOTE: Unimplemented in the circuit - // kyc.pub_op(op!( - // not_contains, - // (sanction_list, "sanctionList"), - // (gov_id, "idNumber") - // ))?; + kyc.pub_op(op!( + set_not_contains, + (sanction_list, "sanctionList"), + (gov_id, "idNumber"), + sanction_set + .middleware_set() + .prove_nonexistence(id_number_value)? + ))?; kyc.pub_op(op!(lt, (gov_id, "dateOfBirth"), now_minus_18y))?; kyc.pub_op(op!( eq, diff --git a/src/backends/plonky2/mock/mainpod/mod.rs b/src/backends/plonky2/mock/mainpod/mod.rs index 81dc15f..5a16e88 100644 --- a/src/backends/plonky2/mock/mainpod/mod.rs +++ b/src/backends/plonky2/mock/mainpod/mod.rs @@ -9,7 +9,7 @@ use std::any::Any; use std::fmt; use crate::{ - backends::plonky2::primitives::merkletree::MerkleProof, + backends::plonky2::primitives::merkletree, middleware::{ self, hash_str, AnchoredKey, Hash, MainPodInputs, NativeOperation, NativePredicate, NonePod, OperationType, Params, Pod, PodId, PodProver, PodType, Predicate, StatementArg, @@ -44,7 +44,7 @@ pub struct MockMainPod { statements: Vec, // All Merkle proofs // TODO: Use a backend-specific representation - merkle_proofs: Vec, + merkle_proofs: Vec, } impl fmt::Display for MockMainPod { @@ -227,6 +227,52 @@ impl MockMainPod { statements } + /// Extracts and pads Merkle proofs from Contains/NotContains ops. + pub(crate) fn extract_merkle_proofs( + params: &Params, + operations: &[middleware::Operation], + ) -> Result> { + let mut merkle_proofs = operations + .iter() + .flat_map(|op| match op { + middleware::Operation::ContainsFromEntries( + middleware::Statement::ValueOf(_, root), + middleware::Statement::ValueOf(_, key), + middleware::Statement::ValueOf(_, value), + pf, + ) => Some(MerkleClaimAndProof::try_from_middleware( + params, + root, + key, + Some(value), + pf, + )), + middleware::Operation::NotContainsFromEntries( + middleware::Statement::ValueOf(_, root), + middleware::Statement::ValueOf(_, key), + pf, + ) => Some(MerkleClaimAndProof::try_from_middleware( + params, root, key, None, pf, + )), + _ => None, + }) + .collect::>>()?; + if merkle_proofs.len() > params.max_merkle_proofs { + return Err(anyhow!( + "The number of required Merkle proofs ({}) exceeds the maximum number ({}).", + merkle_proofs.len(), + params.max_merkle_proofs + )); + } else { + fill_pad( + &mut merkle_proofs, + MerkleClaimAndProof::empty(params.max_depth_mt_gadget), + params.max_merkle_proofs, + ); + Ok(merkle_proofs) + } + } + fn find_op_arg( statements: &[Statement], op_arg: &middleware::Statement, @@ -248,7 +294,7 @@ impl MockMainPod { } fn find_op_aux( - merkle_proofs: &[MerkleProof], + merkle_proofs: &[MerkleClaimAndProof], op_aux: &middleware::OperationAux, ) -> Result { match op_aux { @@ -256,7 +302,14 @@ impl MockMainPod { middleware::OperationAux::MerkleProof(pf_arg) => merkle_proofs .iter() .enumerate() - .find_map(|(i, pf)| (pf == pf_arg).then_some(i)) + .find_map(|(i, pf)| { + pf.clone() + .try_into() + .ok() + .and_then(|mid_pf: merkletree::MerkleProof| { + (&mid_pf == pf_arg).then_some(i) + }) + }) .map(OperationAux::MerkleProofIndex) .ok_or(anyhow!( "Merkle proof corresponding to op arg {} not found", @@ -268,7 +321,7 @@ impl MockMainPod { pub(crate) fn process_private_statements_operations( params: &Params, statements: &[Statement], - merkle_proofs: &[MerkleProof], + merkle_proofs: &[MerkleClaimAndProof], input_operations: &[middleware::Operation], ) -> Result> { let mut operations = Vec::new(); @@ -336,15 +389,9 @@ impl MockMainPod { // TODO: Insert a new public statement of ValueOf with `key=KEY_TYPE, // value=PodType::MockMainPod` let statements = Self::layout_statements(params, &inputs); - let merkle_proofs = inputs - .operations - .iter() - .flat_map(|op| match op { - middleware::Operation::ContainsFromEntries(_, _, _, pf) => Some(pf.clone()), - middleware::Operation::NotContainsFromEntries(_, _, pf) => Some(pf.clone()), - _ => None, - }) - .collect::>(); + // Extract Merkle proofs and pad. + let merkle_proofs = Self::extract_merkle_proofs(params, &inputs.operations)?; + let operations = Self::process_private_statements_operations( params, &statements, diff --git a/src/backends/plonky2/mock/mainpod/operation.rs b/src/backends/plonky2/mock/mainpod/operation.rs index b7ad24d..206ffd2 100644 --- a/src/backends/plonky2/mock/mainpod/operation.rs +++ b/src/backends/plonky2/mock/mainpod/operation.rs @@ -1,12 +1,12 @@ use super::Statement; use crate::{ - backends::plonky2::primitives::merkletree::MerkleProof, - middleware::{self, OperationType, Params, ToFields, F}, + backends::plonky2::primitives::merkletree::{self, kv_hash}, + middleware::{self, Hash, OperationType, Params, ToFields, Value, EMPTY_HASH, EMPTY_VALUE, F}, }; use anyhow::{anyhow, Result}; use plonky2::field::types::{Field, PrimeField64}; use serde::{Deserialize, Serialize}; -use std::fmt; +use std::{fmt, iter}; #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub enum OperationArg { @@ -36,6 +36,118 @@ pub enum OperationAux { MerkleProofIndex(usize), } +impl ToFields for OperationAux { + fn to_fields(&self, _params: &Params) -> Vec { + let f = match self { + Self::None => F::ZERO, + Self::MerkleProofIndex(i) => F::from_canonical_usize(*i), + }; + vec![f] + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct MerkleClaimAndProof { + pub enabled: bool, + pub root: Hash, + pub key: Value, + pub value: Value, + pub existence: bool, + pub siblings: Vec, + pub case_ii_selector: bool, + pub other_key: Value, + pub other_value: Value, +} + +impl MerkleClaimAndProof { + pub fn empty(max_depth: usize) -> Self { + Self { + enabled: false, + root: EMPTY_HASH, + key: Value::from(1), + value: EMPTY_VALUE, + existence: false, + siblings: iter::repeat(EMPTY_HASH).take(max_depth).collect(), + case_ii_selector: false, + other_key: EMPTY_VALUE, + other_value: EMPTY_VALUE, + } + } + pub fn try_from_middleware( + params: &Params, + root: &Value, + key: &Value, + value: Option<&Value>, + mid_mp: &merkletree::MerkleProof, + ) -> Result { + if mid_mp.siblings.len() > params.max_depth_mt_gadget { + Err(anyhow!( + "Number of siblings ({}) exceeds maximum depth ({})", + mid_mp.siblings.len(), + params.max_depth_mt_gadget + )) + } else { + let (other_key, other_value) = mid_mp.other_leaf.unwrap_or((EMPTY_VALUE, EMPTY_VALUE)); + Ok(Self { + enabled: true, + root: root.clone().into(), + key: key.clone(), + value: value.cloned().unwrap_or(EMPTY_VALUE), + existence: mid_mp.existence, + siblings: mid_mp + .siblings + .iter() + .cloned() + .chain(iter::repeat(EMPTY_HASH)) + .take(params.max_depth_mt_gadget) + .collect(), + case_ii_selector: mid_mp.other_leaf.is_some(), + other_key, + other_value, + }) + } + } +} + +impl TryFrom for merkletree::MerkleProof { + type Error = anyhow::Error; + fn try_from(mp: MerkleClaimAndProof) -> Result { + if !mp.enabled { + return Err(anyhow!("Not a valid Merkle proof.")); + } + let existence = mp.existence; + let other_leaf = if mp.case_ii_selector { + Some((mp.other_key, mp.other_value)) + } else { + None + }; + // Trim padding (if any). + let siblings = mp + .siblings + .into_iter() + .rev() + .skip_while(|s| s == &EMPTY_HASH) + .collect::>() + .into_iter() + .rev() + .collect(); + Ok(merkletree::MerkleProof { + existence, + siblings, + other_leaf, + }) + } +} + +impl fmt::Display for MerkleClaimAndProof { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match merkletree::MerkleProof::try_from(self.clone()) { + Err(_) => write!(f, "∅"), + Ok(mp) => mp.fmt(f), + } + } +} + #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct Operation(pub OperationType, pub Vec, pub OperationAux); @@ -46,10 +158,13 @@ impl Operation { pub fn args(&self) -> &[OperationArg] { &self.1 } + pub fn aux(&self) -> &OperationAux { + &self.2 + } pub fn deref( &self, statements: &[Statement], - merkle_proofs: &[MerkleProof], + merkle_proofs: &[MerkleClaimAndProof], ) -> Result { let deref_args = self .1 @@ -65,7 +180,10 @@ impl Operation { .get(i) .cloned() .ok_or(anyhow!("Missing Merkle proof index {}", i)) - .map(crate::middleware::OperationAux::MerkleProof), + .and_then(|mp| { + mp.try_into() + .map(crate::middleware::OperationAux::MerkleProof) + }), }?; middleware::Operation::op(self.0.clone(), &deref_args, &deref_aux) } diff --git a/src/backends/plonky2/primitives/merkletree.rs b/src/backends/plonky2/primitives/merkletree.rs index 999cee1..c5497e1 100644 --- a/src/backends/plonky2/primitives/merkletree.rs +++ b/src/backends/plonky2/primitives/merkletree.rs @@ -23,13 +23,21 @@ pub struct MerkleTree { impl MerkleTree { /// builds a new `MerkleTree` where the leaves contain the given key-values pub fn new(max_depth: usize, kvs: &HashMap) -> Result { - let mut root = Node::Intermediate(Intermediate::empty()); + // Construct leaves. + let mut leaves: Vec<_> = kvs + .iter() + .map(|(k, v)| Leaf::new(max_depth, *k, *v)) + .collect::>()?; - for (k, v) in kvs.iter() { - let leaf = Leaf::new(max_depth, *k, *v)?; + // Start with a leaf or conclude with an empty node as root. + let mut root = leaves.pop().map(|l| Node::Leaf(l)).unwrap_or(Node::None); + + // Iterate over remaining leaves (if any) and add them. + for leaf in leaves.into_iter() { root.add_leaf(0, max_depth, leaf)?; } + // Fill in hashes. let _ = root.compute_hash(); Ok(Self { max_depth, root }) } diff --git a/src/backends/plonky2/primitives/merkletree_circuit.rs b/src/backends/plonky2/primitives/merkletree_circuit.rs index 3dbe9e2..d7e9762 100644 --- a/src/backends/plonky2/primitives/merkletree_circuit.rs +++ b/src/backends/plonky2/primitives/merkletree_circuit.rs @@ -35,8 +35,9 @@ pub struct MerkleProofGadget { pub max_depth: usize, } -pub struct MerkleProofTarget { - max_depth: usize, +#[derive(Clone)] +pub struct MerkleClaimAndProofTarget { + pub(crate) max_depth: usize, // `enabled` determines if the merkleproof verification is enabled pub(crate) enabled: BoolTarget, pub(crate) root: HashOutTarget, @@ -51,7 +52,7 @@ pub struct MerkleProofTarget { impl MerkleProofGadget { /// creates the targets and defines the logic of the circuit - pub fn eval(&self, builder: &mut CircuitBuilder) -> Result { + pub fn eval(&self, builder: &mut CircuitBuilder) -> Result { let enabled = builder.add_virtual_bool_target_safe(); let root = builder.add_virtual_hash(); let key = builder.add_virtual_value(); @@ -133,7 +134,7 @@ impl MerkleProofGadget { builder.connect(computed_root[j], expected_root[j]); } - Ok(MerkleProofTarget { + Ok(MerkleClaimAndProofTarget { max_depth: self.max_depth, enabled, existence, @@ -148,7 +149,7 @@ impl MerkleProofGadget { } } -impl MerkleProofTarget { +impl MerkleClaimAndProofTarget { /// assigns the given values to the targets pub fn set_targets( &self, diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs index 59189ad..5a4c0fb 100644 --- a/src/middleware/mod.rs +++ b/src/middleware/mod.rs @@ -20,6 +20,8 @@ use std::any::Any; use std::collections::HashMap; use std::fmt; +use crate::backends::plonky2::primitives::merkletree::MerkleProof; + pub const SELF: PodId = PodId(SELF_ID_HASH); impl fmt::Display for PodId { @@ -105,6 +107,8 @@ pub struct Params { // in a custom predicate pub max_custom_predicate_arity: usize, pub max_custom_batch_size: usize, + // maximum number of merkle proofs + pub max_merkle_proofs: usize, // maximum depth for merkle tree gadget pub max_depth_mt_gadget: usize, } @@ -121,6 +125,7 @@ impl Default for Params { max_operation_args: 5, max_custom_predicate_arity: 5, max_custom_batch_size: 5, + max_merkle_proofs: 5, max_depth_mt_gadget: 32, } } diff --git a/src/middleware/operation.rs b/src/middleware/operation.rs index 669de0d..b14a5a7 100644 --- a/src/middleware/operation.rs +++ b/src/middleware/operation.rs @@ -5,7 +5,10 @@ use serde::{Deserialize, Serialize}; use std::fmt; use std::iter; +use super::Hash; use super::{CustomPredicateRef, NativePredicate, Statement, StatementArg, ToFields, F}; +use crate::middleware::EMPTY_HASH; +use crate::middleware::EMPTY_VALUE; use crate::{ backends::plonky2::primitives::merkletree::{MerkleProof, MerkleTree}, middleware::{AnchoredKey, Params, Predicate, Value, SELF}, diff --git a/src/middleware/statement.rs b/src/middleware/statement.rs index e8eb2d8..d8379c9 100644 --- a/src/middleware/statement.rs +++ b/src/middleware/statement.rs @@ -15,6 +15,7 @@ pub const KEY_SIGNER: &str = "_signer"; pub const KEY_TYPE: &str = "_type"; pub const STATEMENT_ARG_F_LEN: usize = 8; pub const OPERATION_ARG_F_LEN: usize = 1; +pub const OPERATION_AUX_F_LEN: usize = 1; #[derive(Clone, Copy, Debug, FromRepr, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)] pub enum NativePredicate {