Towards Contains/NotContains in middleware and backend

ax0 · ax0 · commit d7d704f934c9 · 2025-03-19T07:04:15.000+10:00
diff --git a/src/backends/plonky2/mock_main/mod.rs b/src/backends/plonky2/mock_main/mod.rs
@@ -222,21 +222,24 @@ impl MockMainPod {
 
     fn find_op_arg(
         statements: &[Statement],
-        op_arg: &middleware::Statement,
+        op_arg: &middleware::OperationArg,
     ) -> 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(anyhow!(
-                    "Statement corresponding to op arg {} not found",
-                    op_arg
-                )),
+            middleware::OperationArg::Statement(s_arg) => match s_arg {
+                middleware::Statement::None => Ok(OperationArg::None),
+                _ => statements
+                    .iter()
+                    .enumerate()
+                    .find_map(|(i, s)| {
+                        (&middleware::Statement::try_from(s.clone()).ok()? == s_arg).then_some(i)
+                    })
+                    .map(OperationArg::Index)
+                    .ok_or(anyhow!(
+                        "Statement corresponding to op arg {} not found",
+                        op_arg
+                    )),
+            },
+            _ => todo!(),
         }
     }
 
@@ -285,7 +288,10 @@ impl MockMainPod {
                 Operation(
                     OperationType::Native(NativeOperation::CopyStatement),
                     // TODO
-                    vec![Self::find_op_arg(statements, &mid_arg.try_into().unwrap())?],
+                    vec![Self::find_op_arg(
+                        statements,
+                        &middleware::OperationArg::Statement(mid_arg.try_into().unwrap()),
+                    )?],
                 )
             };
             fill_pad(&mut op.1, OperationArg::None, params.max_operation_args);
diff --git a/src/backends/plonky2/mock_main/operation.rs b/src/backends/plonky2/mock_main/operation.rs
@@ -2,12 +2,17 @@ use anyhow::Result;
 use std::fmt;
 
 use super::Statement;
-use crate::middleware::{self, OperationType};
+use crate::{
+    backends::plonky2::primitives::merkletree::MerkleProof,
+    middleware::{self, OperationType},
+};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum OperationArg {
     None,
     Index(usize),
+    // TODO: Replace with `MerkleProofIndex`.
+    MerkleProof(MerkleProof),
 }
 
 impl OperationArg {
@@ -26,9 +31,17 @@ impl Operation {
             .iter()
             .flat_map(|arg| match arg {
                 OperationArg::None => None,
-                OperationArg::Index(i) => Some(statements[*i].clone().try_into()),
+                OperationArg::Index(i) => Some(
+                    statements[*i]
+                        .clone()
+                        .try_into()
+                        .map(|s| crate::middleware::OperationArg::Statement(s)),
+                ),
+                OperationArg::MerkleProof(pf) => {
+                    Some(Ok(crate::middleware::OperationArg::MerkleProof(pf.clone())))
+                }
             })
-            .collect::<Result<Vec<crate::middleware::Statement>>>()?;
+            .collect::<Result<Vec<_>>>()?;
         middleware::Operation::op(self.0.clone(), &deref_args)
     }
 }
@@ -44,6 +57,7 @@ impl fmt::Display for Operation {
                 match arg {
                     OperationArg::None => write!(f, "none")?,
                     OperationArg::Index(i) => write!(f, "{:02}", i)?,
+                    OperationArg::MerkleProof(pf) => write!(f, "{}", pf)?,
                 }
             }
         }
diff --git a/src/backends/plonky2/primitives/merkletree.rs b/src/backends/plonky2/primitives/merkletree.rs
@@ -207,7 +207,7 @@ impl fmt::Display for MerkleTree {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct MerkleProof {
     // note: currently we don't use the `_existence` field, we would use if we merge the methods
     // `verify` and `verify_nonexistence` into a single one
diff --git a/src/middleware/operation.rs b/src/middleware/operation.rs
@@ -3,14 +3,33 @@ use std::fmt;
 use anyhow::{anyhow, Result};
 
 use super::{CustomPredicateRef, NativePredicate, Statement, StatementArg};
-use crate::middleware::{AnchoredKey, Params, Predicate, Value, SELF};
+use crate::{
+    backends::plonky2::primitives::merkletree::{MerkleProof, MerkleTree},
+    middleware::{AnchoredKey, Params, Predicate, Value, SELF},
+};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum OperationType {
     Native(NativeOperation),
     Custom(CustomPredicateRef),
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum OperationArg {
+    Statement(Statement),
+    MerkleProof(MerkleProof),
+}
+
+impl fmt::Display for OperationArg {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::Statement(s) => write!(f, "{}", s)?,
+            Self::MerkleProof(pf) => write!(f, "merkle_proof({})", pf)?,
+        }
+        Ok(())
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum NativeOperation {
     None = 0,
@@ -81,8 +100,17 @@ pub enum Operation {
     TransitiveEqualFromStatements(Statement, Statement),
     GtToNotEqual(Statement),
     LtToNotEqual(Statement),
-    ContainsFromEntries(Statement, Statement, Statement),
-    NotContainsFromEntries(Statement, 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),
@@ -104,72 +132,95 @@ impl Operation {
             Self::TransitiveEqualFromStatements(_, _) => OT::Native(TransitiveEqualFromStatements),
             Self::GtToNotEqual(_) => OT::Native(GtToNotEqual),
             Self::LtToNotEqual(_) => OT::Native(LtToNotEqual),
-            Self::ContainsFromEntries(_, _, _) => OT::Native(ContainsFromEntries),
-            Self::NotContainsFromEntries(_, _) => OT::Native(NotContainsFromEntries),
+            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::Custom(cpr, _) => OT::Custom(cpr.clone()),
         }
     }
 
-    pub fn args(&self) -> Vec<Statement> {
+    pub fn args(&self) -> Vec<OperationArg> {
+        use OperationArg::{MerkleProof as MP, Statement as S};
         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::GtFromEntries(s1, s2) => vec![s1, s2],
-            Self::LtFromEntries(s1, s2) => vec![s1, s2],
-            Self::TransitiveEqualFromStatements(s1, s2) => vec![s1, s2],
-            Self::GtToNotEqual(s) => vec![s],
-            Self::LtToNotEqual(s) => vec![s],
-            Self::ContainsFromEntries(s1, s2, s3) => vec![s1, s2, s3],
-            Self::NotContainsFromEntries(s1, s2) => 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::Custom(_, args) => args,
+            Self::CopyStatement(s) => vec![S(s)],
+            Self::EqualFromEntries(s1, s2) => vec![S(s1), S(s2)],
+            Self::NotEqualFromEntries(s1, s2) => vec![S(s1), S(s2)],
+            Self::GtFromEntries(s1, s2) => vec![S(s1), S(s2)],
+            Self::LtFromEntries(s1, s2) => vec![S(s1), S(s2)],
+            Self::TransitiveEqualFromStatements(s1, s2) => vec![S(s1), S(s2)],
+            Self::GtToNotEqual(s) => vec![S(s)],
+            Self::LtToNotEqual(s) => vec![S(s)],
+            Self::ContainsFromEntries(s1, s2, s3, pf) => vec![S(s1), S(s2), S(s3), MP(pf)],
+            Self::NotContainsFromEntries(s1, s2, pf) => vec![S(s1), S(s2), MP(pf)],
+            Self::SumOf(s1, s2, s3) => vec![S(s1), S(s2), S(s3)],
+            Self::ProductOf(s1, s2, s3) => vec![S(s1), S(s2), S(s3)],
+            Self::MaxOf(s1, s2, s3) => vec![S(s1), S(s2), S(s3)],
+            Self::Custom(_, args) => args.into_iter().map(|s| S(s)).collect(),
         }
     }
     /// Forms operation from op-code and arguments.
-    pub fn op(op_code: OperationType, args: &[Statement]) -> Result<Self> {
+    pub fn op(op_code: OperationType, args: &[OperationArg]) -> Result<Self> {
+        use OperationArg::{MerkleProof as MP, Statement as S};
         type NO = NativeOperation;
         let arg_tup = (
             args.first().cloned(),
             args.get(1).cloned(),
             args.get(2).cloned(),
+            args.get(3).cloned(),
         );
         Ok(match op_code {
             OperationType::Native(o) => match (o, arg_tup, args.len()) {
-                (NO::None, (None, None, None), 0) => Self::None,
-                (NO::NewEntry, (None, None, None), 0) => Self::NewEntry,
-                (NO::CopyStatement, (Some(s), None, None), 1) => Self::CopyStatement(s),
-                (NO::EqualFromEntries, (Some(s1), Some(s2), None), 2) => {
+                (NO::None, (None, None, None, None), 0) => Self::None,
+                (NO::NewEntry, (None, None, None, None), 0) => Self::NewEntry,
+                (NO::CopyStatement, (Some(S(s)), None, None, None), 1) => Self::CopyStatement(s),
+                (NO::EqualFromEntries, (Some(S(s1)), Some(S(s2)), None, None), 2) => {
                     Self::EqualFromEntries(s1, s2)
                 }
-                (NO::NotEqualFromEntries, (Some(s1), Some(s2), None), 2) => {
+                (NO::NotEqualFromEntries, (Some(S(s1)), Some(S(s2)), None, None), 2) => {
                     Self::NotEqualFromEntries(s1, s2)
                 }
-                (NO::GtFromEntries, (Some(s1), Some(s2), None), 2) => Self::GtFromEntries(s1, s2),
-                (NO::LtFromEntries, (Some(s1), Some(s2), None), 2) => Self::LtFromEntries(s1, s2),
-                (NO::ContainsFromEntries, (Some(s1), Some(s2), Some(s3)), 3) => {
-                    Self::ContainsFromEntries(s1, s2, s3)
+                (NO::GtFromEntries, (Some(S(s1)), Some(S(s2)), None, None), 2) => {
+                    Self::GtFromEntries(s1, s2)
+                }
+                (NO::LtFromEntries, (Some(S(s1)), Some(S(s2)), None, None), 2) => {
+                    Self::LtFromEntries(s1, s2)
+                }
+                (
+                    NO::ContainsFromEntries,
+                    (Some(S(s1)), Some(S(s2)), Some(S(s3)), Some(MP(pf))),
+                    4,
+                ) => Self::ContainsFromEntries(s1, s2, s3, pf),
+                (NO::NotContainsFromEntries, (Some(S(s1)), Some(S(s2)), Some(MP(pf)), None), 3) => {
+                    Self::NotContainsFromEntries(s1, s2, pf)
+                }
+                (NO::SumOf, (Some(S(s1)), Some(S(s2)), Some(S(s3)), None), 3) => {
+                    Self::SumOf(s1, s2, s3)
                 }
-                (NO::NotContainsFromEntries, (Some(s1), Some(s2), None), 2) => {
-                    Self::NotContainsFromEntries(s1, s2)
+                (NO::ProductOf, (Some(S(s1)), Some(S(s2)), Some(S(s3)), None), 3) => {
+                    Self::ProductOf(s1, s2, s3)
+                }
+                (NO::MaxOf, (Some(S(s1)), Some(S(s2)), Some(S(s3)), None), 3) => {
+                    Self::MaxOf(s1, s2, s3)
                 }
-                (NO::SumOf, (Some(s1), Some(s2), Some(s3)), 3) => Self::SumOf(s1, s2, s3),
-                (NO::ProductOf, (Some(s1), Some(s2), Some(s3)), 3) => Self::ProductOf(s1, s2, s3),
-                (NO::MaxOf, (Some(s1), Some(s2), Some(s3)), 3) => Self::MaxOf(s1, s2, s3),
                 _ => Err(anyhow!(
                     "Ill-formed operation {:?} with arguments {:?}.",
                     op_code,
                     args
                 ))?,
             },
-            OperationType::Custom(cpr) => Self::Custom(cpr, args.to_vec()),
+            OperationType::Custom(cpr) => Self::Custom(
+                cpr,
+                args.iter()
+                    .map(|a| match a {
+                        S(s) => Ok(s.clone()),
+                        _ => Err(anyhow!("Invalid argument to custom operation: {:?}", a)),
+                    })
+                    .collect::<Result<Vec<_>>>()?,
+            ),
         })
     }
     /// Gives the output statement of the given operation, where determined
@@ -245,20 +296,25 @@ impl Operation {
             Self::LtToNotEqual(_) => {
                 return Err(anyhow!("Invalid operation"));
             }
-            Self::ContainsFromEntries(ValueOf(ak1, v1), ValueOf(ak2, v2), ValueOf(ak3, v3)) =>
-            /* TODO */
+            Self::ContainsFromEntries(ValueOf(ak1, v1), ValueOf(ak2, v2), ValueOf(ak3, v3), pf)
+                if MerkleTree::verify(pf.siblings.len(), (*v1).into(), &pf, v2, v3)? == () =>
             {
-                Some(vec![StatementArg::Key(*ak1), StatementArg::Key(*ak2)])
+                Some(vec![
+                    StatementArg::Key(*ak1),
+                    StatementArg::Key(*ak2),
+                    StatementArg::Key(*ak3),
+                ])
             }
-            Self::ContainsFromEntries(_, _, _) => {
+            Self::ContainsFromEntries(_, _, _, _) => {
                 return Err(anyhow!("Invalid operation"));
             }
-            Self::NotContainsFromEntries(ValueOf(ak1, v1), ValueOf(ak2, v2)) =>
-            /* TODO */
+            Self::NotContainsFromEntries(ValueOf(ak1, v1), ValueOf(ak2, v2), pf)
+                if MerkleTree::verify_nonexistence(pf.siblings.len(), (*v1).into(), &pf, v2)?
+                    == () =>
             {
                 Some(vec![StatementArg::Key(*ak1), StatementArg::Key(*ak2)])
             }
-            Self::NotContainsFromEntries(_, _) => {
+            Self::NotContainsFromEntries(_, _, _) => {
                 return Err(anyhow!("Invalid operation"));
             }
             Self::SumOf(ValueOf(ak1, v1), ValueOf(ak2, v2), ValueOf(ak3, v3)) => {
@@ -327,12 +383,12 @@ impl Operation {
             (Self::LtFromEntries(ValueOf(ak1, v1), ValueOf(ak2, v2)), Lt(ak3, ak4)) => {
                 Ok(v1 < v2 && ak3 == ak1 && ak4 == ak2)
             }
-            (Self::ContainsFromEntries(_, _, _), Contains(_, _, _)) =>
+            (Self::ContainsFromEntries(_, _, _, _), Contains(_, _, _)) =>
             /* TODO */
             {
                 Ok(true)
             }
-            (Self::NotContainsFromEntries(_, _), NotContains(_, _)) =>
+            (Self::NotContainsFromEntries(_, _, _), NotContains(_, _)) =>
             /* TODO */
             {
                 Ok(true)
diff --git a/src/middleware/statement.rs b/src/middleware/statement.rs
@@ -39,8 +39,12 @@ pub enum Statement {
     NotEqual(AnchoredKey, AnchoredKey),
     Gt(AnchoredKey, AnchoredKey),
     Lt(AnchoredKey, AnchoredKey),
-    Contains(AnchoredKey, AnchoredKey, AnchoredKey),
-    NotContains(AnchoredKey, AnchoredKey),
+    Contains(
+        /* root  */ AnchoredKey,
+        /* key   */ AnchoredKey,
+        /* value */ AnchoredKey,
+    ),
+    NotContains(/* root  */ AnchoredKey, /* key   */ AnchoredKey),
     SumOf(AnchoredKey, AnchoredKey, AnchoredKey),
     ProductOf(AnchoredKey, AnchoredKey, AnchoredKey),
     MaxOf(AnchoredKey, AnchoredKey, AnchoredKey),

Original file line number	Diff line number	Diff line change
`@@ -207,7 +207,7 @@ impl fmt::Display for MerkleTree {`
`207`	`207`	`}`
`208`	`208`	`}`
`209`	`209`
`210`		`-#[derive(Clone, Debug)]`
	`210`	`+#[derive(Clone, Debug, PartialEq, Eq)]`
`211`	`211`	`pub struct MerkleProof {`
`212`	`212`	// note: currently we don't use the `_existence` field, we would use if we merge the methods
`213`	`213`	// `verify` and `verify_nonexistence` into a single one