From d8c04e8143e2ee1bc958ea3996bfe3bdc8454b45 Mon Sep 17 00:00:00 2001
From: NikVolf <nikvolf@gmail.com>
Date: Tue, 3 Sep 2019 17:22:21 +0300
Subject: [PATCH] node serialization

---
 Cargo.toml       |  4 ++-
 src/lib.rs       |  4 ++-
 src/node_data.rs | 65 +++++++++++++++++++++++++++++++++++++++++++++---
 src/tree.rs      | 12 ++++-----
 4 files changed, 73 insertions(+), 12 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 219918b..355bd09 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,4 +9,6 @@ assert_matches = "1.3.0"
 quickcheck = "0.8"
 
 [dependencies]
-derive_more = "0.15"
\ No newline at end of file
+derive_more = "0.15"
+bigint = "4"
+byteorder = "1"
\ No newline at end of file
diff --git a/src/lib.rs b/src/lib.rs
index 1511ea7..dbc5b91 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -5,6 +5,8 @@
 #[cfg(test)] #[macro_use] extern crate assert_matches;
 #[cfg(test)] #[macro_use] extern crate quickcheck;
 extern crate derive_more;
+extern crate bigint;
+extern crate byteorder;
 
 mod tree;
 mod node_data;
@@ -53,7 +55,7 @@ impl Entry {
         leaves & (leaves - 1) == 0
     }
 
-    pub fn leaf_count(&self) -> u32 {
+    pub fn leaf_count(&self) -> u64 {
         self.data.end_height - self.data.start_height + 1
     }
 
diff --git a/src/node_data.rs b/src/node_data.rs
index 8d81c71..ba7877f 100644
--- a/src/node_data.rs
+++ b/src/node_data.rs
@@ -1,3 +1,7 @@
+use byteorder::{LittleEndian, WriteBytesExt};
+
+use bigint::U256;
+
 /// Node metadata.
 #[repr(C)]
 #[derive(Debug)]
@@ -9,13 +13,15 @@ pub struct NodeData {
     pub end_target: u32,
     pub start_sapling_root: [u8; 32],
     pub end_sapling_root: [u8; 32],
-    pub subtree_total_work: u64,
-    pub start_height: u32,
-    pub end_height: u32,
+    pub subtree_total_work: U256,
+    pub start_height: u64,
+    pub end_height: u64,
     pub shielded_tx: u64,
 }
 
 impl NodeData {
+    pub const MAX_SERIALIZED_SIZE: usize = 32 + 4 + 4 + 4 + 4 + 32 + 32 + 32 + 9 + 9 + 9; // =171;
+
     pub fn combine(left: &NodeData, right: &NodeData) -> NodeData {
         NodeData {
             // TODO: hash children
@@ -28,10 +34,61 @@ impl NodeData {
             end_sapling_root: right.end_sapling_root,
 
             // TODO: sum work?
-            subtree_total_work: 0,
+            subtree_total_work: left.subtree_total_work + right.subtree_total_work,
             start_height: left.start_height,
             end_height: right.end_height,
             shielded_tx: left.shielded_tx + right.shielded_tx,
         }
     }
+
+    fn write_compact<W: std::io::Write>(w: &mut W, compact: u64) -> std::io::Result<()> {
+        match compact {
+            0..=0xfc => {
+                w.write_all(&[compact as u8])?
+            },
+            0xfd..=0xffff => {
+                w.write_all(&[0xfd])?;
+                w.write_u16::<LittleEndian>(compact as u16)?;
+            },
+            0x10000..=0xffff_ffff => {
+                w.write_all(&[0xfe])?;
+                w.write_u32::<LittleEndian>(compact as u32)?;
+            },
+            _ => {
+                w.write_all(&[0xff])?;
+                w.write_u64::<LittleEndian>(compact)?;
+            }
+        }
+        Ok(())
+    }
+
+    pub fn write<W: std::io::Write>(&self, w: &mut W) -> std::io::Result<()> {
+        w.write_all(&self.subtree_commitment)?;
+        w.write_u32::<LittleEndian>(self.start_time)?;
+        w.write_u32::<LittleEndian>(self.end_time)?;
+        w.write_u32::<LittleEndian>(self.start_target)?;
+        w.write_u32::<LittleEndian>(self.end_target)?;
+        w.write_all(&self.start_sapling_root)?;
+        w.write_all(&self.end_sapling_root)?;
+
+        let mut work_buf = [0u8; 32];
+        self.subtree_total_work.to_little_endian(&mut work_buf[..]);
+        w.write_all(&work_buf)?;
+
+        Self::write_compact(w, self.start_height)?;
+        Self::write_compact(w, self.end_height)?;
+        Self::write_compact(w, self.shielded_tx)?;
+        Ok(())
+    }
+
+    pub fn to_bytes(&self) -> Vec<u8> {
+        let mut buf = [0u8; Self::MAX_SERIALIZED_SIZE];
+        let pos = {
+            let mut cursor = std::io::Cursor::new(&mut buf[..]);
+            self.write(&mut cursor).expect("Cursor cannot fail");
+            cursor.position() as usize
+        };
+
+        buf[0..pos].to_vec()
+    }
 }
\ No newline at end of file
diff --git a/src/tree.rs b/src/tree.rs
index 9be1b5e..e2b1c5b 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -292,14 +292,14 @@ mod tests {
             end_target: 0,
             start_sapling_root: [0u8; 32],
             end_sapling_root: [0u8; 32],
-            subtree_total_work: 0,
-            start_height: height,
-            end_height: height,
+            subtree_total_work: 0.into(),
+            start_height: height as u64,
+            end_height: height as u64,
             shielded_tx: 7,
         }
     }
 
-    fn node(start_height: u32, end_height: u32) -> NodeData {
+    fn node(start_height: u64, end_height: u64) -> NodeData {
         NodeData {
             subtree_commitment: [0u8; 32],
             start_time: 0,
@@ -308,7 +308,7 @@ mod tests {
             end_target: 0,
             start_sapling_root: [0u8; 32],
             end_sapling_root: [0u8; 32],
-            subtree_total_work: 0,
+            subtree_total_work: 0.into(),
             start_height: start_height,
             end_height: end_height,
             shielded_tx: 7,
@@ -679,7 +679,7 @@ mod tests {
                         .node
                         .leaf_count()
                         ==
-                        number
+                        number as u64
                 )
             }
         }