Check a single POD against a POD Request (#359)

src/frontend/pod_request.rs

@@ -0,0 +1,255 @@
+use std::{collections::HashMap, fmt::Display};
+
+use crate::{
+    frontend::{Error, Result},
+    lang::PrettyPrint,
+    middleware::{Pod, Statement, StatementArg, StatementTmpl, StatementTmplArg, Value},
+};
+
+/// Represents a request for a POD, in terms of a set of statement templates.
+/// The response should be a POD (or PODs) containing a set of statements which
+/// satisfy the templates, with consistent wildcard bindings across all templates.
+#[derive(Debug, Clone, PartialEq)]
+pub struct PodRequest {
+    pub request_templates: Vec<StatementTmpl>,
+}
+
+impl PodRequest {
+    pub fn new(request_templates: Vec<StatementTmpl>) -> Self {
+        Self { request_templates }
+    }
+
+    /// Checks if the request is fully satisfied by a single supplied POD.
+    /// This checks for exact matches to the statement templates; that is to say
+    /// that it performs a "syntactic" match, not a "semantic" match; no
+    /// processing of the semantics of the statements is performed.
+    pub fn exact_match_pod(&self, pod: &dyn Pod) -> Result<HashMap<String, Value>> {
+        let pod_statements = pod.pub_statements();
+        let mut bindings: HashMap<String, Value> = HashMap::new();
+
+        if self.dfs_match_all(&pod_statements, &mut bindings, 0) {
+            Ok(bindings)
+        } else {
+            Err(Error::pod_request_validation("No match found".to_string()))
+        }
+    }
+
+    /// Performs a depth-first search through the statement templates, trying to
+    /// match each template to a statement in the POD.
+    /// Returns true if all templates are matched, false otherwise.
+    /// The bindings map is used to store the bindings of the wildcards to the
+    /// values in the POD.
+    /// The template_idx is used to track the current template being matched.
+    fn dfs_match_all(
+        &self,
+        pod_statements: &[Statement],
+        bindings: &mut HashMap<String, Value>,
+        template_idx: usize,
+    ) -> bool {
+        // Base case: all templates matched
+        if template_idx >= self.request_templates.len() {
+            return true;
+        }
+
+        let template = &self.request_templates[template_idx];
+
+        // Try to match this template with each statement in the POD
+        for stmt in pod_statements {
+            if let Some(new_bindings) = self.try_match_template(template, stmt, bindings) {
+                let original_bindings = bindings.clone();
+                bindings.extend(new_bindings);
+
+                if self.dfs_match_all(pod_statements, bindings, template_idx + 1) {
+                    return true;
+                }
+
+                *bindings = original_bindings;
+            }
+        }
+
+        false
+    }
+
+    fn try_match_template(
+        &self,
+        template: &StatementTmpl,
+        statement: &Statement,
+        current_bindings: &HashMap<String, Value>,
+    ) -> Option<HashMap<String, Value>> {
+        if template.pred != statement.predicate() {
+            return None;
+        }
+
+        let template_args = template.args();
+        let statement_args = statement.args();
+
+        if template_args.len() != statement_args.len() {
+            return None;
+        }
+
+        let mut new_bindings = HashMap::new();
+
+        for (tmpl_arg, stmt_arg) in template_args.iter().zip(statement_args.iter()) {
+            if !self.try_match_arg(tmpl_arg, stmt_arg, current_bindings, &mut new_bindings) {
+                return None;
+            }
+        }
+
+        Some(new_bindings)
+    }
+
+    fn try_match_arg(
+        &self,
+        template_arg: &StatementTmplArg,
+        statement_arg: &StatementArg,
+        current_bindings: &HashMap<String, Value>,
+        new_bindings: &mut HashMap<String, Value>,
+    ) -> bool {
+        match (template_arg, statement_arg) {
+            // Literal must match exactly
+            (StatementTmplArg::Literal(tmpl_val), StatementArg::Literal(stmt_val)) => {
+                tmpl_val == stmt_val
+            }
+
+            // Wildcard can bind to any literal value
+            (StatementTmplArg::Wildcard(wildcard), StatementArg::Literal(stmt_val)) => self
+                .try_bind_wildcard(
+                    &wildcard.name,
+                    stmt_val.clone(),
+                    current_bindings,
+                    new_bindings,
+                ),
+
+            // AnchoredKey wildcard must match statement's anchored key
+            (StatementTmplArg::AnchoredKey(wildcard, tmpl_key), StatementArg::Key(stmt_key)) => {
+                // Check if keys match
+                if tmpl_key != &stmt_key.key {
+                    return false;
+                }
+
+                // Try to bind wildcard to the POD ID
+                let pod_id_value = Value::from(stmt_key.pod_id);
+                self.try_bind_wildcard(&wildcard.name, pod_id_value, current_bindings, new_bindings)
+            }
+
+            // Other combinations don't match
+            _ => false,
+        }
+    }
+
+    fn try_bind_wildcard(
+        &self,
+        wildcard_name: &str,
+        value: Value,
+        current_bindings: &HashMap<String, Value>,
+        new_bindings: &mut HashMap<String, Value>,
+    ) -> bool {
+        // Check if wildcard is already bound
+        if let Some(existing_value) = current_bindings.get(wildcard_name) {
+            // Must match existing binding
+            return existing_value == &value;
+        }
+
+        // Check if we're trying to bind it in new_bindings
+        if let Some(existing_value) = new_bindings.get(wildcard_name) {
+            // Must match existing binding in this attempt
+            return existing_value == &value;
+        }
+
+        // Bind the wildcard
+        new_bindings.insert(wildcard_name.to_string(), value);
+        true
+    }
+
+    pub fn templates(&self) -> &[StatementTmpl] {
+        &self.request_templates
+    }
+}
+
+impl Display for PodRequest {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.fmt_podlang(f)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        backends::plonky2::{
+            mock::mainpod::MockProver, primitives::ec::schnorr::SecretKey, signedpod::Signer,
+        },
+        examples::{zu_kyc_pod_builder, zu_kyc_pod_request, zu_kyc_sign_pod_builders, MOCK_VD_SET},
+        frontend::{MainPodBuilder, Operation},
+        lang::parse,
+        middleware::Params,
+    };
+
+    #[test]
+    fn test_pod_request_exact_match_pod() {
+        let params = Params::default();
+        let vd_set = &*MOCK_VD_SET;
+
+        let (gov_id, pay_stub) = zu_kyc_sign_pod_builders(&params);
+        let gov_id = gov_id.sign(&Signer(SecretKey(1u32.into()))).unwrap();
+        let pay_stub = pay_stub.sign(&Signer(SecretKey(2u32.into()))).unwrap();
+        let builder = zu_kyc_pod_builder(&Params::default(), vd_set, &gov_id, &pay_stub).unwrap();
+        let prover = MockProver {};
+        let kyc = builder.prove(&prover, &params).unwrap();
+
+        // This request matches the POD
+        let request = zu_kyc_pod_request(
+            gov_id.get("_signer").unwrap(),
+            pay_stub.get("_signer").unwrap(),
+        )
+        .unwrap();
+        assert!(request.exact_match_pod(&*kyc.pod).is_ok());
+
+        // This request does not match the POD, because the POD does not contain a NotEqual statement.
+        let non_matching_request = parse(
+            r#"
+        REQUEST(
+            NotEqual(4, 5)
+        )
+        "#,
+            &params,
+            &[],
+        )
+        .unwrap()
+        .request;
+        assert!(non_matching_request.exact_match_pod(&*kyc.pod).is_err());
+    }
+
+    #[test]
+    fn test_ambiguous_pod() {
+        let params = Params::default();
+        let vd_set = &*MOCK_VD_SET;
+
+        let mut builder = MainPodBuilder::new(&params, vd_set);
+        let _sum_of_stmt_1 = builder.pub_op(Operation::sum_of(11, 1, 10));
+        let _sum_of_stmt_2 = builder.pub_op(Operation::sum_of(10, 9, 1));
+        let _eq_stmt = builder.pub_op(Operation::eq(10, 10));
+
+        let prover = MockProver {};
+
+        let pod = builder.prove(&prover, &params).unwrap();
+
+        println!("{pod}");
+
+        let request = parse(
+            r#"
+        REQUEST(
+            SumOf(?a, ?b, ?c)
+            Equal(?a, 10)
+        )
+        "#,
+            &params,
+            &[],
+        )
+        .unwrap();
+
+        let bindings = request.request.exact_match_pod(&*pod.pod).unwrap();
+        assert_eq!(*bindings.get("a").unwrap(), 10.into());
+        assert_eq!(*bindings.get("b").unwrap(), 9.into());
+        assert_eq!(*bindings.get("c").unwrap(), 1.into());
+    }
+}