Refactor HistoryTree into NonEmptyHistoryTree and HistoryTree (#2582)

* Refactor HistoryTree into NonEmptyHistoryTree and HistoryTree

* HistoryTree: use Deref instead of AsRef; remove unneeded PartialEq

zebra-chain/src/history_tree.rs

@@ -6,6 +6,7 @@ mod tests;
 use std::{
     collections::{BTreeMap, HashSet},
     io,
+    ops::Deref,
     sync::Arc,
 };
 
@@ -33,6 +34,7 @@ pub enum HistoryTreeError {
 }
 
 /// The inner [Tree] in one of its supported versions.
+#[derive(Debug)]
 enum InnerHistoryTree {
     /// A pre-Orchard tree.
     PreOrchard(Tree<PreOrchard>),
@@ -41,8 +43,9 @@ enum InnerHistoryTree {
 }
 
 /// History tree (Merkle mountain range) structure that contains information about
-// the block history, as specified in [ZIP-221][https://zips.z.cash/zip-0221].
-pub struct HistoryTree {
+/// the block history, as specified in [ZIP-221][https://zips.z.cash/zip-0221].
+#[derive(Debug)]
+pub struct NonEmptyHistoryTree {
     network: Network,
     network_upgrade: NetworkUpgrade,
     /// Merkle mountain range tree from `zcash_history`.
@@ -58,7 +61,7 @@ pub struct HistoryTree {
     current_height: Height,
 }
 
-impl HistoryTree {
+impl NonEmptyHistoryTree {
     /// Recreate a [`HistoryTree`] from previously saved data.
     ///
     /// The parameters must come from the values of [HistoryTree::size],
@@ -68,7 +71,7 @@ impl HistoryTree {
         size: u32,
         peaks: BTreeMap<u32, Entry>,
         current_height: Height,
-    ) -> Result<Self, io::Error> {
+    ) -> Result<Self, HistoryTreeError> {
         let network_upgrade = NetworkUpgrade::current(network, current_height);
         let inner = match network_upgrade {
             NetworkUpgrade::Genesis
@@ -119,7 +122,7 @@ impl HistoryTree {
         block: Arc<Block>,
         sapling_root: &sapling::tree::Root,
         orchard_root: &orchard::tree::Root,
-    ) -> Result<Self, io::Error> {
+    ) -> Result<Self, HistoryTreeError> {
         let height = block
             .coinbase_height()
             .expect("block must have coinbase height during contextual verification");
@@ -153,7 +156,7 @@ impl HistoryTree {
         };
         let mut peaks = BTreeMap::new();
         peaks.insert(0u32, entry);
-        Ok(HistoryTree {
+        Ok(NonEmptyHistoryTree {
             network,
             network_upgrade,
             inner: tree,
@@ -359,7 +362,7 @@ impl HistoryTree {
     }
 }
 
-impl Clone for HistoryTree {
+impl Clone for NonEmptyHistoryTree {
     fn clone(&self) -> Self {
         let tree = match self.inner {
             InnerHistoryTree::PreOrchard(_) => InnerHistoryTree::PreOrchard(
@@ -383,7 +386,7 @@ impl Clone for HistoryTree {
                 .expect("rebuilding an existing tree should always work"),
             ),
         };
-        HistoryTree {
+        NonEmptyHistoryTree {
             network: self.network,
             network_upgrade: self.network_upgrade,
             inner: tree,
@@ -393,3 +396,68 @@ impl Clone for HistoryTree {
         }
     }
 }
+
+/// A History Tree that keeps track of its own creation in the Heartwood
+/// activation block, being empty beforehand.
+#[derive(Debug, Default, Clone)]
+pub struct HistoryTree(Option<NonEmptyHistoryTree>);
+
+impl HistoryTree {
+    /// Push a block to a maybe-existing HistoryTree, handling network upgrades.
+    ///
+    /// The tree is updated in-place. It is created when pushing the Heartwood
+    /// activation block.
+    pub fn push(
+        &mut self,
+        network: Network,
+        block: Arc<Block>,
+        sapling_root: sapling::tree::Root,
+        orchard_root: orchard::tree::Root,
+    ) -> Result<(), HistoryTreeError> {
+        let heartwood_height = NetworkUpgrade::Heartwood
+            .activation_height(network)
+            .expect("Heartwood height is known");
+        match block
+            .coinbase_height()
+            .expect("must have height")
+            .cmp(&heartwood_height)
+        {
+            std::cmp::Ordering::Less => {
+                assert!(
+                    self.0.is_none(),
+                    "history tree must not exist pre-Heartwood"
+                );
+            }
+            std::cmp::Ordering::Equal => {
+                let tree = Some(NonEmptyHistoryTree::from_block(
+                    network,
+                    block.clone(),
+                    &sapling_root,
+                    &orchard_root,
+                )?);
+                // Replace the current object with the new tree
+                *self = HistoryTree(tree);
+            }
+            std::cmp::Ordering::Greater => {
+                self.0
+                    .as_mut()
+                    .expect("history tree must exist Heartwood-onward")
+                    .push(block.clone(), &sapling_root, &orchard_root)?;
+            }
+        };
+        Ok(())
+    }
+}
+
+impl From<NonEmptyHistoryTree> for HistoryTree {
+    fn from(tree: NonEmptyHistoryTree) -> Self {
+        HistoryTree(Some(tree))
+    }
+}
+
+impl Deref for HistoryTree {
+    type Target = Option<NonEmptyHistoryTree>;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}

zebra-chain/src/history_tree/tests/vectors.rs

@@ -5,7 +5,7 @@ use crate::{
         Block,
         Commitment::{self, ChainHistoryActivationReserved},
     },
-    history_tree::HistoryTree,
+    history_tree::NonEmptyHistoryTree,
     parameters::{Network, NetworkUpgrade},
     sapling,
     serialization::ZcashDeserializeInto,
@@ -63,7 +63,7 @@ fn push_and_prune_for_network_upgrade(
     // Build initial history tree tree with only the first block
     let first_sapling_root =
         sapling::tree::Root(**sapling_roots.get(&height).expect("test vector exists"));
-    let mut tree = HistoryTree::from_block(
+    let mut tree = NonEmptyHistoryTree::from_block(
         network,
         first_block,
         &first_sapling_root,
@@ -140,8 +140,12 @@ fn upgrade_for_network_upgrade(network: Network, network_upgrade: NetworkUpgrade
     // network upgrade), so we won't be able to check if its root is correct.
     let sapling_root_prev =
         sapling::tree::Root(**sapling_roots.get(&height).expect("test vector exists"));
-    let mut tree =
-        HistoryTree::from_block(network, block_prev, &sapling_root_prev, &Default::default())?;
+    let mut tree = NonEmptyHistoryTree::from_block(
+        network,
+        block_prev,
+        &sapling_root_prev,
+        &Default::default(),
+    )?;
 
     assert_eq!(tree.size(), 1);
     assert_eq!(tree.peaks().len(), 1);

zebra-chain/src/primitives/zcash_history.rs

@@ -62,7 +62,7 @@ impl From<&zcash_history::NodeData> for NodeData {
 /// An encoded entry in the tree.
 ///
 /// Contains the node data and information about its position in the tree.
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct Entry {
     #[serde(with = "BigArray")]
     inner: [u8; zcash_history::MAX_ENTRY_SIZE],
@@ -235,6 +235,15 @@ impl<V: Version> Tree<V> {
     }
 }
 
+impl<V: zcash_history::Version> std::fmt::Debug for Tree<V> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Tree")
+            .field("network", &self.network)
+            .field("network_upgrade", &self.network_upgrade)
+            .finish()
+    }
+}
+
 impl Version for zcash_history::V1 {
     /// Convert a Block into a V1::NodeData used in the MMR tree.
     ///

zebra-state/src/service/finalized_state.rs

@@ -9,9 +9,9 @@ use std::{collections::HashMap, convert::TryInto, path::Path, sync::Arc};
 
 use zebra_chain::{
     block::{self, Block},
-    history_tree::HistoryTree,
+    history_tree::{HistoryTree, NonEmptyHistoryTree},
     orchard,
-    parameters::{Network, NetworkUpgrade, GENESIS_PREVIOUS_BLOCK_HASH},
+    parameters::{Network, GENESIS_PREVIOUS_BLOCK_HASH},
     sapling, sprout,
     transaction::{self, Transaction},
     transparent,
@@ -363,32 +363,12 @@ impl FinalizedState {
             let sapling_root = sapling_note_commitment_tree.root();
             let orchard_root = orchard_note_commitment_tree.root();
 
-            // Create the history tree if it's the Heartwood activation block.
-            let heartwood_height = NetworkUpgrade::Heartwood
-                .activation_height(self.network)
-                .expect("Heartwood height is known");
-            match height.cmp(&heartwood_height) {
-                std::cmp::Ordering::Less => assert!(
-                    history_tree.is_none(),
-                    "history tree must not exist pre-Heartwood"
-                ),
-                std::cmp::Ordering::Equal => {
-                    history_tree = Some(HistoryTree::from_block(
-                        self.network,
-                        block.clone(),
-                        &sapling_root,
-                        &orchard_root,
-                    )?);
-                }
-                std::cmp::Ordering::Greater => history_tree
-                    .as_mut()
-                    .expect("history tree must exist Heartwood-onward")
-                    .push(block.clone(), &sapling_root, &orchard_root)?,
-            }
+            history_tree.push(self.network, block.clone(), sapling_root, orchard_root)?;
 
             // Compute the new anchors and index them
-            batch.zs_insert(sapling_anchors, sapling_note_commitment_tree.root(), ());
-            batch.zs_insert(orchard_anchors, orchard_note_commitment_tree.root(), ());
+            // Note: if the root hasn't changed, we write the same value again.
+            batch.zs_insert(sapling_anchors, sapling_root, ());
+            batch.zs_insert(orchard_anchors, orchard_root, ());
 
             // Update the trees in state
             if let Some(h) = finalized_tip_height {
@@ -406,7 +386,7 @@ impl FinalizedState {
                 height,
                 orchard_note_commitment_tree,
             );
-            if let Some(history_tree) = history_tree {
+            if let Some(history_tree) = history_tree.as_ref() {
                 batch.zs_insert(history_tree_cf, height, history_tree);
             }
 
@@ -550,10 +530,20 @@ impl FinalizedState {
 
     /// Returns the ZIP-221 history tree of the finalized tip or `None`
     /// if it does not exist yet in the state (pre-Heartwood).
-    pub fn history_tree(&self) -> Option<HistoryTree> {
-        let height = self.finalized_tip_height()?;
-        let history_tree = self.db.cf_handle("history_tree").unwrap();
-        self.db.zs_get(history_tree, &height)
+    pub fn history_tree(&self) -> HistoryTree {
+        match self.finalized_tip_height() {
+            Some(height) => {
+                let history_tree_cf = self.db.cf_handle("history_tree").unwrap();
+                let history_tree: Option<NonEmptyHistoryTree> =
+                    self.db.zs_get(history_tree_cf, &height);
+                if let Some(non_empty_tree) = history_tree {
+                    HistoryTree::from(non_empty_tree)
+                } else {
+                    Default::default()
+                }
+            }
+            None => Default::default(),
+        }
     }
 
     /// If the database is `ephemeral`, delete it.

zebra-state/src/service/finalized_state/disk_format.rs

@@ -6,7 +6,7 @@ use zebra_chain::{
     amount::NonNegative,
     block,
     block::{Block, Height},
-    history_tree::HistoryTree,
+    history_tree::NonEmptyHistoryTree,
     orchard,
     parameters::Network,
     primitives::zcash_history,
@@ -321,7 +321,7 @@ struct HistoryTreeParts {
     current_height: Height,
 }
 
-impl IntoDisk for HistoryTree {
+impl IntoDisk for NonEmptyHistoryTree {
     type Bytes = Vec<u8>;
 
     fn as_bytes(&self) -> Self::Bytes {
@@ -337,15 +337,20 @@ impl IntoDisk for HistoryTree {
     }
 }
 
-impl FromDisk for HistoryTree {
+impl FromDisk for NonEmptyHistoryTree {
     fn from_bytes(bytes: impl AsRef<[u8]>) -> Self {
         let parts: HistoryTreeParts = bincode::DefaultOptions::new()
             .deserialize(bytes.as_ref())
             .expect(
                 "deserialization format should match the serialization format used by IntoDisk",
             );
-        HistoryTree::from_cache(parts.network, parts.size, parts.peaks, parts.current_height)
-            .expect("deserialization format should match the serialization format used by IntoDisk")
+        NonEmptyHistoryTree::from_cache(
+            parts.network,
+            parts.size,
+            parts.peaks,
+            parts.current_height,
+        )
+        .expect("deserialization format should match the serialization format used by IntoDisk")
     }
 }