Constraints for custom predicates (#227)

* add target types for custom predicates

* simplify

* fix clippy

* fix typo

* don't use ref for NativePredicate

* fix wrong len

* precalculate CustomPredicateBatch id

* wip

* wip

* move code back

* great progress

* wip

* code complete, hopefully; missing tests

* fill aux for custom predicate op

* fix clippy warnings

* fix typos

* fix test import

* fix missing assignment in lt_mask, test custom_operation_verify_gadget

* fix mistake

* wip

* fix

* debug revert except for let entry = CustomPredicateVerifyEntryTarget

* fix batch_id calculation by fixing padding

* oops

* remove completed TODOs
This commit is contained in:
Eduard S. 2025-05-13 11:00:45 +02:00 committed by GitHub
parent 4fa9e20ecd
commit 024ed8bd04
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 1597 additions and 291 deletions

View file

@ -1,6 +1,6 @@
pub mod operation;
pub mod statement;
use std::any::Any;
use std::{any::Any, sync::Arc};
use itertools::Itertools;
pub use operation::*;
@ -17,14 +17,17 @@ pub use statement::*;
use crate::{
backends::plonky2::{
basetypes::{C, D},
circuits::mainpod::{MainPodVerifyCircuit, MainPodVerifyInput},
circuits::mainpod::{
CustomPredicateVerification, MainPodVerifyCircuit, MainPodVerifyInput,
},
error::{Error, Result},
primitives::merkletree::MerkleClaimAndProof,
signedpod::SignedPod,
},
middleware::{
self, AnchoredKey, DynError, Hash, MainPodInputs, NativeOperation, NonePod, OperationType,
Params, Pod, PodId, PodProver, PodType, StatementArg, ToFields, F, KEY_TYPE, SELF,
self, resolve_wildcard_values, AnchoredKey, CustomPredicateBatch, DynError, Hash,
MainPodInputs, NativeOperation, NonePod, OperationType, Params, Pod, PodId, PodProver,
PodType, StatementArg, ToFields, F, KEY_TYPE, SELF,
},
};
@ -37,7 +40,71 @@ pub(crate) fn hash_statements(statements: &[Statement], _params: &Params) -> mid
Hash(PoseidonHash::hash_no_pad(&field_elems).elements)
}
/// Extracts and pads Merkle proofs from Contains/NotContains ops.
/// Extracts unique `CustomPredicateBatch`es from Custom ops.
pub(crate) fn extract_custom_predicate_batches(
params: &Params,
operations: &[middleware::Operation],
) -> Result<Vec<Arc<CustomPredicateBatch>>> {
let custom_predicate_batches: Vec<_> = operations
.iter()
.flat_map(|op| match op {
middleware::Operation::Custom(cpr, _) => Some(cpr.batch.clone()),
_ => None,
})
.unique_by(|cpr| cpr.id())
.collect();
if custom_predicate_batches.len() > params.max_custom_predicate_batches {
return Err(Error::custom(format!(
"The number of required `CustomPredicateBatch`es ({}) exceeds the maximum number ({}).",
custom_predicate_batches.len(),
params.max_custom_predicate_batches
)));
}
Ok(custom_predicate_batches)
}
/// Extracts all custom predicate operations with all the data required to verify them.
pub(crate) fn extract_custom_predicate_verifications(
params: &Params,
operations: &[middleware::Operation],
custom_predicate_batches: &[Arc<CustomPredicateBatch>],
) -> Result<Vec<CustomPredicateVerification>> {
let custom_predicate_data: Vec<_> = operations
.iter()
.flat_map(|op| match op {
middleware::Operation::Custom(cpr, sts) => Some((cpr, sts)),
_ => None,
})
.map(|(cpr, sts)| {
let wildcard_values =
resolve_wildcard_values(params, cpr.predicate(), sts).expect("resolved wildcards");
let sts = sts.iter().map(|s| Statement::from(s.clone())).collect();
let batch_index = custom_predicate_batches
.iter()
.enumerate()
.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;
CustomPredicateVerification {
custom_predicate_table_index,
custom_predicate: cpr.clone(),
args: wildcard_values,
op_args: sts,
}
})
.collect();
if custom_predicate_data.len() > params.max_custom_predicate_verifications {
return Err(Error::custom(format!(
"The number of required custom predicate verifications ({}) exceeds the maximum number ({}).",
custom_predicate_data.len(),
params.max_custom_predicate_verifications
)));
}
Ok(custom_predicate_data)
}
/// Extracts Merkle proofs from Contains/NotContains ops.
pub(crate) fn extract_merkle_proofs(
params: &Params,
operations: &[middleware::Operation],
@ -98,11 +165,32 @@ fn find_op_arg(statements: &[Statement], op_arg: &middleware::Statement) -> Resu
}
/// Find the operation auxiliary data in the list of auxiliary data and return the index.
// NOTE: The `custom_predicate_verifications` is optional because in the MainPod we want to store
// the index of a custom predicate verification in the aux data, but in the MockMainPod we don't
// need that because we keep a reference to the custom predicate in the operation type, which
// removes the need for indexing. We could change the OperationType and Predicate for the backend
// to not keep a reference to the custom predicate and instead just keep the id and index and then
// do the same double indexing that the MainPod does to verify custom predicates.
fn find_op_aux(
merkle_proofs: &[MerkleClaimAndProof],
op_aux: &middleware::OperationAux,
custom_predicate_verifications: Option<&[CustomPredicateVerification]>,
op: &middleware::Operation,
) -> Result<OperationAux> {
match op_aux {
let op_aux = op.aux();
let op_type = op.op_type();
if let (OperationType::Custom(cpr), Some(cpvs)) = (op_type, 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)
})
.map(OperationAux::CustomPredVerifyIndex)
.expect("custom predicate verification in the list"));
}
match &op_aux {
middleware::OperationAux::None => Ok(OperationAux::None),
middleware::OperationAux::MerkleProof(pf_arg) => merkle_proofs
.iter()
@ -217,6 +305,7 @@ pub(crate) fn process_private_statements_operations(
params: &Params,
statements: &[Statement],
merkle_proofs: &[MerkleClaimAndProof],
custom_predicate_verifications: Option<&[CustomPredicateVerification]>,
input_operations: &[middleware::Operation],
) -> Result<Vec<Operation>> {
let mut operations = Vec::new();
@ -231,8 +320,7 @@ pub(crate) fn process_private_statements_operations(
.map(|mid_arg| find_op_arg(statements, mid_arg))
.collect::<Result<Vec<_>>>()?;
let mid_aux = op.aux();
let aux = find_op_aux(merkle_proofs, &mid_aux)?;
let aux = find_op_aux(merkle_proofs, custom_predicate_verifications, &op)?;
pad_operation_args(params, &mut args);
operations.push(Operation(op.op_type(), args, aux));
@ -301,12 +389,19 @@ impl Prover {
.collect_vec();
let merkle_proofs = extract_merkle_proofs(params, inputs.operations)?;
let custom_predicate_batches = extract_custom_predicate_batches(params, inputs.operations)?;
let custom_predicate_verifications = extract_custom_predicate_verifications(
params,
inputs.operations,
&custom_predicate_batches,
)?;
let statements = layout_statements(params, &inputs);
let operations = process_private_statements_operations(
params,
&statements,
&merkle_proofs,
Some(&custom_predicate_verifications),
inputs.operations,
)?;
let operations = process_public_statements_operations(params, &statements, operations)?;
@ -321,6 +416,8 @@ impl Prover {
statements: statements[statements.len() - params.max_statements..].to_vec(),
operations,
merkle_proofs,
custom_predicate_batches,
custom_predicate_verifications,
};
main_pod.set_targets(&mut pw, &input)?;
@ -505,4 +602,41 @@ pub mod tests {
let pod = (kyc_pod.pod as Box<dyn Any>).downcast::<MainPod>().unwrap();
pod.verify().unwrap()
}
#[test]
fn test_mainpod_small_empty() {
let params = middleware::Params {
max_input_signed_pods: 0,
max_input_main_pods: 0,
max_statements: 5,
max_signed_pod_values: 2,
max_public_statements: 2,
max_statement_args: 2,
max_operation_args: 3,
max_custom_predicate_batches: 2,
max_custom_predicate_verifications: 2,
max_custom_predicate_arity: 2,
max_custom_predicate_wildcards: 2,
max_custom_batch_size: 2,
max_merkle_proofs: 2,
max_depth_mt_gadget: 4,
};
let pod_builder = frontend::MainPodBuilder::new(&params);
// Mock
let mut prover = MockProver {};
let kyc_pod = pod_builder.prove(&mut prover, &params).unwrap();
let pod = (kyc_pod.pod as Box<dyn Any>)
.downcast::<MockMainPod>()
.unwrap();
pod.verify().unwrap();
println!("{:#}", pod);
// Real
let mut prover = Prover {};
let kyc_pod = pod_builder.prove(&mut prover, &params).unwrap();
let pod = (kyc_pod.pod as Box<dyn Any>).downcast::<MainPod>().unwrap();
pod.verify().unwrap()
}
}