limit the number of StatementTmpl in CustomPredicate:

arnaucube · arnaucube · commit c302dada3f2b · 2025-02-28T15:57:02.000+01:00
- add constructor method for CustomPredicate - make size checks at the CustomPredicate creation, so that once instantiated we can assume that contains valid data This resolves #79
diff --git a/src/backends/plonky2/mock_main/mod.rs b/src/backends/plonky2/mock_main/mod.rs
@@ -435,7 +435,7 @@ impl Pod for MockMainPod {
                 self.operations[i]
                     .deref(&self.statements[..input_statement_offset + i])
                     .unwrap()
-                    .check(&s.clone().try_into().unwrap())
+                    .check(&self.params, &s.clone().try_into().unwrap())
             })
             .collect::<Result<Vec<_>>>()
             .unwrap();
diff --git a/src/frontend/custom.rs b/src/frontend/custom.rs
@@ -1,8 +1,9 @@
 #![allow(unused)]
+use anyhow::Result;
 use std::sync::Arc;
 
 use crate::middleware::{
-    hash_str, CustomPredicate, CustomPredicateBatch, Hash, HashOrWildcard, NativePredicate,
+    hash_str, CustomPredicate, CustomPredicateBatch, Hash, HashOrWildcard, NativePredicate, Params,
     Predicate, StatementTmpl, StatementTmplArg, ToFields, Value, F,
 };
 
@@ -96,31 +97,34 @@ impl CustomPredicateBatchBuilder {
 
     fn predicate_and(
         &mut self,
+        params: &Params,
         args: &[&str],
         priv_args: &[&str],
         sts: &[StatementTmplBuilder],
-    ) -> Predicate {
-        self.predicate(true, args, priv_args, sts)
+    ) -> Result<Predicate> {
+        self.predicate(params, true, args, priv_args, sts)
     }
 
     fn predicate_or(
         &mut self,
+        params: &Params,
         args: &[&str],
         priv_args: &[&str],
         sts: &[StatementTmplBuilder],
-    ) -> Predicate {
-        self.predicate(false, args, priv_args, sts)
+    ) -> Result<Predicate> {
+        self.predicate(params, false, args, priv_args, sts)
     }
 
     /// creates the custom predicate from the given input, adds it to the
     /// self.predicates, and returns the index of the created predicate
     fn predicate(
         &mut self,
+        params: &Params,
         conjunction: bool,
         args: &[&str],
         priv_args: &[&str],
         sts: &[StatementTmplBuilder],
-    ) -> Predicate {
+    ) -> Result<Predicate> {
         let statements = sts
             .iter()
             .map(|sb| {
@@ -138,13 +142,9 @@ impl CustomPredicateBatchBuilder {
                 StatementTmpl(sb.predicate.clone(), args)
             })
             .collect();
-        let custom_predicate = CustomPredicate {
-            conjunction,
-            statements,
-            args_len: args.len(),
-        };
+        let custom_predicate = CustomPredicate::new(params, conjunction, statements, args.len())?;
         self.predicates.push(custom_predicate);
-        Predicate::BatchSelf(self.predicates.len() - 1)
+        Ok(Predicate::BatchSelf(self.predicates.len() - 1))
     }
 
     fn finish(self) -> Arc<CustomPredicateBatch> {
@@ -174,7 +174,7 @@ mod tests {
     use crate::middleware::{CustomPredicateRef, Params, PodType};
 
     #[test]
-    fn test_custom_pred() {
+    fn test_custom_pred() -> Result<()> {
         use NativePredicate as NP;
         use StatementTmplBuilder as STB;
 
@@ -183,6 +183,7 @@ mod tests {
 
         let mut builder = CustomPredicateBatchBuilder::new("eth_friend".into());
         let _eth_friend = builder.predicate_and(
+            &params,
             // arguments:
             &["src_ori", "src_key", "dst_ori", "dst_key"],
             // private arguments:
@@ -202,7 +203,7 @@ mod tests {
                     .arg(("attestation_pod", literal("attestation")))
                     .arg(("dst_ori", "dst_key")),
             ],
-        );
+        )?;
 
         println!("a.0. eth_friend = {}", builder.predicates.last().unwrap());
         let eth_friend = builder.finish();
@@ -216,6 +217,7 @@ mod tests {
         // >
         let mut builder = CustomPredicateBatchBuilder::new("eth_dos_distance_base".into());
         let eth_dos_distance_base = builder.predicate_and(
+            &params,
             &[
                 // arguments:
                 "src_ori",
@@ -236,7 +238,7 @@ mod tests {
                     .arg(("distance_ori", "distance_key"))
                     .arg(0),
             ],
-        );
+        )?;
         println!(
             "b.0. eth_dos_distance_base = {}",
             builder.predicates.last().unwrap()
@@ -246,6 +248,7 @@ mod tests {
 
         // next chunk builds:
         let eth_dos_distance_ind = builder.predicate_and(
+            &params,
             &[
                 // arguments:
                 "src_ori",
@@ -281,14 +284,15 @@ mod tests {
                     .arg(("intermed_ori", "intermed_key"))
                     .arg(("dst_ori", "dst_key")),
             ],
-        );
+        )?;
 
         println!(
             "b.1. eth_dos_distance_ind = {}",
             builder.predicates.last().unwrap()
         );
 
         let _eth_dos_distance = builder.predicate_or(
+            &params,
             &[
                 "src_ori",
                 "src_key",
@@ -308,7 +312,7 @@ mod tests {
                     .arg(("dst_ori", "dst_key"))
                     .arg(("distance_ori", "distance_key")),
             ],
-        );
+        )?;
 
         println!(
             "b.2. eth_dos_distance = {}",
@@ -318,5 +322,7 @@ mod tests {
         let eth_dos_batch_b = builder.finish();
         let fields = eth_dos_batch_b.to_fields(&params);
         println!("Batch b, serialized: {:?}", fields);
+
+        Ok(())
     }
 }
diff --git a/src/frontend/operation.rs b/src/frontend/operation.rs
@@ -1,7 +1,7 @@
 use std::fmt;
 
 use super::{AnchoredKey, SignedPod, Statement, StatementArg, Value};
-use crate::middleware::{hash_str, NativeOperation, NativePredicate, OperationType, Predicate};
+use crate::middleware::{hash_str, NativePredicate, OperationType, Predicate};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum OperationArg {
diff --git a/src/middleware/custom.rs b/src/middleware/custom.rs
@@ -195,14 +195,42 @@ impl ToFields for StatementTmpl {
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct CustomPredicate {
+    /// NOTE: fields are not public (outside of crate) to enforce the struct instantiation through
+    /// the `::and/or` methods, which performs checks on the values.
+
     /// true for "and", false for "or"
-    pub conjunction: bool,
-    pub statements: Vec<StatementTmpl>,
-    pub args_len: usize,
+    pub(crate) conjunction: bool,
+    pub(crate) statements: Vec<StatementTmpl>,
+    pub(crate) args_len: usize,
     // TODO: Add private args length?
     // TODO: Add args type information?
 }
 
+impl CustomPredicate {
+    pub fn and(params: &Params, statements: Vec<StatementTmpl>, args_len: usize) -> Result<Self> {
+        Self::new(params, true, statements, args_len)
+    }
+    pub fn or(params: &Params, statements: Vec<StatementTmpl>, args_len: usize) -> Result<Self> {
+        Self::new(params, false, statements, args_len)
+    }
+    pub fn new(
+        params: &Params,
+        conjunction: bool,
+        statements: Vec<StatementTmpl>,
+        args_len: usize,
+    ) -> Result<Self> {
+        if statements.len() > params.max_custom_predicate_arity {
+            return Err(anyhow!("Custom predicate depends on too many statements"));
+        }
+
+        Ok(Self {
+            conjunction,
+            statements,
+            args_len,
+        })
+    }
+}
+
 impl ToFields for CustomPredicate {
     fn to_fields(&self, params: &Params) -> (Vec<F>, usize) {
         // serialize as:
@@ -212,9 +240,9 @@ impl ToFields for CustomPredicate {
         //   (params.max_custom_predicate_arity * params.statement_tmpl_size())
         //   field elements
 
-        // TODO think if this check should go into the StatementTmpl creation,
-        // instead of at the `to_fields` method, where we should assume that the
-        // values are already valid
+        // NOTE: this method assumes that the self.params.len() is inside the
+        // expected bound, as Self should be instantiated with the constructor
+        // method `new` which performs the check.
         if self.statements.len() > params.max_custom_predicate_arity {
             panic!("Custom predicate depends on too many statements");
         }
@@ -353,7 +381,7 @@ mod tests {
 
     use crate::middleware::{
         AnchoredKey, CustomPredicate, CustomPredicateBatch, CustomPredicateRef, Hash,
-        HashOrWildcard, NativePredicate, Operation, PodId, PodType, Predicate, Statement,
+        HashOrWildcard, NativePredicate, Operation, Params, PodId, PodType, Predicate, Statement,
         StatementTmpl, StatementTmplArg, SELF,
     };
 
@@ -418,13 +446,16 @@ mod tests {
             ],
         );
 
-        assert!(custom_deduction.check(&custom_statement)?);
+        let params = Params::default();
+        assert!(custom_deduction.check(&params, &custom_statement)?);
 
         Ok(())
     }
 
     #[test]
     fn ethdos_test() -> Result<()> {
+        let params = Params::default();
+
         let eth_friend_cp = CustomPredicate {
             conjunction: true,
             statements: vec![
@@ -561,7 +592,7 @@ mod tests {
         );
 
         // Copies should work.
-        assert!(Operation::CopyStatement(ethdos_example.clone()).check(&ethdos_example)?);
+        assert!(Operation::CopyStatement(ethdos_example.clone()).check(&params, &ethdos_example)?);
 
         // This could arise as the inductive step.
         let ethdos_ind_example = Statement::Custom(
@@ -577,7 +608,7 @@ mod tests {
             CustomPredicateRef(eth_dos_distance_batch.clone(), 2),
             vec![ethdos_ind_example.clone()]
         )
-        .check(&ethdos_example)?);
+        .check(&params, &ethdos_example)?);
 
         // And the inductive step would arise as follows: Say the
         // ETHDoS distance from Alice to Charlie is 6, which is one
@@ -610,7 +641,7 @@ mod tests {
             CustomPredicateRef(eth_dos_distance_batch.clone(), 1),
             ethdos_facts
         )
-        .check(&ethdos_ind_example)?);
+        .check(&params, &ethdos_ind_example)?);
 
         Ok(())
     }
diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs
@@ -92,6 +92,22 @@ pub struct Params {
     pub max_custom_batch_size: usize,
 }
 
+impl Default for Params {
+    fn default() -> Self {
+        Self {
+            max_input_signed_pods: 3,
+            max_input_main_pods: 3,
+            max_statements: 20,
+            max_signed_pod_values: 8,
+            max_public_statements: 10,
+            max_statement_args: 5,
+            max_operation_args: 5,
+            max_custom_predicate_arity: 5,
+            max_custom_batch_size: 5,
+        }
+    }
+}
+
 impl Params {
     pub fn max_priv_statements(&self) -> usize {
         self.max_statements - self.max_public_statements
@@ -134,22 +150,6 @@ impl Params {
     }
 }
 
-impl Default for Params {
-    fn default() -> Self {
-        Self {
-            max_input_signed_pods: 3,
-            max_input_main_pods: 3,
-            max_statements: 20,
-            max_signed_pod_values: 8,
-            max_public_statements: 10,
-            max_statement_args: 5,
-            max_operation_args: 5,
-            max_custom_predicate_arity: 5,
-            max_custom_batch_size: 5,
-        }
-    }
-}
-
 pub trait Pod: fmt::Debug + DynClone {
     fn verify(&self) -> bool;
     fn id(&self) -> PodId;
diff --git a/src/middleware/operation.rs b/src/middleware/operation.rs
@@ -4,7 +4,9 @@ use anyhow::{anyhow, Result};
 
 use super::{CustomPredicateRef, Statement};
 use crate::{
-    middleware::{AnchoredKey, CustomPredicate, PodId, Predicate, StatementTmpl, Value, SELF},
+    middleware::{
+        AnchoredKey, CustomPredicate, Params, PodId, Predicate, StatementTmpl, Value, SELF,
+    },
     util::hashmap_insert_no_dupe,
 };
 
@@ -145,7 +147,7 @@ impl Operation {
         })
     }
     /// Checks the given operation against a statement.
-    pub fn check(&self, output_statement: &Statement) -> Result<bool> {
+    pub fn check(&self, params: &Params, output_statement: &Statement) -> Result<bool> {
         use Statement::*;
         match (self, output_statement) {
             (Self::None, None) => Ok(true),
@@ -211,10 +213,11 @@ impl Operation {
                 // references with custom predicate references.
                 let custom_predicate = {
                     let cp = (**cpb).predicates[*i].clone();
-                    CustomPredicate {
-                        conjunction: cp.conjunction,
-                        statements: cp
-                            .statements
+                    CustomPredicate::new(
+                        params,
+                        cp.conjunction,
+                        // statments:
+                        cp.statements
                             .into_iter()
                             .map(|StatementTmpl(p, args)| {
                                 StatementTmpl(
@@ -228,8 +231,8 @@ impl Operation {
                                 )
                             })
                             .collect(),
-                        args_len: cp.args_len,
-                    }
+                        cp.args_len,
+                    )?
                 };
                 match custom_predicate.conjunction {
                     true if custom_predicate.statements.len() == args.len() => {
