Self-referential predicate hashes as statement template args (#494)

* Support quoted predicate hashes, including self-referential predicates

* Clippy

* Review feedback

src/frontend/custom.rs

@@ -18,6 +18,8 @@ pub enum BuilderArg {
     /// Key: (origin, key), where origin is Wildcard and key is Key
     Key(String, String),
     WildcardLiteral(String),
+    /// Reference to a same-batch predicate's identity hash (resolved by name in finish()).
+    SelfPredicateHash(String),
 }
 
 /// When defining a `BuilderArg`, it can be done from 3 different inputs:
@@ -130,6 +132,8 @@ pub struct CustomPredicateBatchBuilder {
     params: Params,
     pub name: String,
     pub predicates: Vec<CustomPredicate>,
+    /// Forward references to resolve in finish(): (predicate_idx, statement_idx, arg_idx, name)
+    pending_self_pred_hashes: Vec<(usize, usize, usize, String)>,
 }
 
 impl CustomPredicateBatchBuilder {
@@ -138,6 +142,7 @@ impl CustomPredicateBatchBuilder {
             params,
             name,
             predicates: Vec::new(),
+            pending_self_pred_hashes: Vec::new(),
         }
     }
 
@@ -194,14 +199,18 @@ impl CustomPredicateBatchBuilder {
             ));
         }
 
+        let pred_idx = self.predicates.len();
+        let mut pending = Vec::new();
         let statements = sts
             .iter()
-            .map(|sb| {
+            .enumerate()
+            .map(|(stmt_idx, sb)| {
                 let stb = sb.clone().desugar();
                 let st_tmpl_args = stb
                     .args
                     .iter()
-                    .map(|a| {
+                    .enumerate()
+                    .map(|(arg_idx, a)| {
                         Ok::<_, Error>(match a {
                             BuilderArg::Literal(v) => StatementTmplArg::Literal(v.clone()),
                             BuilderArg::Key(root_wc, key_str) => StatementTmplArg::AnchoredKey(
@@ -211,6 +220,22 @@ impl CustomPredicateBatchBuilder {
                             BuilderArg::WildcardLiteral(v) => {
                                 StatementTmplArg::Wildcard(resolve_wildcard(args, priv_args, v)?)
                             }
+                            BuilderArg::SelfPredicateHash(pred_name) => {
+                                // Try backward reference first
+                                match self.predicates.iter().position(|p| p.name == *pred_name) {
+                                    Some(index) => StatementTmplArg::SelfPredicateHash(index),
+                                    None => {
+                                        // Forward reference - placeholder, resolved in finish()
+                                        pending.push((
+                                            pred_idx,
+                                            stmt_idx,
+                                            arg_idx,
+                                            pred_name.clone(),
+                                        ));
+                                        StatementTmplArg::SelfPredicateHash(0)
+                                    }
+                                }
+                            }
                         })
                     })
                     .collect::<Result<_>>()?;
@@ -240,11 +265,27 @@ impl CustomPredicateBatchBuilder {
                 .collect(),
         )?;
         self.predicates.push(custom_predicate);
+        self.pending_self_pred_hashes.extend(pending);
         Ok(Predicate::BatchSelf(self.predicates.len() - 1))
     }
 
-    pub fn finish(self) -> Arc<CustomPredicateBatch> {
-        CustomPredicateBatch::new(self.name, self.predicates)
+    pub fn finish(mut self) -> Result<Arc<CustomPredicateBatch>> {
+        // Resolve forward references for SelfPredicateHash
+        for (pred_idx, stmt_idx, arg_idx, ref name) in &self.pending_self_pred_hashes {
+            let target_idx = self
+                .predicates
+                .iter()
+                .position(|p| p.name == *name)
+                .ok_or_else(|| {
+                    Error::custom(format!(
+                        "SelfPredicateHash references unknown predicate '{}'",
+                        name
+                    ))
+                })?;
+            self.predicates[*pred_idx].statements[*stmt_idx].args[*arg_idx] =
+                StatementTmplArg::SelfPredicateHash(target_idx);
+        }
+        Ok(CustomPredicateBatch::new(self.name, self.predicates))
     }
 }
 
@@ -306,7 +347,7 @@ mod tests {
             .arg("s2");
 
         builder.predicate_and("gt_custom_pred", &["s1", "s2"], &[], &[gt_stb])?;
-        let batch = builder.finish();
+        let batch = builder.finish()?;
         let batch_clone = batch.clone();
         let gt_custom_pred = CustomPredicateRef::new(batch, 0);
 
@@ -356,7 +397,7 @@ mod tests {
             &[],
             &[set_contains_stb],
         )?;
-        let batch = builder.finish();
+        let batch = builder.finish()?;
         let batch_clone = batch.clone();
 
         let mut mp_builder = MainPodBuilder::new(&params, vd_set);
@@ -386,4 +427,83 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_builder_self_predicate_hash_unknown_ref() {
+        let params = Params::default();
+        let mut builder = CustomPredicateBatchBuilder::new(params.clone(), "batch".into());
+
+        let stb = StatementTmplBuilder::new_from_pred(NativePredicate::Equal)
+            .arg("x")
+            .arg(BuilderArg::SelfPredicateHash("nonexistent".into()));
+        builder
+            .predicate_and("pred_A", &["x"], &[], &[stb])
+            .unwrap();
+
+        // finish() should fail because "nonexistent" was never defined
+        assert!(builder.finish().is_err());
+    }
+
+    /// Tests cyclic SelfPredicateHash references end-to-end:
+    /// pred_A references pred_B's hash (forward ref), pred_B references pred_A's hash (backward
+    /// ref). Exercises forward reference resolution in finish(), then builds and verifies a POD
+    /// using pred_A via MockProver.
+    #[test]
+    fn test_builder_self_predicate_hash_e2e() -> Result<()> {
+        let params = Params::default();
+        let vd_set = &*MOCK_VD_SET;
+
+        let mut builder = CustomPredicateBatchBuilder::new(params.clone(), "batch".into());
+
+        // pred_A references pred_B's hash (forward ref, pred_B not yet defined)
+        let stb_a = StatementTmplBuilder::new_from_pred(NativePredicate::Equal)
+            .arg("x")
+            .arg(BuilderArg::SelfPredicateHash("pred_B".into()));
+        builder.predicate_and("pred_A", &["x"], &[], &[stb_a])?;
+
+        // pred_B references pred_A's hash (backward ref, pred_A already defined)
+        let stb_b = StatementTmplBuilder::new_from_pred(NativePredicate::Equal)
+            .arg("x")
+            .arg(BuilderArg::SelfPredicateHash("pred_A".into()));
+        builder.predicate_and("pred_B", &["x"], &[], &[stb_b])?;
+
+        let batch = builder.finish()?;
+
+        // Verify resolution: pred_A references pred_B (index 1), pred_B references pred_A (index 0)
+        assert_eq!(
+            batch.predicates()[0].statements[0].args[1],
+            StatementTmplArg::SelfPredicateHash(1)
+        );
+        assert_eq!(
+            batch.predicates()[1].statements[0].args[1],
+            StatementTmplArg::SelfPredicateHash(0)
+        );
+
+        // Compute concrete hashes
+        let pred_a_ref = CustomPredicateRef::new(batch.clone(), 0);
+        let pred_b_ref = CustomPredicateRef::new(batch.clone(), 1);
+        let pred_b_hash = Value::from(Predicate::Custom(pred_b_ref.clone()).hash());
+
+        // Build a POD using pred_A: Equal(pred_b_hash, pred_b_hash)
+        let mut mp_builder = MainPodBuilder::new(&params, vd_set);
+        let eq_st = mp_builder.priv_op(Operation::eq(pred_b_hash.clone(), pred_b_hash.clone()))?;
+        mp_builder.pub_op(Operation::custom(pred_a_ref, [eq_st]))?;
+
+        // Prove and verify
+        let prover = MockProver {};
+        let proof = mp_builder.prove(&prover)?;
+        proof.pod.verify()?;
+
+        // Verify the public statement contains pred_b_hash as its argument
+        let pub_sts = proof.pod.pub_self_statements();
+        let custom_st = pub_sts
+            .iter()
+            .find(|s| matches!(s, middleware::Statement::Custom(_, _)))
+            .expect("should have a custom statement");
+        if let middleware::Statement::Custom(_, args) = custom_st {
+            assert_eq!(args[0], pred_b_hash);
+        }
+
+        Ok(())
+    }
 }