zcash_history: Implement v2 history tree with Orchard support

zcash_history/CHANGELOG.md

@@ -6,6 +6,18 @@ and this library adheres to Rust's notion of
 [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Added
+- Support for multiple history tree versions:
+  - `zcash_history::Version` trait.
+  - `zcash_history::V1`, marking the original history tree version.
+  - `zcash_history::V2`, marking the history tree version from NU5.
+- `zcash_history::Entry::new_leaf`
+
+### Changed
+- `zcash_history::{Entry, IndexedNode, Tree}` now have a `Version` parameter.
+
+### Removed
+- `impl From<NodeData> for Entry` (replaced by `Entry::new_leaf`).
 
 ## [0.2.0] - 2020-03-13
 No changes, just a version bump.

zcash_history/examples/lib/shared.rs

@@ -1,8 +1,8 @@
-use zcash_history::{Entry, EntryLink, NodeData, Tree};
+use zcash_history::{Entry, EntryLink, NodeData, Tree, V1};
 
 pub struct NodeDataIterator {
     return_stack: Vec<NodeData>,
-    tree: Tree,
+    tree: Tree<V1>,
     cursor: usize,
     leaf_cursor: usize,
 }
@@ -56,7 +56,7 @@ impl NodeDataIterator {
         let tree = Tree::new(
             3,
             vec![(2, root)],
-            vec![(0, leaf(1).into()), (1, leaf(2).into())],
+            vec![(0, Entry::new_leaf(leaf(1))), (1, Entry::new_leaf(leaf(2)))],
         );
 
         NodeDataIterator {

zcash_history/examples/long.rs

@@ -1,12 +1,12 @@
-use zcash_history::{Entry, EntryLink, NodeData, Tree};
+use zcash_history::{Entry, EntryLink, NodeData, Tree, V1};
 
 #[path = "lib/shared.rs"]
 mod share;
 
-fn draft(into: &mut Vec<(u32, Entry)>, vec: &[NodeData], peak_pos: usize, h: u32) {
+fn draft(into: &mut Vec<(u32, Entry<V1>)>, vec: &[NodeData], peak_pos: usize, h: u32) {
     let node_data = vec[peak_pos - 1].clone();
-    let peak: Entry = match h {
-        0 => node_data.into(),
+    let peak = match h {
+        0 => Entry::new_leaf(node_data),
         _ => Entry::new(
             node_data,
             EntryLink::Stored((peak_pos - (1 << h) - 1) as u32),
@@ -19,7 +19,7 @@ fn draft(into: &mut Vec<(u32, Entry)>, vec: &[NodeData], peak_pos: usize, h: u32
     into.push(((peak_pos - 1) as u32, peak));
 }
 
-fn prepare_tree(vec: &[NodeData]) -> Tree {
+fn prepare_tree(vec: &[NodeData]) -> Tree<V1> {
     assert!(!vec.is_empty());
 
     // integer log2 of (vec.len()+1), -1

zcash_history/src/entry.rs

@@ -1,26 +1,34 @@
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 
-use crate::{EntryKind, EntryLink, Error, NodeData, MAX_NODE_DATA_SIZE};
+use crate::{EntryKind, EntryLink, Error, Version, MAX_NODE_DATA_SIZE};
 
 /// Max serialized length of entry data.
 pub const MAX_ENTRY_SIZE: usize = MAX_NODE_DATA_SIZE + 9;
 
 /// MMR Entry.
 #[derive(Debug)]
-pub struct Entry {
+pub struct Entry<V: Version> {
     pub(crate) kind: EntryKind,
-    pub(crate) data: NodeData,
+    pub(crate) data: V::NodeData,
 }
 
-impl Entry {
+impl<V: Version> Entry<V> {
     /// New entry of type node.
-    pub fn new(data: NodeData, left: EntryLink, right: EntryLink) -> Self {
+    pub fn new(data: V::NodeData, left: EntryLink, right: EntryLink) -> Self {
         Entry {
             kind: EntryKind::Node(left, right),
             data,
         }
     }
 
+    /// Creates a new leaf.
+    pub fn new_leaf(data: V::NodeData) -> Self {
+        Entry {
+            kind: EntryKind::Leaf,
+            data,
+        }
+    }
+
     /// Returns if is this node complete (has total of 2^N leaves)
     pub fn complete(&self) -> bool {
         let leaves = self.leaf_count();
@@ -29,7 +37,7 @@ impl Entry {
 
     /// Number of leaves under this node.
     pub fn leaf_count(&self) -> u64 {
-        self.data.end_height - (self.data.start_height - 1)
+        V::end_height(&self.data) - (V::start_height(&self.data) - 1)
     }
 
     /// Is this node a leaf.
@@ -67,7 +75,7 @@ impl Entry {
             }
         };
 
-        let data = NodeData::read(consensus_branch_id, r)?;
+        let data = V::read(consensus_branch_id, r)?;
 
         Ok(Entry { kind, data })
     }
@@ -88,7 +96,7 @@ impl Entry {
             }
         }
 
-        self.data.write(w)?;
+        V::write(&self.data, w)?;
 
         Ok(())
     }
@@ -100,16 +108,7 @@ impl Entry {
     }
 }
 
-impl From<NodeData> for Entry {
-    fn from(s: NodeData) -> Self {
-        Entry {
-            kind: EntryKind::Leaf,
-            data: s,
-        }
-    }
-}
-
-impl std::fmt::Display for Entry {
+impl<V: Version> std::fmt::Display for Entry<V> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self.kind {
             EntryKind::Node(l, r) => write!(f, "node({}, {}, ..)", l, r),

zcash_history/src/lib.rs

@@ -9,10 +9,12 @@
 mod entry;
 mod node_data;
 mod tree;
+mod version;
 
 pub use entry::{Entry, MAX_ENTRY_SIZE};
 pub use node_data::{NodeData, MAX_NODE_DATA_SIZE};
 pub use tree::Tree;
+pub use version::{Version, V1, V2};
 
 /// Crate-level error type
 #[derive(Debug)]

zcash_history/src/node_data.rs

@@ -1,6 +1,7 @@
 use bigint::U256;
-use blake2::Params as Blake2Params;
-use byteorder::{ByteOrder, LittleEndian, ReadBytesExt, WriteBytesExt};
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+
+use crate::Version;
 
 /// Maximum serialized size of the node metadata.
 pub const MAX_NODE_DATA_SIZE: usize = 32 + // subtree commitment
@@ -13,10 +14,13 @@ pub const MAX_NODE_DATA_SIZE: usize = 32 + // subtree commitment
     32 + // subtree total work
     9 +  // start height (compact uint)
     9 +  // end height (compact uint)
-    9; // Sapling tx count (compact uint)
-       // = total of 171
+    9 + // Sapling tx count (compact uint)
+    32 + // start Orchard tree root
+    32 + // end Orchard tree root
+    9; // Orchard tx count (compact uint)
+       // = total of 244
 
-/// Node metadata.
+/// V1 node metadata.
 #[repr(C)]
 #[derive(Debug, Clone, Default)]
 #[cfg_attr(test, derive(PartialEq))]
@@ -47,49 +51,20 @@ pub struct NodeData {
     pub sapling_tx: u64,
 }
 
-fn blake2b_personal(personalization: &[u8], input: &[u8]) -> [u8; 32] {
-    let hash_result = Blake2Params::new()
-        .hash_length(32)
-        .personal(personalization)
-        .to_state()
-        .update(input)
-        .finalize();
-    let mut result = [0u8; 32];
-    result.copy_from_slice(hash_result.as_bytes());
-    result
-}
-
-fn personalization(branch_id: u32) -> [u8; 16] {
-    let mut result = [0u8; 16];
-    result[..12].copy_from_slice(b"ZcashHistory");
-    LittleEndian::write_u32(&mut result[12..], branch_id);
-    result
-}
-
 impl NodeData {
     /// Combine two nodes metadata.
     pub fn combine(left: &NodeData, right: &NodeData) -> NodeData {
-        assert_eq!(left.consensus_branch_id, right.consensus_branch_id);
-
-        let mut hash_buf = [0u8; MAX_NODE_DATA_SIZE * 2];
-        let size = {
-            let mut cursor = ::std::io::Cursor::new(&mut hash_buf[..]);
-            left.write(&mut cursor)
-                .expect("Writing to memory buf with enough length cannot fail; qed");
-            right
-                .write(&mut cursor)
-                .expect("Writing to memory buf with enough length cannot fail; qed");
-            cursor.position() as usize
-        };
-
-        let hash = blake2b_personal(
-            &personalization(left.consensus_branch_id),
-            &hash_buf[..size],
-        );
+        crate::V1::combine(left, right)
+    }
 
+    pub(crate) fn combine_inner(
+        subtree_commitment: [u8; 32],
+        left: &NodeData,
+        right: &NodeData,
+    ) -> NodeData {
         NodeData {
             consensus_branch_id: left.consensus_branch_id,
-            subtree_commitment: hash,
+            subtree_commitment,
             start_time: left.start_time,
             end_time: right.end_time,
             start_target: left.start_target,
@@ -180,27 +155,64 @@ impl NodeData {
 
     /// Convert to byte representation.
     pub fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = [0u8; MAX_NODE_DATA_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()
+        crate::V1::to_bytes(self)
     }
 
     /// Convert from byte representation.
     pub fn from_bytes<T: AsRef<[u8]>>(consensus_branch_id: u32, buf: T) -> std::io::Result<Self> {
-        let mut cursor = std::io::Cursor::new(buf);
-        Self::read(consensus_branch_id, &mut cursor)
+        crate::V1::from_bytes(consensus_branch_id, buf)
     }
 
     /// Hash node metadata
     pub fn hash(&self) -> [u8; 32] {
-        let bytes = self.to_bytes();
+        crate::V1::hash(self)
+    }
+}
+
+/// V2 node metadata.
+#[derive(Debug, Clone, Default)]
+#[cfg_attr(test, derive(PartialEq))]
+pub struct V2 {
+    /// The V1 node data retained in V2.
+    pub v1: NodeData,
+    /// Start Orchard tree root.
+    pub start_orchard_root: [u8; 32],
+    /// End Orchard tree root.
+    pub end_orchard_root: [u8; 32],
+    /// Number of Orchard transactions.
+    pub orchard_tx: u64,
+}
+
+impl V2 {
+    pub(crate) fn combine_inner(subtree_commitment: [u8; 32], left: &V2, right: &V2) -> V2 {
+        V2 {
+            v1: NodeData::combine_inner(subtree_commitment, &left.v1, &right.v1),
+            start_orchard_root: left.start_orchard_root,
+            end_orchard_root: right.end_orchard_root,
+            orchard_tx: left.orchard_tx + right.orchard_tx,
+        }
+    }
 
-        blake2b_personal(&personalization(self.consensus_branch_id), &bytes)
+    /// Write to the byte representation.
+    pub fn write<W: std::io::Write>(&self, w: &mut W) -> std::io::Result<()> {
+        self.v1.write(w)?;
+        w.write_all(&self.start_orchard_root)?;
+        w.write_all(&self.end_orchard_root)?;
+        NodeData::write_compact(w, self.orchard_tx)?;
+        Ok(())
+    }
+
+    /// Read from the byte representation.
+    pub fn read<R: std::io::Read>(consensus_branch_id: u32, r: &mut R) -> std::io::Result<Self> {
+        let mut data = V2 {
+            v1: NodeData::read(consensus_branch_id, r)?,
+            ..Default::default()
+        };
+        r.read_exact(&mut data.start_orchard_root)?;
+        r.read_exact(&mut data.end_orchard_root)?;
+        data.orchard_tx = NodeData::read_compact(r)?;
+
+        Ok(data)
     }
 }