Re-implement serialization (#201)

* Serialization tests now pass again

* Tidy up and test more edge-cases

* Use attributes rather than custom serializer for arrays

* Add JSON Schema support

* Tests for JSON Schema generation and validation

* Add comments

* Support custom predicates

* Clippy fixes

* Make deserialization/constructor functions pub(crate)

src/middleware/containers.rs

@@ -3,7 +3,10 @@ use std::collections::{HashMap, HashSet};
 /// This file implements the types defined at
 /// https://0xparc.github.io/pod2/values.html#dictionary-array-set .
 use anyhow::{anyhow, Result};
+use schemars::JsonSchema;
+use serde::{Deserialize, Deserializer, Serialize};
 
+use super::serialization::{ordered_map, ordered_set};
 #[cfg(feature = "backend_plonky2")]
 use crate::backends::plonky2::primitives::merkletree::{MerkleProof, MerkleTree};
 use crate::{
@@ -14,9 +17,12 @@ use crate::{
 /// Dictionary: the user original keys and values are hashed to be used in the leaf.
 ///    leaf.key=hash(original_key)
 ///    leaf.value=hash(original_value)
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Serialize)]
+#[serde(transparent)]
 pub struct Dictionary {
+    #[serde(skip)]
     mt: MerkleTree,
+    #[serde(serialize_with = "ordered_map")]
     kvs: HashMap<Key, Value>,
 }
 
@@ -76,12 +82,36 @@ impl PartialEq for Dictionary {
 }
 impl Eq for Dictionary {}
 
+impl<'de> Deserialize<'de> for Dictionary {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let kvs: HashMap<Key, Value> = HashMap::deserialize(deserializer)?;
+        Dictionary::new(kvs).map_err(serde::de::Error::custom)
+    }
+}
+
+impl JsonSchema for Dictionary {
+    fn schema_name() -> String {
+        "Dictionary".to_string()
+    }
+
+    fn json_schema(gen: &mut schemars::gen::SchemaGenerator) -> schemars::schema::Schema {
+        // Just use the schema of HashMap<Key, Value> since that's what we're actually serializing
+        <HashMap<Key, Value>>::json_schema(gen)
+    }
+}
+
 /// Set: the value field of the leaf is unused, and the key contains the hash of the element.
 ///    leaf.key=hash(original_value)
 ///    leaf.value=0
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Serialize)]
+#[serde(transparent)]
 pub struct Set {
+    #[serde(skip)]
     mt: MerkleTree,
+    #[serde(serialize_with = "ordered_set")]
     set: HashSet<Value>,
 }
 
@@ -134,12 +164,38 @@ impl PartialEq for Set {
 }
 impl Eq for Set {}
 
+impl<'de> Deserialize<'de> for Set {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        // Deserialize the set directly
+        let set: HashSet<Value> = HashSet::deserialize(deserializer)?;
+
+        // Create a new Set using the set field
+        Set::new(set).map_err(serde::de::Error::custom)
+    }
+}
+
+impl JsonSchema for Set {
+    fn schema_name() -> String {
+        "Set".to_string()
+    }
+
+    fn json_schema(gen: &mut schemars::gen::SchemaGenerator) -> schemars::schema::Schema {
+        // Just use the schema of HashSet<Value> since that's what we're actually serializing
+        <HashSet<Value>>::json_schema(gen)
+    }
+}
+
 /// Array: the elements are placed at the value field of each leaf, and the key field is just the
 /// array index (integer).
 ///    leaf.key=i
 ///    leaf.value=original_value
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Serialize)]
+#[serde(transparent)]
 pub struct Array {
+    #[serde(skip)]
     mt: MerkleTree,
     array: Vec<Value>,
 }
@@ -190,3 +246,24 @@ impl PartialEq for Array {
     }
 }
 impl Eq for Array {}
+
+impl<'de> Deserialize<'de> for Array {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let array: Vec<Value> = Vec::deserialize(deserializer)?;
+        Array::new(array).map_err(serde::de::Error::custom)
+    }
+}
+
+impl JsonSchema for Array {
+    fn schema_name() -> String {
+        "Array".to_string()
+    }
+
+    fn json_schema(gen: &mut schemars::gen::SchemaGenerator) -> schemars::schema::Schema {
+        // Just use the schema of Vec<Value> since that's what we're actually serializing
+        <Vec<Value>>::json_schema(gen)
+    }
+}