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
This commit is contained in:
Ahmad Afuni 2025-04-08 02:15:46 +10:00 committed by GitHub
parent adad695ba5
commit 6528914366
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 502 additions and 48 deletions

View file

@ -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<OperationVerifyTarget> {
let _true = builder._true();
let _false = builder._false();
@ -54,6 +63,12 @@ impl OperationVerifyGadget {
.collect::<Vec<_>>()
};
// 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<F, D>,
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::<Vec<_>>();
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<F, D>,
@ -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::<Result<_>>()?;
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<StatementTarget>,
operations: Vec<OperationTarget>,
merkle_proofs: Vec<MerkleClaimAndProofTarget>,
op_verifications: Vec<OperationVerifyTarget>,
}
@ -379,6 +489,7 @@ pub struct MainPodVerifyInput {
pub signed_pods: Vec<SignedPod>,
pub statements: Vec<mainpod::Statement>,
pub operations: Vec<mainpod::Operation>,
pub merkle_proofs: Vec<MerkleClaimAndProof>,
}
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<mainpod::Statement>,
merkle_proofs: Vec<mainpod::MerkleClaimAndProof>,
) -> 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::<F, D>::new(config);
@ -448,6 +581,15 @@ mod tests {
let prev_statements_target: Vec<_> = (0..prev_statements.len())
.map(|_| builder.add_virtual_statement(&params))
.collect();
let merkle_proofs_target: Vec<_> = merkle_proofs
.iter()
.map(|_| mp_gadget.eval(&mut builder))
.collect::<Result<_>>()?;
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::<F>::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, &params, 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(
&params, &root, &key, None, &no_key_pf,
)?];
let prev_statements = vec![root_st, key_st];
operation_verify(st, op, prev_statements, merkle_proofs.clone())?;
Ok(())
}