fix pollard serialization/deserialization

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rustreexo"
-version = "0.1.0"
+version = "0.2.0"
 authors = ["Calvin Kim <[email protected]>"]
 edition = "2018"
 description = "A Rust implementation of Utreexo"

src/accumulator/node_hash.rs

@@ -229,6 +229,46 @@ impl NodeHash {
     pub const fn placeholder() -> Self {
         NodeHash::Placeholder
     }
+
+    /// write to buffer
+    pub(super) fn write<W>(&self, writer: &mut W) -> std::io::Result<()>
+    where
+        W: std::io::Write,
+    {
+        match self {
+            Self::Empty => writer.write_all(&[0]),
+            Self::Placeholder => writer.write_all(&[1]),
+            Self::Some(hash) => {
+                writer.write_all(&[2])?;
+                writer.write_all(hash)
+            }
+        }
+    }
+
+    /// Read from buffer
+    pub(super) fn read<R>(reader: &mut R) -> std::io::Result<Self>
+    where
+        R: std::io::Read,
+    {
+        let mut tag = [0];
+        reader.read_exact(&mut tag)?;
+        match tag {
+            [0] => Ok(Self::Empty),
+            [1] => Ok(Self::Placeholder),
+            [2] => {
+                let mut hash = [0; 32];
+                reader.read_exact(&mut hash)?;
+                Ok(Self::Some(hash))
+            }
+            [_] => {
+                let err = std::io::Error::new(
+                    std::io::ErrorKind::InvalidData,
+                    "unexpected tag for NodeHash",
+                );
+                Err(err)
+            }
+        }
+    }
 }
 
 #[cfg(test)]

src/accumulator/pollard.rs

@@ -91,7 +91,7 @@ impl Node {
             NodeType::Branch => writer.write_all(&0_u64.to_le_bytes())?,
             NodeType::Leaf => writer.write_all(&1_u64.to_le_bytes())?,
         }
-        writer.write_all(&*self.data.get())?;
+        self.data.get().write(writer)?;
         self.left
             .borrow()
             .as_ref()
@@ -101,7 +101,7 @@ impl Node {
         self.right
             .borrow()
             .as_ref()
-            .map(|l| l.write_one(writer))
+            .map(|r| r.write_one(writer))
             .transpose()?;
         Ok(())
     }
@@ -118,10 +118,9 @@ impl Node {
             reader: &mut R,
             index: &mut HashMap<NodeHash, Weak<Node>>,
         ) -> std::io::Result<Rc<Node>> {
-            let mut data = [0u8; 32];
             let mut ty = [0u8; 8];
             reader.read_exact(&mut ty)?;
-            reader.read_exact(&mut data)?;
+            let data = NodeHash::read(reader)?;
 
             let ty = match u64::from_le_bytes(ty) {
                 0 => NodeType::Branch,
@@ -131,7 +130,7 @@ impl Node {
             if ty == NodeType::Leaf {
                 let leaf = Rc::new(Node {
                     ty,
-                    data: Cell::new(data.into()),
+                    data: Cell::new(data),
                     parent: RefCell::new(ancestor.map(|a| Rc::downgrade(&a))),
                     left: RefCell::new(None),
                     right: RefCell::new(None),
@@ -141,16 +140,17 @@ impl Node {
             }
             let node = Rc::new(Node {
                 ty: NodeType::Branch,
-                data: Cell::new(data.into()),
+                data: Cell::new(data),
                 parent: RefCell::new(ancestor.map(|a| Rc::downgrade(&a))),
                 left: RefCell::new(None),
                 right: RefCell::new(None),
             });
-            let left = _read_one(Some(node.clone()), reader, index)?;
-            let right = _read_one(Some(node.clone()), reader, index)?;
-            node.left.replace(Some(left));
-            node.right.replace(Some(right));
-
+            if !data.is_empty() {
+                let left = _read_one(Some(node.clone()), reader, index)?;
+                let right = _read_one(Some(node.clone()), reader, index)?;
+                node.left.replace(Some(left));
+                node.right.replace(Some(right));
+            }
             node.left
                 .borrow()
                 .as_ref()
@@ -972,4 +972,19 @@ mod test {
     fn get_hash_vec_of(elements: &[u8]) -> Vec<NodeHash> {
         elements.iter().map(|el| hash_from_u8(*el)).collect()
     }
+
+    #[test]
+    fn test_serialization_roundtrip() {
+        let mut p = Pollard::new();
+        let values = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
+        let hashes: Vec<NodeHash> = values
+            .into_iter()
+            .map(|i| NodeHash::from([i; 32]))
+            .collect();
+        p.modify(&hashes, &[]).expect("modify should work");
+        p.modify(&[], &hashes).expect("modify should work");
+        let mut serialized = Vec::<u8>::new();
+        p.serialize(&mut serialized).expect("serialize should work");
+        let _deserialized = Pollard::deserialize(&*serialized).expect("deserialize should work");
+    }
 }