add target types for custom predicates (#223)

* add target types for custom predicates

* simplify

* fix clippy

* fix typo

* don't use ref for NativePredicate

* fix wrong len

* apply feedback from @ax0

src/middleware/statement.rs

@@ -6,8 +6,8 @@ use serde::{Deserialize, Serialize};
 use strum_macros::FromRepr;
 
 use crate::middleware::{
-    AnchoredKey, CustomPredicateRef, Error, Key, Params, PodId, Predicate, RawValue, Result,
-    ToFields, Value, F, VALUE_SIZE,
+    AnchoredKey, CustomPredicateRef, Error, Key, Params, PodId, RawValue, Result, ToFields, Value,
+    F, VALUE_SIZE,
 };
 
 // TODO: Maybe store KEY_SIGNER and KEY_TYPE as Key with lazy_static
@@ -84,6 +84,78 @@ impl ToFields for WildcardValue {
     }
 }
 
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
+#[serde(tag = "type", content = "value")]
+pub enum Predicate {
+    Native(NativePredicate),
+    BatchSelf(usize),
+    Custom(CustomPredicateRef),
+}
+
+impl From<NativePredicate> for Predicate {
+    fn from(v: NativePredicate) -> Self {
+        Self::Native(v)
+    }
+}
+
+#[derive(Clone, Copy)]
+pub enum PredicatePrefix {
+    Native = 1,
+    BatchSelf = 2,
+    Custom = 3,
+}
+
+impl From<PredicatePrefix> for F {
+    fn from(prefix: PredicatePrefix) -> Self {
+        Self::from_canonical_usize(prefix as usize)
+    }
+}
+
+impl ToFields for Predicate {
+    fn to_fields(&self, params: &Params) -> Vec<F> {
+        // serialize:
+        // NativePredicate(id) as (1, id, 0, 0, 0, 0) -- id: usize
+        // BatchSelf(i) as (2, i, 0, 0, 0, 0) -- i: usize
+        // CustomPredicateRef(pb, i) as
+        // (3, [hash of pb], i) -- pb hashes to 4 field elements
+        //                      -- i: usize
+
+        // in every case: pad to (hash_size + 2) field elements
+        let mut fields: Vec<F> = match self {
+            Self::Native(p) => iter::once(F::from(PredicatePrefix::Native))
+                .chain(p.to_fields(params))
+                .collect(),
+            Self::BatchSelf(i) => iter::once(F::from(PredicatePrefix::BatchSelf))
+                .chain(iter::once(F::from_canonical_usize(*i)))
+                .collect(),
+            Self::Custom(CustomPredicateRef { batch, index }) => {
+                iter::once(F::from(PredicatePrefix::Custom))
+                    .chain(batch.id(params).0)
+                    .chain(iter::once(F::from_canonical_usize(*index)))
+                    .collect()
+            }
+        };
+        fields.resize_with(Params::predicate_size(), || F::from_canonical_u64(0));
+        fields
+    }
+}
+
+impl fmt::Display for Predicate {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::Native(p) => write!(f, "{:?}", p),
+            Self::BatchSelf(i) => write!(f, "self.{}", i),
+            Self::Custom(CustomPredicateRef { batch, index }) => {
+                write!(
+                    f,
+                    "{}.{}[{}]",
+                    batch.name, index, batch.predicates[*index].name
+                )
+            }
+        }
+    }
+}
+
 /// Type encapsulating statements with their associated arguments.
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
 #[serde(tag = "predicate", content = "args")]