diff --git a/zebra-chain/src/history_tree.rs b/zebra-chain/src/history_tree.rs new file mode 100644 index 00000000..b53fc6aa --- /dev/null +++ b/zebra-chain/src/history_tree.rs @@ -0,0 +1,246 @@ +//! History tree (Merkle mountain range) structure that contains information about +//! the block history as specified in ZIP-221. + +use std::{ + collections::{BTreeMap, HashSet}, + io, + sync::Arc, +}; + +use thiserror::Error; + +use crate::{ + block::{Block, ChainHistoryMmrRootHash, Height}, + orchard, + parameters::{Network, NetworkUpgrade}, + primitives::zcash_history::{Entry, Tree as InnerHistoryTree}, + sapling, +}; + +/// An error describing why a history tree operation failed. +#[derive(Debug, Error)] +#[non_exhaustive] +#[allow(missing_docs)] +pub enum HistoryTreeError { + #[error("zcash_history error: {inner:?}")] + #[non_exhaustive] + InnerError { inner: zcash_history::Error }, + + #[error("I/O error")] + IOError(#[from] io::Error), +} + +/// 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 { + network: Network, + network_upgrade: NetworkUpgrade, + /// Merkle mountain range tree from `zcash_history`. + /// This is a "runtime" structure used to add / remove nodes, and it's not + /// persistent. + inner: InnerHistoryTree, + /// The number of nodes in the tree. + size: u32, + /// The peaks of the tree, indexed by their position in the array representation + /// of the tree. This can be persisted to save the tree. + peaks: BTreeMap, + /// The height of the most recent block added to the tree. + current_height: Height, +} + +impl HistoryTree { + /// Create a new history tree with a single block. + pub fn from_block( + network: Network, + block: Arc, + sapling_root: &sapling::tree::Root, + _orchard_root: Option<&orchard::tree::Root>, + ) -> Result { + let height = block + .coinbase_height() + .expect("block must have coinbase height during contextual verification"); + let network_upgrade = NetworkUpgrade::current(network, height); + // TODO: handle Orchard root, see https://github.com/ZcashFoundation/zebra/issues/2283 + let (tree, entry) = InnerHistoryTree::new_from_block(network, block, sapling_root)?; + let mut peaks = BTreeMap::new(); + peaks.insert(0u32, entry); + Ok(HistoryTree { + network, + network_upgrade, + inner: tree, + size: 1, + peaks, + current_height: height, + }) + } + + /// Add block data to the tree. + /// + /// # Panics + /// + /// If the block height is not one more than the previously pushed block. + pub fn push( + &mut self, + block: Arc, + sapling_root: &sapling::tree::Root, + _orchard_root: Option<&orchard::tree::Root>, + ) -> Result<(), HistoryTreeError> { + // Check if the block has the expected height. + // librustzcash assumes the heights are correct and corrupts the tree if they are wrong, + // resulting in a confusing error, which we prevent here. + let height = block + .coinbase_height() + .expect("block must have coinbase height during contextual verification"); + if height - self.current_height != 1 { + panic!( + "added block with height {:?} but it must be {:?}+1", + height, self.current_height + ); + } + + // TODO: handle orchard root + let new_entries = self + .inner + .append_leaf(block, sapling_root) + .map_err(|e| HistoryTreeError::InnerError { inner: e })?; + for entry in new_entries { + // Not every entry is a peak; those will be trimmed later + self.peaks.insert(self.size, entry); + self.size += 1; + } + self.prune()?; + // TODO: implement network upgrade logic: drop previous history, start new history + self.current_height = height; + Ok(()) + } + + /// Extend the history tree with the given blocks. + pub fn try_extend< + 'a, + T: IntoIterator< + Item = ( + Arc, + &'a sapling::tree::Root, + Option<&'a orchard::tree::Root>, + ), + >, + >( + &mut self, + iter: T, + ) -> Result<(), HistoryTreeError> { + for (block, sapling_root, orchard_root) in iter { + self.push(block, sapling_root, orchard_root)?; + } + Ok(()) + } + + /// Prune tree, removing all non-peak entries. + fn prune(&mut self) -> Result<(), io::Error> { + // Go through all the peaks of the tree. + // This code is based on a librustzcash example: + // https://github.com/zcash/librustzcash/blob/02052526925fba9389f1428d6df254d4dec967e6/zcash_history/examples/long.rs + // The explanation of how it works is from zcashd: + // https://github.com/zcash/zcash/blob/0247c0c682d59184a717a6536edb0d18834be9a7/src/coins.cpp#L351 + + let mut peak_pos_set = HashSet::new(); + + // Assume the following example peak layout with 14 leaves, and 25 stored nodes in + // total (the "tree length"): + // + // P + // /\ + // / \ + // / \ \ + // / \ \ Altitude + // _A_ \ \ 3 + // _/ \_ B \ 2 + // / \ / \ / \ C 1 + // /\ /\ /\ /\ /\ /\ /\ 0 + // + // We start by determining the altitude of the highest peak (A). + let mut alt = (32 - ((self.size + 1) as u32).leading_zeros() - 1) - 1; + + // We determine the position of the highest peak (A) by pretending it is the right + // sibling in a tree, and its left-most leaf has position 0. Then the left sibling + // of (A) has position -1, and so we can "jump" to the peak's position by computing + // -1 + 2^(alt + 1) - 1. + let mut peak_pos = (1 << (alt + 1)) - 2; + + // Now that we have the position and altitude of the highest peak (A), we collect + // the remaining peaks (B, C). We navigate the peaks as if they were nodes in this + // Merkle tree (with additional imaginary nodes 1 and 2, that have positions beyond + // the MMR's length): + // + // / \ + // / \ + // / \ + // / \ + // A ==========> 1 + // / \ // \ + // _/ \_ B ==> 2 + // /\ /\ /\ // + // / \ / \ / \ C + // /\ /\ /\ /\ /\ /\ /\ + // + loop { + // If peak_pos is out of bounds of the tree, we compute the position of its left + // child, and drop down one level in the tree. + if peak_pos >= self.size { + // left child, -2^alt + peak_pos -= 1 << alt; + alt -= 1; + } + + // If the peak exists, we take it and then continue with its right sibling. + if peak_pos < self.size { + // There is a peak at index `peak_pos` + peak_pos_set.insert(peak_pos); + + // right sibling + peak_pos = peak_pos + (1 << (alt + 1)) - 1; + } + + if alt == 0 { + break; + } + } + + // Remove all non-peak entries + self.peaks.retain(|k, _| peak_pos_set.contains(k)); + // Rebuild tree + self.inner = InnerHistoryTree::new_from_cache( + self.network, + self.network_upgrade, + self.size, + &self.peaks, + &Default::default(), + )?; + Ok(()) + } + + /// Return the hash of the tree root. + pub fn hash(&self) -> ChainHistoryMmrRootHash { + self.inner.hash() + } +} + +impl Clone for HistoryTree { + fn clone(&self) -> Self { + let tree = InnerHistoryTree::new_from_cache( + self.network, + self.network_upgrade, + self.size, + &self.peaks, + &Default::default(), + ) + .expect("rebuilding an existing tree should always work"); + HistoryTree { + network: self.network, + network_upgrade: self.network_upgrade, + inner: tree, + size: self.size, + peaks: self.peaks.clone(), + current_height: self.current_height, + } + } +} diff --git a/zebra-chain/src/lib.rs b/zebra-chain/src/lib.rs index 4ac1e9f2..9d7f8f9e 100644 --- a/zebra-chain/src/lib.rs +++ b/zebra-chain/src/lib.rs @@ -23,6 +23,7 @@ extern crate bitflags; pub mod amount; pub mod block; pub mod fmt; +pub mod history_tree; pub mod orchard; pub mod parameters; pub mod primitives; diff --git a/zebra-chain/src/primitives/zcash_history.rs b/zebra-chain/src/primitives/zcash_history.rs index 1c96ab47..6956b58b 100644 --- a/zebra-chain/src/primitives/zcash_history.rs +++ b/zebra-chain/src/primitives/zcash_history.rs @@ -6,7 +6,7 @@ mod tests; -use std::{collections::HashMap, convert::TryInto, io, sync::Arc}; +use std::{collections::BTreeMap, convert::TryInto, io, sync::Arc}; use crate::{ block::{Block, ChainHistoryMmrRootHash}, @@ -43,7 +43,9 @@ 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)] pub struct Entry { inner: [u8; zcash_history::MAX_ENTRY_SIZE], } @@ -101,12 +103,12 @@ impl Tree { /// # Panics /// /// Will panic if `peaks` is empty. - fn new_from_cache( + pub fn new_from_cache( network: Network, network_upgrade: NetworkUpgrade, length: u32, - peaks: &HashMap, - extra: &HashMap, + peaks: &BTreeMap, + extra: &BTreeMap, ) -> Result { let branch_id = network_upgrade .branch_id() @@ -132,19 +134,24 @@ impl Tree { /// Create a single-node MMR tree for the given block. /// /// `sapling_root` is the root of the Sapling note commitment tree of the block. - fn new_from_block( + pub fn new_from_block( network: Network, block: Arc, sapling_root: &sapling::tree::Root, - ) -> Result { + ) -> Result<(Self, Entry), io::Error> { let height = block .coinbase_height() .expect("block must have coinbase height during contextual verification"); let network_upgrade = NetworkUpgrade::current(network, height); let entry0 = Entry::new_leaf(block, network, sapling_root); - let mut peaks = HashMap::new(); - peaks.insert(0u32, &entry0); - Tree::new_from_cache(network, network_upgrade, 1, &peaks, &HashMap::new()) + let mut peaks = BTreeMap::new(); + peaks.insert(0u32, entry0); + Ok(( + Tree::new_from_cache(network, network_upgrade, 1, &peaks, &BTreeMap::new())?, + peaks + .remove(&0u32) + .expect("must work since it was just added"), + )) } /// Append a new block to the tree, as a new leaf. @@ -157,18 +164,18 @@ impl Tree { /// /// Panics if the network upgrade of the given block is different from /// the network upgrade of the other blocks in the tree. - fn append_leaf( + pub fn append_leaf( &mut self, block: Arc, sapling_root: &sapling::tree::Root, - ) -> Result, zcash_history::Error> { + ) -> Result, zcash_history::Error> { let height = block .coinbase_height() .expect("block must have coinbase height during contextual verification"); let network_upgrade = NetworkUpgrade::current(self.network, height); if self.network_upgrade != network_upgrade { panic!( - "added block from network upgrade {:?} but MMR tree is restricted to {:?}", + "added block from network upgrade {:?} but history tree is restricted to {:?}", network_upgrade, self.network_upgrade ); } @@ -177,17 +184,17 @@ impl Tree { let appended = self.inner.append_leaf(node_data)?; let mut new_nodes = Vec::new(); - for entry in appended { - let mut node = NodeData { - inner: [0; zcash_history::MAX_NODE_DATA_SIZE], + for entry_link in appended { + let mut entry = Entry { + inner: [0; zcash_history::MAX_ENTRY_SIZE], }; self.inner - .resolve_link(entry) + .resolve_link(entry_link) .expect("entry was just generated so it must be valid") - .data() - .write(&mut &mut node.inner[..]) + .node() + .write(&mut &mut entry.inner[..]) .expect("buffer was created with enough capacity"); - new_nodes.push(node); + new_nodes.push(entry); } Ok(new_nodes) } @@ -196,7 +203,7 @@ impl Tree { fn append_leaf_iter( &mut self, vals: impl Iterator, sapling::tree::Root)>, - ) -> Result, zcash_history::Error> { + ) -> Result, zcash_history::Error> { let mut new_nodes = Vec::new(); for (block, root) in vals { new_nodes.append(&mut self.append_leaf(block, &root)?); @@ -212,7 +219,7 @@ impl Tree { } /// Return the root hash of the tree, i.e. `hashChainHistoryRoot`. - fn hash(&self) -> ChainHistoryMmrRootHash { + pub fn hash(&self) -> ChainHistoryMmrRootHash { // Both append_leaf() and truncate_leaf() leave a root node, so it should // always exist. self.inner diff --git a/zebra-chain/src/primitives/zcash_history/tests/vectors.rs b/zebra-chain/src/primitives/zcash_history/tests/vectors.rs index 88bf6b21..6fd92764 100644 --- a/zebra-chain/src/primitives/zcash_history/tests/vectors.rs +++ b/zebra-chain/src/primitives/zcash_history/tests/vectors.rs @@ -49,7 +49,7 @@ fn tree_for_network_upgrade(network: Network, network_upgrade: NetworkUpgrade) - // Build initial MMR tree with only Block 0 let sapling_root0 = sapling::tree::Root(**sapling_roots.get(&height).expect("test vector exists")); - let mut tree = Tree::new_from_block(network, block0, &sapling_root0)?; + let (mut tree, _) = Tree::new_from_block(network, block0, &sapling_root0)?; // Compute root hash of the MMR tree, which will be included in the next block let hash0 = tree.hash();