pod2/src/middleware/operation.rs
Ahmad Afuni 1c9b735da6
chore(middleware): add container op checks and tests (#347)
* Add container op checks

* Formatting & comments

* Formatting
2025-07-19 10:25:42 +10:00

638 lines
24 KiB
Rust

use std::{fmt, iter};
use log::error;
use plonky2::field::types::Field;
use serde::{Deserialize, Serialize};
use crate::{
backends::plonky2::primitives::merkletree::{MerkleProof, MerkleTree},
middleware::{
hash_values, AnchoredKey, CustomPredicate, CustomPredicateRef, Error, NativePredicate,
Params, Predicate, Result, Statement, StatementArg, StatementTmplArg, ToFields, Value,
ValueRef, Wildcard, F, SELF,
},
};
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum OperationType {
Native(NativeOperation),
Custom(CustomPredicateRef),
}
#[derive(Clone, Debug, PartialEq)]
pub enum OperationAux {
None,
MerkleProof(MerkleProof),
}
impl fmt::Display for OperationAux {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
Self::None => write!(f, "<no aux>")?,
Self::MerkleProof(pf) => write!(f, "merkle_proof({})", pf)?,
}
Ok(())
}
}
impl ToFields for OperationType {
/// Encoding:
/// - Native(native_op) => [1, [native_op], 0, 0, 0, 0]
/// - Custom(batch, index) => [3, [batch.id], index]
fn to_fields(&self, params: &Params) -> Vec<F> {
let mut fields: Vec<F> = match self {
Self::Native(p) => iter::once(F::from_canonical_u64(1))
.chain(p.to_fields(params))
.collect(),
Self::Custom(CustomPredicateRef { batch, index }) => {
iter::once(F::from_canonical_u64(3))
.chain(batch.id().0)
.chain(iter::once(F::from_canonical_usize(*index)))
.collect()
}
};
fields.resize_with(Params::operation_type_size(), || F::from_canonical_u64(0));
fields
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum NativeOperation {
None = 0,
NewEntry = 1,
CopyStatement = 2,
EqualFromEntries = 3,
NotEqualFromEntries = 4,
LtEqFromEntries = 5,
LtFromEntries = 6,
TransitiveEqualFromStatements = 7,
LtToNotEqual = 8,
ContainsFromEntries = 9,
NotContainsFromEntries = 10,
SumOf = 11,
ProductOf = 12,
MaxOf = 13,
HashOf = 14,
// Syntactic sugar operations. These operations are not supported by the backend. The
// frontend compiler is responsible of translating these operations into the operations above.
DictContainsFromEntries = 1001,
DictNotContainsFromEntries = 1002,
SetContainsFromEntries = 1003,
SetNotContainsFromEntries = 1004,
ArrayContainsFromEntries = 1005,
GtEqFromEntries = 1006,
GtFromEntries = 1007,
GtToNotEqual = 1008,
}
impl NativeOperation {
pub fn is_syntactic_sugar(self) -> bool {
(self as usize) >= 1000
}
}
impl ToFields for NativeOperation {
fn to_fields(&self, _params: &Params) -> Vec<F> {
vec![F::from_canonical_u64(*self as u64)]
}
}
impl OperationType {
/// Gives the type of predicate that the operation will output, if known.
/// CopyStatement may output any predicate (it will match the statement copied),
/// so output_predicate returns None on CopyStatement.
pub fn output_predicate(&self) -> Option<Predicate> {
match self {
OperationType::Native(native_op) => match native_op {
NativeOperation::None => Some(Predicate::Native(NativePredicate::None)),
NativeOperation::NewEntry => Some(Predicate::Native(NativePredicate::Equal)),
NativeOperation::CopyStatement => None,
NativeOperation::EqualFromEntries => {
Some(Predicate::Native(NativePredicate::Equal))
}
NativeOperation::NotEqualFromEntries => {
Some(Predicate::Native(NativePredicate::NotEqual))
}
NativeOperation::LtEqFromEntries => Some(Predicate::Native(NativePredicate::LtEq)),
NativeOperation::LtFromEntries => Some(Predicate::Native(NativePredicate::Lt)),
NativeOperation::TransitiveEqualFromStatements => {
Some(Predicate::Native(NativePredicate::Equal))
}
NativeOperation::LtToNotEqual => Some(Predicate::Native(NativePredicate::NotEqual)),
NativeOperation::ContainsFromEntries => {
Some(Predicate::Native(NativePredicate::Contains))
}
NativeOperation::NotContainsFromEntries => {
Some(Predicate::Native(NativePredicate::NotContains))
}
NativeOperation::SumOf => Some(Predicate::Native(NativePredicate::SumOf)),
NativeOperation::ProductOf => Some(Predicate::Native(NativePredicate::ProductOf)),
NativeOperation::MaxOf => Some(Predicate::Native(NativePredicate::MaxOf)),
NativeOperation::HashOf => Some(Predicate::Native(NativePredicate::HashOf)),
no => unreachable!("Unexpected syntactic sugar op {:?}", no),
},
OperationType::Custom(cpr) => Some(Predicate::Custom(cpr.clone())),
}
}
}
// TODO: Refine this enum.
#[derive(Clone, Debug, PartialEq)]
pub enum Operation {
None,
NewEntry,
CopyStatement(Statement),
EqualFromEntries(Statement, Statement),
NotEqualFromEntries(Statement, Statement),
LtEqFromEntries(Statement, Statement),
LtFromEntries(Statement, Statement),
TransitiveEqualFromStatements(Statement, Statement),
LtToNotEqual(Statement),
ContainsFromEntries(
/* root */ Statement,
/* key */ Statement,
/* value */ Statement,
/* proof */ MerkleProof,
),
NotContainsFromEntries(
/* root */ Statement,
/* key */ Statement,
/* proof */ MerkleProof,
),
SumOf(Statement, Statement, Statement),
ProductOf(Statement, Statement, Statement),
MaxOf(Statement, Statement, Statement),
HashOf(Statement, Statement, Statement),
Custom(CustomPredicateRef, Vec<Statement>),
}
pub(crate) fn sum_op(x: i64, y: i64) -> i64 {
x + y
}
pub(crate) fn prod_op(x: i64, y: i64) -> i64 {
x * y
}
pub(crate) fn max_op(x: i64, y: i64) -> i64 {
x.max(y)
}
pub(crate) fn hash_op(x: Value, y: Value) -> Value {
Value::from(hash_values(&[x, y]))
}
impl Operation {
pub fn op_type(&self) -> OperationType {
type OT = OperationType;
use NativeOperation::*;
match self {
Self::None => OT::Native(None),
Self::NewEntry => OT::Native(NewEntry),
Self::CopyStatement(_) => OT::Native(CopyStatement),
Self::EqualFromEntries(_, _) => OT::Native(EqualFromEntries),
Self::NotEqualFromEntries(_, _) => OT::Native(NotEqualFromEntries),
Self::LtEqFromEntries(_, _) => OT::Native(LtEqFromEntries),
Self::LtFromEntries(_, _) => OT::Native(LtFromEntries),
Self::TransitiveEqualFromStatements(_, _) => OT::Native(TransitiveEqualFromStatements),
Self::LtToNotEqual(_) => OT::Native(LtToNotEqual),
Self::ContainsFromEntries(_, _, _, _) => OT::Native(ContainsFromEntries),
Self::NotContainsFromEntries(_, _, _) => OT::Native(NotContainsFromEntries),
Self::SumOf(_, _, _) => OT::Native(SumOf),
Self::ProductOf(_, _, _) => OT::Native(ProductOf),
Self::MaxOf(_, _, _) => OT::Native(MaxOf),
Self::HashOf(_, _, _) => OT::Native(HashOf),
Self::Custom(cpr, _) => OT::Custom(cpr.clone()),
}
}
pub fn args(&self) -> Vec<Statement> {
match self.clone() {
Self::None => vec![],
Self::NewEntry => vec![],
Self::CopyStatement(s) => vec![s],
Self::EqualFromEntries(s1, s2) => vec![s1, s2],
Self::NotEqualFromEntries(s1, s2) => vec![s1, s2],
Self::LtEqFromEntries(s1, s2) => vec![s1, s2],
Self::LtFromEntries(s1, s2) => vec![s1, s2],
Self::TransitiveEqualFromStatements(s1, s2) => vec![s1, s2],
Self::LtToNotEqual(s) => vec![s],
Self::ContainsFromEntries(s1, s2, s3, _pf) => vec![s1, s2, s3],
Self::NotContainsFromEntries(s1, s2, _pf) => vec![s1, s2],
Self::SumOf(s1, s2, s3) => vec![s1, s2, s3],
Self::ProductOf(s1, s2, s3) => vec![s1, s2, s3],
Self::MaxOf(s1, s2, s3) => vec![s1, s2, s3],
Self::HashOf(s1, s2, s3) => vec![s1, s2, s3],
Self::Custom(_, args) => args,
}
}
/// Extracts auxiliary data from operation.
pub fn aux(&self) -> OperationAux {
match self {
Self::ContainsFromEntries(_, _, _, mp) => OperationAux::MerkleProof(mp.clone()),
Self::NotContainsFromEntries(_, _, mp) => OperationAux::MerkleProof(mp.clone()),
_ => OperationAux::None,
}
}
/// Forms operation from op-code and arguments.
pub fn op(op_code: OperationType, args: &[Statement], aux: &OperationAux) -> Result<Self> {
type OA = OperationAux;
type NO = NativeOperation;
Ok(match op_code {
OperationType::Native(o) => match (o, &args, aux.clone()) {
(NO::None, &[], OA::None) => Self::None,
(NO::NewEntry, &[], OA::None) => Self::NewEntry,
(NO::CopyStatement, &[s], OA::None) => Self::CopyStatement(s.clone()),
(NO::EqualFromEntries, &[s1, s2], OA::None) => {
Self::EqualFromEntries(s1.clone(), s2.clone())
}
(NO::NotEqualFromEntries, &[s1, s2], OA::None) => {
Self::NotEqualFromEntries(s1.clone(), s2.clone())
}
(NO::LtEqFromEntries, &[s1, s2], OA::None) => {
Self::LtEqFromEntries(s1.clone(), s2.clone())
}
(NO::LtFromEntries, &[s1, s2], OA::None) => {
Self::LtFromEntries(s1.clone(), s2.clone())
}
(NO::ContainsFromEntries, &[s1, s2, s3], OA::MerkleProof(pf)) => {
Self::ContainsFromEntries(s1.clone(), s2.clone(), s3.clone(), pf)
}
(NO::NotContainsFromEntries, &[s1, s2], OA::MerkleProof(pf)) => {
Self::NotContainsFromEntries(s1.clone(), s2.clone(), pf)
}
(NO::SumOf, &[s1, s2, s3], OA::None) => {
Self::SumOf(s1.clone(), s2.clone(), s3.clone())
}
(NO::ProductOf, &[s1, s2, s3], OA::None) => {
Self::ProductOf(s1.clone(), s2.clone(), s3.clone())
}
(NO::MaxOf, &[s1, s2, s3], OA::None) => {
Self::MaxOf(s1.clone(), s2.clone(), s3.clone())
}
(NO::HashOf, &[s1, s2, s3], OA::None) => {
Self::HashOf(s1.clone(), s2.clone(), s3.clone())
}
_ => Err(Error::custom(format!(
"Ill-formed operation {:?} with {} arguments {:?} and aux {:?}.",
op_code,
args.len(),
args,
aux
)))?,
},
OperationType::Custom(cpr) => Self::Custom(cpr, args.to_vec()),
})
}
/// Checks the given operation against a statement, and prints information if the check does not pass
pub fn check_and_log(&self, params: &Params, output_statement: &Statement) -> Result<bool> {
let valid: bool = self.check(params, output_statement)?;
if !valid {
error!("Check failed on the following statement");
error!("{}", output_statement);
}
Ok(valid)
}
pub(crate) fn check_int_fn(
v1: &Value,
v2: &Value,
v3: &Value,
f: impl FnOnce(i64, i64) -> i64,
) -> Result<bool> {
let i1: i64 = v1.typed().try_into()?;
let i2: i64 = v2.typed().try_into()?;
let i3: i64 = v3.typed().try_into()?;
Ok(i1 == f(i2, i3))
}
/// Checks the given operation against a statement.
pub fn check(&self, params: &Params, output_statement: &Statement) -> Result<bool> {
use Statement::*;
let deduction_err = || Error::invalid_deduction(self.clone(), output_statement.clone());
let val = |v, s| value_from_op(s, v).ok_or_else(deduction_err);
let b = match (self, output_statement) {
(Self::None, None) => true,
(Self::NewEntry, Equal(ValueRef::Key(AnchoredKey { pod_id, .. }), _)) => {
pod_id == &SELF
}
(Self::CopyStatement(s1), s2) => s1 == s2,
(Self::EqualFromEntries(s1, s2), Equal(v3, v4)) => val(v3, s1)? == val(v4, s2)?,
(Self::NotEqualFromEntries(s1, s2), NotEqual(v3, v4)) => val(v3, s1)? != val(v4, s2)?,
(Self::LtEqFromEntries(s1, s2), LtEq(v3, v4)) => val(v3, s1)? <= val(v4, s2)?,
(Self::LtFromEntries(s1, s2), Lt(v3, v4)) => val(v3, s1)? < val(v4, s2)?,
(
Self::ContainsFromEntries(root_s, key_s, val_s, pf),
Contains(root_v, key_v, val_v),
) => {
let root = val(root_v, root_s)?;
let key = val(key_v, key_s)?;
let value = val(val_v, val_s)?;
MerkleTree::verify(
params.max_depth_mt_containers,
root.raw().into(),
pf,
&key.raw(),
&value.raw(),
)?;
true
}
(Self::NotContainsFromEntries(root_s, key_s, pf), NotContains(root_v, key_v)) => {
let root = val(root_v, root_s)?;
let key = val(key_v, key_s)?;
MerkleTree::verify_nonexistence(
params.max_depth_mt_containers,
root.raw().into(),
pf,
&key.raw(),
)?;
true
}
(
Self::TransitiveEqualFromStatements(Equal(ak1, ak2), Equal(ak3, ak4)),
Equal(ak5, ak6),
) => ak2 == ak3 && ak5 == ak1 && ak6 == ak4,
(Self::LtToNotEqual(Lt(ak1, ak2)), NotEqual(ak3, ak4)) => ak1 == ak3 && ak2 == ak4,
(Self::SumOf(s1, s2, s3), SumOf(v4, v5, v6)) => {
Self::check_int_fn(&val(v4, s1)?, &val(v5, s2)?, &val(v6, s3)?, sum_op)?
}
(Self::ProductOf(s1, s2, s3), ProductOf(v4, v5, v6)) => {
Self::check_int_fn(&val(v4, s1)?, &val(v5, s2)?, &val(v6, s3)?, prod_op)?
}
(Self::MaxOf(s1, s2, s3), MaxOf(v4, v5, v6)) => {
Self::check_int_fn(&val(v4, s1)?, &val(v5, s2)?, &val(v6, s3)?, max_op)?
}
(Self::HashOf(s1, s2, s3), HashOf(v4, v5, v6)) => {
val(v4, s1)? == hash_op(val(v5, s2)?, val(v6, s3)?)
}
(Self::Custom(CustomPredicateRef { batch, index }, args), Custom(cpr, s_args))
if batch == &cpr.batch && index == &cpr.index =>
{
check_custom_pred(params, cpr, args, s_args).map(|_| true)?
}
_ => return Err(deduction_err()),
};
Ok(b)
}
}
/// Check that a StatementArg follows a StatementTmplArg based on the currently mapped wildcards.
/// Update the wildcard map with newly found wildcards.
pub fn check_st_tmpl(
st_tmpl_arg: &StatementTmplArg,
st_arg: &StatementArg,
// Map from wildcards to values that we have seen so far.
wildcard_map: &mut [Option<Value>],
) -> Result<()> {
// Check that the value `v` at wildcard `wc` exists in the map or set it.
fn check_or_set(v: Value, wc: &Wildcard, wildcard_map: &mut [Option<Value>]) -> Result<()> {
if let Some(prev) = &wildcard_map[wc.index] {
if *prev != v {
return Err(Error::invalid_wildcard_assignment(
wc.clone(),
v,
prev.clone(),
));
}
} else {
wildcard_map[wc.index] = Some(v);
}
Ok(())
}
match (st_tmpl_arg, st_arg) {
(StatementTmplArg::None, StatementArg::None) => Ok(()),
(StatementTmplArg::Literal(lhs), StatementArg::Literal(rhs)) if lhs == rhs => Ok(()),
(
StatementTmplArg::AnchoredKey(pod_id_wc, key_tmpl),
StatementArg::Key(AnchoredKey { pod_id, key }),
) => {
let pod_id_ok = check_or_set(Value::from(*pod_id), pod_id_wc, wildcard_map);
pod_id_ok.and_then(|_| {
(key_tmpl == key).then_some(()).ok_or(
Error::mismatched_anchored_key_in_statement_tmpl_arg(
pod_id_wc.clone(),
*pod_id,
key_tmpl.clone(),
key.clone(),
),
)
})
}
(StatementTmplArg::Wildcard(wc), StatementArg::Literal(v)) => {
check_or_set(v.clone(), wc, wildcard_map)
}
_ => Err(Error::mismatched_statement_tmpl_arg(
st_tmpl_arg.clone(),
st_arg.clone(),
)),
}
}
pub fn resolve_wildcard_values(
params: &Params,
pred: &CustomPredicate,
args: &[Statement],
) -> Result<Vec<Value>> {
// Check that all wildcard have consistent values as assigned in the statements while storing a
// map of their values.
// NOTE: We assume the statements have the same order as defined in the custom predicate. For
// disjunctions we expect Statement::None for the unused statements.
let mut wildcard_map = vec![None; params.max_custom_predicate_wildcards];
for (st_tmpl, st) in pred.statements.iter().zip(args) {
let st_args = st.args();
st_tmpl
.args
.iter()
.zip(&st_args)
.try_for_each(|(st_tmpl_arg, st_arg)| {
check_st_tmpl(st_tmpl_arg, st_arg, &mut wildcard_map)
})?;
}
// NOTE: We set unresolved wildcard slots with an empty value. They can be unresolved because
// they are beyond the number of used wildcards in this custom predicate, or they could be
// private arguments that are unused in a particular disjunction.
Ok(wildcard_map
.into_iter()
.map(|opt| opt.unwrap_or(Value::from(0)))
.collect())
}
fn check_custom_pred(
params: &Params,
custom_pred_ref: &CustomPredicateRef,
args: &[Statement],
s_args: &[Value],
) -> Result<()> {
let pred = custom_pred_ref.predicate();
if pred.statements.len() != args.len() {
return Err(Error::diff_amount(
"custom predicate operation".to_string(),
"statements".to_string(),
pred.statements.len(),
args.len(),
));
}
if pred.args_len != s_args.len() {
return Err(Error::diff_amount(
"custom predicate statement".to_string(),
"args".to_string(),
pred.args_len,
s_args.len(),
));
}
// Count the number of statements that match the templates by predicate.
let mut num_matches = 0;
for (st_tmpl, st) in pred.statements.iter().zip(args) {
let st_tmpl_pred = match &st_tmpl.pred {
Predicate::BatchSelf(i) => Predicate::Custom(CustomPredicateRef {
batch: custom_pred_ref.batch.clone(),
index: *i,
}),
p => p.clone(),
};
if st_tmpl_pred == st.predicate() {
num_matches += 1;
}
}
let wildcard_map = resolve_wildcard_values(params, pred, args)?;
// Check that the resolved wildcards match the statement arguments.
for (arg_index, (s_arg, wc_value)) in s_args.iter().zip(wildcard_map.iter()).enumerate() {
if *wc_value != *s_arg {
return Err(Error::mismatched_wildcard_value_and_statement_arg(
wc_value.clone(),
s_arg.clone(),
arg_index,
pred.clone(),
));
}
}
if pred.conjunction {
if num_matches != pred.statements.len() {
return Err(Error::unsatisfied_custom_predicate_conjunction(
pred.clone(),
));
}
} else if num_matches == 0 {
return Err(Error::unsatisfied_custom_predicate_disjunction(
pred.clone(),
));
}
Ok(())
}
impl ToFields for Operation {
fn to_fields(&self, _params: &Params) -> Vec<F> {
todo!()
}
}
impl fmt::Display for Operation {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
writeln!(f, "middleware::Operation:")?;
writeln!(f, " {:?} ", self.op_type())?;
for arg in self.args().iter() {
writeln!(f, " {}", arg)?;
}
Ok(())
}
}
/// Returns the value associated with `output_ref`.
/// If `output_ref` is a concrete value, returns that value.
/// Otherwise, `output_ref` was constructed using an `Equal` statement, and `input_st`
/// must be that statement.
pub(crate) fn value_from_op(input_st: &Statement, output_ref: &ValueRef) -> Option<Value> {
match (input_st, output_ref) {
(Statement::None, ValueRef::Literal(v)) => Some(v.clone()),
(Statement::Equal(r1, ValueRef::Literal(v)), r2) if r1 == r2 => Some(v.clone()),
_ => None,
}
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use crate::{
backends::plonky2::primitives::merkletree::MerkleTree,
middleware::{
hash_value, AnchoredKey, Error, Key, Operation, Params, PodId, Result, Statement,
},
};
#[test]
fn check_container_ops() -> Result<()> {
let params = Params::default();
let pod_id = PodId::default();
let root_ak = AnchoredKey::new(pod_id, Key::new("root".into()));
let key_ak = AnchoredKey::new(pod_id, Key::new("key".into()));
let val_ak = AnchoredKey::new(pod_id, Key::new("value".into()));
// Form Merkle tree
let kvs = (0..10)
.map(|i| (hash_value(&i.into()).into(), i.into()))
.collect::<HashMap<_, _>>();
let mt = MerkleTree::new(params.max_depth_mt_containers, &kvs)?;
let root_s = Statement::Equal(root_ak.clone().into(), mt.root().into());
// Check existence proofs
kvs.iter().try_for_each(|(k, v)| {
// Form op args
let key_s = Statement::Equal(key_ak.clone().into(), (*k).into());
let value_s = Statement::Equal(val_ak.clone().into(), (*v).into());
let (_, pf) = mt.prove(k)?;
// Form op
let op = Operation::ContainsFromEntries(root_s.clone(), key_s, value_s, pf);
// Form output statement
let st = Statement::Contains(
root_ak.clone().into(),
key_ak.clone().into(),
val_ak.clone().into(),
);
// Check op against output statement
op.check(&params, &st).and_then(|ind| {
if ind {
Ok(())
} else {
Err(Error::custom(format!(
"ContainedFromEntries check failed for pair ({},{})",
k, v
)))
}
})
})?;
// Check non-existence proofs similarly
(50..60).try_for_each(|k| {
let key_s = Statement::Equal(key_ak.clone().into(), k.into());
let pf = mt.prove_nonexistence(&k.into())?;
let op = Operation::NotContainsFromEntries(root_s.clone(), key_s, pf);
let st = Statement::NotContains(root_ak.clone().into(), key_ak.clone().into());
op.check(&params, &st).and_then(|ind| {
if ind {
Ok(())
} else {
Err(Error::custom(format!(
"NotContainedFromEntries check failed for key {}",
k
)))
}
})
})
}
}