Fix custom predicate circuits and add tests for them (#235)

* add tests, fix custom predicates

* wip

* wip

* fix custom predicates

* modularize code

* fix typos

* remove scratch file

* update

* Update src/backends/plonky2/circuits/mainpod.rs

Co-authored-by: Ahmad Afuni <root@ahmadafuni.com>

---------

Co-authored-by: Ahmad Afuni <root@ahmadafuni.com>
This commit is contained in:
Eduard S. 2025-05-16 13:17:14 +02:00 committed by GitHub
parent f5a1aa7523
commit def0730462
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 629 additions and 153 deletions

View file

@ -87,7 +87,7 @@ pub(crate) fn extract_custom_predicate_verifications(
.find_map(|(i, cpb)| (cpb.id() == cpr.batch.id()).then_some(i))
.expect("find the custom predicate from the extracted unique list");
let custom_predicate_table_index =
batch_index * params.max_custom_predicate_batches + cpr.index;
batch_index * params.max_custom_batch_size + cpr.index;
CustomPredicateVerification {
custom_predicate_table_index,
custom_predicate: cpr.clone(),
@ -150,20 +150,18 @@ pub(crate) fn extract_merkle_proofs(
/// Find the operation argument statement in the list of previous statements and return the index.
fn find_op_arg(statements: &[Statement], op_arg: &middleware::Statement) -> Result<OperationArg> {
match op_arg {
middleware::Statement::None => Ok(OperationArg::None),
_ => statements
.iter()
.enumerate()
.find_map(|(i, s)| {
(&middleware::Statement::try_from(s.clone()).ok()? == op_arg).then_some(i)
})
.map(OperationArg::Index)
.ok_or(Error::custom(format!(
"Statement corresponding to op arg {} not found",
op_arg
))),
}
// NOTE: The `None` `Statement` always exists as a constant at index 0
statements
.iter()
.enumerate()
.find_map(|(i, s)| {
(&middleware::Statement::try_from(s.clone()).ok()? == op_arg).then_some(i)
})
.map(OperationArg::Index)
.ok_or(Error::custom(format!(
"Statement corresponding to op arg {} not found",
op_arg
)))
}
/// Find the operation auxiliary data in the list of auxiliary data and return the index.
@ -179,15 +177,21 @@ fn find_op_aux(
op: &middleware::Operation,
) -> Result<OperationAux> {
let op_aux = op.aux();
let op_type = op.op_type();
if let (OperationType::Custom(cpr), Some(cpvs)) = (op_type, custom_predicate_verifications) {
if let (middleware::Operation::Custom(cpr, op_args), Some(cpvs)) =
(op, custom_predicate_verifications)
{
return Ok(cpvs
.iter()
.enumerate()
.find_map(|(i, cpv)| {
(cpv.custom_predicate.batch.id() == cpr.batch.id()
&& cpv.custom_predicate.index == cpr.index)
.then_some(i)
&& cpv.custom_predicate.index == cpr.index
&& cpv
.op_args
.iter()
.zip_eq(op_args.iter())
.all(|(a0, a1)| a0.0 == a1.predicate() && a0.1 == a1.args()))
.then_some(i)
})
.map(OperationAux::CustomPredVerifyIndex)
.expect("custom predicate verification in the list"));
@ -228,6 +232,10 @@ fn pad_operation_args(params: &Params, args: &mut Vec<OperationArg>) {
pub(crate) fn layout_statements(params: &Params, inputs: &MainPodInputs) -> Vec<Statement> {
let mut statements = Vec::new();
// Statement at index 0 is always None to be used for padding operation arguments in custom
// predicate statements
statements.push(middleware::Statement::None.into());
// Input signed pods region
let none_sig_pod_box: Box<dyn Pod> = Box::new(NonePod {});
let none_sig_pod = none_sig_pod_box.as_ref();
@ -531,10 +539,16 @@ pub mod tests {
primitives::signature::SecretKey,
signedpod::Signer,
},
examples::{zu_kyc_pod_builder, zu_kyc_sign_pod_builders},
frontend::{self},
examples::{
eth_dos_pod_builder, eth_friend_signed_pod_builder, zu_kyc_pod_builder,
zu_kyc_sign_pod_builders,
},
frontend::{
key, literal, CustomPredicateBatchBuilder, MainPodBuilder, StatementTmplBuilder as STB,
{self},
},
middleware,
middleware::RawValue,
middleware::{CustomPredicateRef, NativePredicate as NP, RawValue},
op,
};
@ -644,4 +658,109 @@ pub mod tests {
let pod = (kyc_pod.pod as Box<dyn Any>).downcast::<MainPod>().unwrap();
pod.verify().unwrap()
}
#[test]
fn test_main_ethdos() -> frontend::Result<()> {
let params = Params {
max_input_signed_pods: 2,
max_input_main_pods: 1,
max_statements: 26,
max_public_statements: 5,
max_signed_pod_values: 8,
max_statement_args: 6,
max_operation_args: 4,
max_custom_predicate_arity: 4,
max_custom_batch_size: 3,
max_custom_predicate_wildcards: 12,
max_custom_predicate_verifications: 8,
..Default::default()
};
let mut alice = Signer(SecretKey(RawValue::from(1)));
let bob = Signer(SecretKey(RawValue::from(2)));
let mut charlie = Signer(SecretKey(RawValue::from(3)));
// Alice attests that she is ETH friends with Charlie and Charlie
// attests that he is ETH friends with Bob.
let alice_attestation =
eth_friend_signed_pod_builder(&params, charlie.public_key().0.into())
.sign(&mut alice)?;
let charlie_attestation =
eth_friend_signed_pod_builder(&params, bob.public_key().0.into()).sign(&mut charlie)?;
let alice_bob_ethdos_builder = eth_dos_pod_builder(
&params,
false,
&alice_attestation,
&charlie_attestation,
bob.public_key().0.into(),
)?;
let mut prover = MockProver {};
let pod = alice_bob_ethdos_builder.prove(&mut prover, &params)?;
assert!(pod.pod.verify().is_ok());
let mut prover = Prover {};
let alice_bob_ethdos = alice_bob_ethdos_builder.prove(&mut prover, &params)?;
let pod = (alice_bob_ethdos.pod as Box<dyn Any>)
.downcast::<MainPod>()
.unwrap();
Ok(pod.verify()?)
}
#[test]
fn test_main_mini_custom_1() -> frontend::Result<()> {
let params = Params {
max_input_signed_pods: 0,
max_input_main_pods: 0,
max_statements: 9,
max_public_statements: 4,
max_statement_args: 3,
max_operation_args: 3,
max_custom_predicate_arity: 3,
max_custom_batch_size: 3,
max_custom_predicate_wildcards: 4,
max_custom_predicate_verifications: 2,
..Default::default()
};
let mut cpb_builder = CustomPredicateBatchBuilder::new(params.clone(), "cpb".into());
let stb0 = STB::new(NP::ValueOf)
.arg(("id", key("score")))
.arg(literal(42));
let stb1 = STB::new(NP::Equal)
.arg(("id", "secret_key"))
.arg(("id", key("score")));
let _ = cpb_builder.predicate_and(
"pred_and",
&["id"],
&["secret_key"],
&[stb0.clone(), stb1.clone()],
)?;
let _ = cpb_builder.predicate_or("pred_or", &["id"], &["secret_key"], &[stb0, stb1])?;
let cpb = cpb_builder.finish();
let cpb_and = CustomPredicateRef::new(cpb.clone(), 0);
let _cpb_or = CustomPredicateRef::new(cpb.clone(), 1);
let mut pod_builder = MainPodBuilder::new(&params);
let st0 = pod_builder.priv_op(op!(new_entry, ("score", 42)))?;
let st1 = pod_builder.priv_op(op!(new_entry, ("foo", 42)))?;
let st2 = pod_builder.priv_op(op!(eq, st1.clone(), st0.clone()))?;
let _st3 = pod_builder.priv_op(op!(custom, cpb_and.clone(), st0, st2))?;
let mut prover = MockProver {};
let pod = pod_builder.prove(&mut prover, &params)?;
assert!(pod.pod.verify().is_ok());
let mut prover = Prover {};
let pod = pod_builder.prove(&mut prover, &params)?;
let pod = (pod.pod as Box<dyn Any>).downcast::<MainPod>().unwrap();
Ok(pod.verify()?)
}
}