Add HistoryTree struct (#2396)
* Add HistoryTree struct * Update zebra-chain/src/primitives/zcash_history.rs Co-authored-by: Deirdre Connolly <deirdre@zfnd.org> * fix formatting * Apply suggestions from code review Co-authored-by: Deirdre Connolly <deirdre@zfnd.org> * history_tree.rs: fixes from code review Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>
This commit is contained in:
parent
dd645e7e0c
commit
d430e951c2
|
@ -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<u32, Entry>,
|
||||
/// 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<Block>,
|
||||
sapling_root: &sapling::tree::Root,
|
||||
_orchard_root: Option<&orchard::tree::Root>,
|
||||
) -> Result<Self, io::Error> {
|
||||
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<Block>,
|
||||
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<Block>,
|
||||
&'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,
|
||||
}
|
||||
}
|
||||
}
|
|
@ -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;
|
||||
|
|
|
@ -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<u32, &Entry>,
|
||||
extra: &HashMap<u32, &Entry>,
|
||||
peaks: &BTreeMap<u32, Entry>,
|
||||
extra: &BTreeMap<u32, Entry>,
|
||||
) -> Result<Self, io::Error> {
|
||||
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<Block>,
|
||||
sapling_root: &sapling::tree::Root,
|
||||
) -> Result<Self, io::Error> {
|
||||
) -> 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<Block>,
|
||||
sapling_root: &sapling::tree::Root,
|
||||
) -> Result<Vec<NodeData>, zcash_history::Error> {
|
||||
) -> Result<Vec<Entry>, 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<Item = (Arc<Block>, sapling::tree::Root)>,
|
||||
) -> Result<Vec<NodeData>, zcash_history::Error> {
|
||||
) -> Result<Vec<Entry>, 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
|
||||
|
|
|
@ -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();
|
||||
|
|
Loading…
Reference in New Issue