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

@ -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<Statement>,
// All Merkle proofs
// TODO: Use a backend-specific representation
merkle_proofs: Vec<MerkleProof>,
merkle_proofs: Vec<MerkleClaimAndProof>,
}
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<Vec<MerkleClaimAndProof>> {
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::<Result<Vec<_>>>()?;
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<OperationAux> {
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<Vec<Operation>> {
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::<Vec<_>>();
// Extract Merkle proofs and pad.
let merkle_proofs = Self::extract_merkle_proofs(params, &inputs.operations)?;
let operations = Self::process_private_statements_operations(
params,
&statements,

View file

@ -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<F> {
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<Hash>,
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<Self> {
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<MerkleClaimAndProof> for merkletree::MerkleProof {
type Error = anyhow::Error;
fn try_from(mp: MerkleClaimAndProof) -> Result<Self> {
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::<Vec<_>>()
.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<OperationArg>, 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<crate::middleware::Operation> {
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)
}