Produce JSON schema for serialized types

robknight · robknight · commit 02d6fb3cb434 · 2025-03-15T14:35:39.000+01:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -19,6 +19,7 @@ plonky2 = { git = "https://github.com/0xPolygonZero/plonky2", optional = true }
 serde = "1.0.219"
 serde_json = "1.0.140"
 base64 = "0.22.1"
+schemars = "1.0.0-alpha.17"
 
 [features]
 default = ["backend_plonky2"]
diff --git a/src/backends/plonky2/basetypes.rs b/src/backends/plonky2/basetypes.rs
@@ -14,6 +14,7 @@ use plonky2::hash::poseidon::PoseidonHash;
 use plonky2::plonk::config::Hasher;
 use plonky2::plonk::config::PoseidonGoldilocksConfig;
 use plonky2::plonk::proof::Proof as Plonky2Proof;
+use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use std::cmp::{Ord, Ordering};
 use std::fmt;
@@ -37,12 +38,16 @@ pub const EMPTY: Value = Value([F::ZERO, F::ZERO, F::ZERO, F::ZERO]);
 pub const SELF_ID_HASH: Hash = Hash([F::ONE, F::ZERO, F::ZERO, F::ZERO]);
 pub const NULL: Hash = Hash([F::ZERO, F::ZERO, F::ZERO, F::ZERO]);
 
-#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
+#[schemars(rename = "MiddlewareValue")]
 pub struct Value(
     #[serde(
         serialize_with = "serialize_value_tuple",
         deserialize_with = "deserialize_value_tuple"
     )]
+    // We know that Serde will serialize and deserialize this as a string, so we can
+    // use the JsonSchema to validate the format.
+    #[schemars(with = "String", regex(pattern = r"^[0-9a-fA-F]{64}$"))]
     pub [F; VALUE_SIZE],
 );
 
@@ -126,12 +131,13 @@ impl fmt::Display for Value {
     }
 }
 
-#[derive(Clone, Copy, Debug, Default, Hash, Eq, PartialEq, Serialize, Deserialize)]
+#[derive(Clone, Copy, Debug, Default, Hash, Eq, PartialEq, Serialize, Deserialize, JsonSchema)]
 pub struct Hash(
     #[serde(
         serialize_with = "serialize_hash_tuple",
         deserialize_with = "deserialize_hash_tuple"
     )]
+    #[schemars(with = "String", regex(pattern = r"^[0-9a-fA-F]{64}$"))]
     pub [F; HASH_SIZE],
 );
 
diff --git a/src/frontend/containers.rs b/src/frontend/containers.rs
@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 
+use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
 use super::Value;
@@ -10,7 +11,7 @@ use crate::middleware::{
     hash_str, Value as MiddlewareValue,
 };
 
-#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, JsonSchema)]
 #[serde(transparent)]
 pub struct Set(Vec<Value>, #[serde(skip)] MiddlewareSet);
 
@@ -40,7 +41,7 @@ impl<'de> Deserialize<'de> for Set {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, JsonSchema)]
 #[serde(transparent)]
 pub struct Dictionary(HashMap<String, Value>, #[serde(skip)] MiddlewareDictionary);
 
@@ -75,7 +76,7 @@ impl<'de> Deserialize<'de> for Dictionary {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, JsonSchema)]
 #[serde(transparent)]
 pub struct Array(Vec<Value>, #[serde(skip)] MiddlewareArray);
 
diff --git a/src/frontend/mod.rs b/src/frontend/mod.rs
@@ -10,6 +10,7 @@ use crate::middleware::{OperationType, Predicate, KEY_SIGNER, KEY_TYPE};
 use anyhow::{anyhow, Error, Result};
 use containers::{Array, Dictionary, Set};
 use itertools::Itertools;
+use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::convert::From;
@@ -25,19 +26,20 @@ pub use operation::*;
 pub use statement::*;
 
 /// This type is just for presentation purposes.
-#[derive(Clone, Debug, Default, h::Hash, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Clone, Debug, Default, h::Hash, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 pub enum PodClass {
     #[default]
     Signed,
     Main,
 }
 
 // An Origin, which represents a reference to an ancestor POD.
-#[derive(Clone, Debug, PartialEq, Eq, h::Hash, Default, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, h::Hash, Default, Serialize, Deserialize, JsonSchema)]
 #[serde(into = "OriginSerdeHelper", from = "OriginSerdeHelper")]
 pub struct Origin(pub PodClass, pub PodId);
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, JsonSchema)]
+#[schemars(rename = "Origin")]
 struct OriginSerdeHelper {
     pod_class: PodClass,
     pod_id: PodId,
@@ -58,7 +60,7 @@ impl From<OriginSerdeHelper> for Origin {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 pub enum Value {
     // Serde cares about the order of the enum variants, with untagged variants
     // appearing at the end.
@@ -212,7 +214,8 @@ impl SignedPodBuilder {
 
 /// SignedPod is a wrapper on top of backend::SignedPod, which additionally stores the
 /// string<-->hash relation of the keys.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(try_from = "SignedPodHelper", into = "SignedPodHelper")]
 pub struct SignedPod {
     pub pod: Box<dyn middleware::Pod>,
     /// Key-value pairs as represented in the frontend. These should
@@ -260,9 +263,12 @@ impl SignedPod {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, h::Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, h::Hash, Serialize, Deserialize, JsonSchema)]
 #[serde(into = "AnchoredKeySerdeHelper", from = "AnchoredKeySerdeHelper")]
-pub struct AnchoredKey(pub Origin, pub String);
+pub struct AnchoredKey(
+    #[schemars(with = "OriginSerdeHelper")] pub Origin,
+    pub String,
+);
 
 impl From<AnchoredKey> for middleware::AnchoredKey {
     fn from(ak: AnchoredKey) -> Self {
@@ -273,8 +279,10 @@ impl From<AnchoredKey> for middleware::AnchoredKey {
 // Tuple structs get serialized to bare JSON arrays, which is not very
 // user-friendly. This helper struct allows us to serialize AnchoredKeys
 // into a structure with named fields.
-#[derive(Serialize, Deserialize)]
-struct AnchoredKeySerdeHelper {
+#[derive(Serialize, Deserialize, JsonSchema)]
+#[schemars(rename = "AnchoredKey")]
+pub struct AnchoredKeySerdeHelper {
+    #[schemars(with = "OriginSerdeHelper")]
     origin: Origin,
     key: String,
 }
@@ -807,7 +815,8 @@ impl MainPodBuilder {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(try_from = "MainPodHelper", into = "MainPodHelper")]
 pub struct MainPod {
     pub pod: Box<dyn middleware::Pod>,
     pub public_statements: Vec<Statement>,
diff --git a/src/frontend/serialization.rs b/src/frontend/serialization.rs
@@ -1,112 +1,88 @@
 use std::collections::HashMap;
 
 use base64::prelude::*;
-use serde::ser::{SerializeStruct, Serializer};
+use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
 use crate::backends::plonky2::mock_main::MockMainPod;
 use crate::backends::plonky2::mock_signed::MockSignedPod;
 use crate::frontend::containers::Dictionary;
 use crate::frontend::Statement;
+use crate::frontend::StatementSerdeHelper;
 use crate::middleware::{PodId, F};
 use crate::middleware::{HASH_SIZE, VALUE_SIZE};
 use plonky2::field::types::Field;
 
 use super::Value;
 use super::{MainPod, SignedPod};
 
-impl Serialize for SignedPod {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: Serializer,
-    {
-        let mut state = serializer.serialize_struct("SignedPod", 3)?;
-        state.serialize_field(
-            "entries",
-            &self
-                .kvs
-                .iter()
-                .map(|(k, v)| (k.clone(), v))
-                .collect::<HashMap<String, &Value>>(),
-        )?;
-
-        let signature = self.pod.serialized_proof();
-
-        state.serialize_field("id", &self.id())?;
-        state.serialize_field("proof", &signature)?;
-        state.serialize_field("pod_class", "Signed")?;
-        state.serialize_field("pod_type", "Mock")?;
-        state.end()
-    }
+#[derive(Serialize, Deserialize, JsonSchema)]
+pub struct SignedPodHelper {
+    entries: HashMap<String, Value>,
+    proof: String,
+    pod_class: String,
+    pod_type: String,
 }
 
-impl<'de> Deserialize<'de> for SignedPod {
-    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
-    where
-        D: serde::Deserializer<'de>,
-    {
-        #[derive(Deserialize)]
-        struct SignedPodHelper {
-            entries: HashMap<String, Value>,
-            proof: String,
-            pod_class: String,
-            pod_type: String,
-        }
+impl TryFrom<SignedPodHelper> for SignedPod {
+    type Error = anyhow::Error;
 
-        let helper = SignedPodHelper::deserialize(deserializer)?;
+    fn try_from(helper: SignedPodHelper) -> Result<SignedPod, Self::Error> {
         if helper.pod_class != "Signed" {
-            return Err(serde::de::Error::custom("pod_class is not Signed"));
+            return Err(anyhow::anyhow!("pod_class is not Signed"));
         }
         if helper.pod_type != "Mock" {
-            return Err(serde::de::Error::custom("pod_type is not Mock"));
+            return Err(anyhow::anyhow!("pod_type is not Mock"));
         }
-        let kvs = helper.entries;
-        let dict = Dictionary::new(kvs.clone()).middleware_dict().clone();
+
+        let dict = Dictionary::new(helper.entries.clone())
+            .middleware_dict()
+            .clone();
         let pod = MockSignedPod::new(PodId(dict.commitment()), helper.proof, dict);
+
         Ok(SignedPod {
             pod: Box::new(pod),
-            kvs: kvs,
+            kvs: helper.entries,
         })
     }
 }
 
-impl Serialize for MainPod {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: Serializer,
-    {
-        let mut state = serializer.serialize_struct("MainPod", 2)?;
-        state.serialize_field("public_statements", &self.public_statements)?;
-        state.serialize_field("id", &self.pod.id())?;
-        state.serialize_field("proof", &self.pod.serialized_proof())?;
-        state.serialize_field("pod_class", "Main")?;
-        state.serialize_field("pod_type", "Mock")?;
-        state.end()
+impl From<SignedPod> for SignedPodHelper {
+    fn from(pod: SignedPod) -> Self {
+        SignedPodHelper {
+            entries: pod.kvs,
+            proof: pod.pod.serialized_proof(),
+            pod_class: "Signed".to_string(),
+            pod_type: "Mock".to_string(),
+        }
     }
 }
 
-impl<'de> Deserialize<'de> for MainPod {
-    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
-    where
-        D: serde::Deserializer<'de>,
-    {
-        #[derive(Deserialize)]
-        struct MainPodHelper {
-            public_statements: Vec<Statement>,
-            proof: String,
-            pod_class: String,
-            pod_type: String,
-        }
+#[derive(Serialize, Deserialize, JsonSchema)]
+pub struct MainPodHelper {
+    #[schemars(with = "Vec<StatementSerdeHelper>")]
+    public_statements: Vec<Statement>,
+    proof: String,
+    pod_class: String,
+    pod_type: String,
+}
 
-        let helper = MainPodHelper::deserialize(deserializer)?;
+impl TryFrom<MainPodHelper> for MainPod {
+    type Error = anyhow::Error; // or you can create a custom error type
+
+    fn try_from(helper: MainPodHelper) -> Result<Self, Self::Error> {
         if helper.pod_class != "Main" {
-            return Err(serde::de::Error::custom("pod_class is not Main"));
+            return Err(anyhow::anyhow!("pod_class is not Main"));
         }
         if helper.pod_type != "Mock" {
-            return Err(serde::de::Error::custom("pod_type is not Mock"));
+            return Err(anyhow::anyhow!("pod_type is not Mock"));
         }
-        let proof = String::from_utf8(BASE64_STANDARD.decode(&helper.proof).unwrap()).unwrap();
-        let pod: MockMainPod = serde_json::from_str(&proof).unwrap();
+
+        let proof = String::from_utf8(BASE64_STANDARD.decode(&helper.proof)?)
+            .map_err(|e| anyhow::anyhow!("Invalid base64 encoding: {}", e))?;
+
+        let pod: MockMainPod = serde_json::from_str(&proof)
+            .map_err(|e| anyhow::anyhow!("Failed to parse proof: {}", e))?;
 
         Ok(MainPod {
             pod: Box::new(pod),
@@ -115,6 +91,17 @@ impl<'de> Deserialize<'de> for MainPod {
     }
 }
 
+impl From<MainPod> for MainPodHelper {
+    fn from(pod: MainPod) -> Self {
+        MainPodHelper {
+            public_statements: pod.public_statements,
+            proof: pod.pod.serialized_proof(),
+            pod_class: "Main".to_string(),
+            pod_type: "Mock".to_string(),
+        }
+    }
+}
+
 pub fn serialize_i64<S>(value: &i64, serializer: S) -> Result<S::Ok, S::Error>
 where
     S: serde::Serializer,
@@ -305,4 +292,15 @@ mod tests {
         assert_eq!(kyc_pod.pod.id(), deserialized.pod.id());
         assert_eq!(kyc_pod.pod.verify(), deserialized.pod.verify());
     }
+
+    #[test]
+    fn test_gen_schema() {
+        let schema = schemars::schema_for!(SignedPodHelper);
+        let schema = serde_json::to_string_pretty(&schema).unwrap();
+        println!("schema: {}", schema);
+
+        let schema = schemars::schema_for!(MainPodHelper);
+        let schema = serde_json::to_string_pretty(&schema).unwrap();
+        println!("schema: {}", schema);
+    }
 }
diff --git a/src/frontend/statement.rs b/src/frontend/statement.rs
@@ -1,12 +1,14 @@
-use super::{AnchoredKey, SignedPod, Value};
+use super::{AnchoredKey, AnchoredKeySerdeHelper, SignedPod, Value};
 use crate::middleware::{self, NativePredicate, Predicate};
 use anyhow::{anyhow, Result};
+use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use std::fmt;
 
-#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 pub enum StatementArg {
     Literal(Value),
+    #[schemars(with = "AnchoredKeySerdeHelper")]
     Key(AnchoredKey),
 }
 
@@ -19,12 +21,13 @@ impl fmt::Display for StatementArg {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 #[serde(into = "StatementSerdeHelper", from = "StatementSerdeHelper")]
 pub struct Statement(pub Predicate, pub Vec<StatementArg>);
 
-#[derive(Serialize, Deserialize)]
-struct StatementSerdeHelper {
+#[derive(Serialize, Deserialize, JsonSchema)]
+#[schemars(rename = "Statement")]
+pub struct StatementSerdeHelper {
     predicate: Predicate,
     args: Vec<StatementArg>,
 }
diff --git a/src/middleware/custom.rs b/src/middleware/custom.rs
diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs
diff --git a/src/middleware/statement.rs b/src/middleware/statement.rs

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,6 @@`
`1`	`1`	`use std::collections::HashMap;`
`2`	`2`
	`3`	`+use schemars::JsonSchema;`
`3`	`4`	`use serde::{Deserialize, Serialize};`
`4`	`5`
`5`	`6`	`use super::Value;`
`@@ -10,7 +11,7 @@ use crate::middleware::{`
`10`	`11`	`hash_str, Value as MiddlewareValue,`
`11`	`12`	`};`
`12`	`13`
`13`		`-#[derive(Clone, Debug, PartialEq, Eq, Serialize)]`
	`14`	`+#[derive(Clone, Debug, PartialEq, Eq, Serialize, JsonSchema)]`
`14`	`15`	`#[serde(transparent)]`
`15`	`16`	`pub struct Set(Vec<Value>, #[serde(skip)] MiddlewareSet);`
`16`	`17`
`@@ -40,7 +41,7 @@ impl<'de> Deserialize<'de> for Set {`
`40`	`41`	`}`
`41`	`42`	`}`
`42`	`43`
`43`		`-#[derive(Clone, Debug, PartialEq, Eq, Serialize)]`
	`44`	`+#[derive(Clone, Debug, PartialEq, Eq, Serialize, JsonSchema)]`
`44`	`45`	`#[serde(transparent)]`
`45`	`46`	`pub struct Dictionary(HashMap<String, Value>, #[serde(skip)] MiddlewareDictionary);`
`46`	`47`
`@@ -75,7 +76,7 @@ impl<'de> Deserialize<'de> for Dictionary {`
`75`	`76`	`}`
`76`	`77`	`}`
`77`	`78`
`78`		`-#[derive(Clone, Debug, PartialEq, Eq, Serialize)]`
	`79`	`+#[derive(Clone, Debug, PartialEq, Eq, Serialize, JsonSchema)]`
`79`	`80`	`#[serde(transparent)]`
`80`	`81`	`pub struct Array(Vec<Value>, #[serde(skip)] MiddlewareArray);`
`81`	`82`