Skip to content

Commit

Permalink
Replace Chain.sprout_note_commitment_tree with a lookup method
Browse files Browse the repository at this point in the history
  • Loading branch information
@teor2345
teor2345 committed Feb 7, 2023
1 parent 6c1fbe7 commit 98dfa87
Show file tree
Hide file tree
Showing 4 changed files with 81 additions and 55 deletions.
13 changes: 5 additions & 8 deletions zebra-state/src/service/non_finalized_state.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,6 @@ impl NonFinalizedState {

let parent_chain = self.parent_chain(
parent_hash,
finalized_state.sprout_note_commitment_tree(),
finalized_state.sapling_note_commitment_tree(),
finalized_state.orchard_note_commitment_tree(),
finalized_state.history_tree(),
Expand Down Expand Up @@ -192,6 +191,9 @@ impl NonFinalizedState {
) -> Result<(), ValidateContextError> {
let chain = Chain::new(
self.network,
finalized_state
.finalized_tip_height()
.expect("finalized state contains blocks"),
finalized_state.sprout_note_commitment_tree(),
finalized_state.sapling_note_commitment_tree(),
finalized_state.orchard_note_commitment_tree(),
Expand Down Expand Up @@ -456,16 +458,12 @@ impl NonFinalizedState {

/// Return the chain whose tip block hash is `parent_hash`.
///
/// The chain can be an existing chain in the non-finalized state or a freshly
/// created fork, if needed.
///
/// The trees must be the trees of the finalized tip.
/// They are used to recreate the trees if a fork is needed.
/// The chain can be an existing chain in the non-finalized state, or a freshly
/// created fork.
#[allow(clippy::unwrap_in_result)]
fn parent_chain(
&mut self,
parent_hash: block::Hash,
sprout_note_commitment_tree: Arc<sprout::tree::NoteCommitmentTree>,
sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
orchard_note_commitment_tree: Arc<orchard::tree::NoteCommitmentTree>,
history_tree: Arc<HistoryTree>,
Expand All @@ -484,7 +482,6 @@ impl NonFinalizedState {
chain
.fork(
parent_hash,
sprout_note_commitment_tree.clone(),
sapling_note_commitment_tree.clone(),
orchard_note_commitment_tree.clone(),
history_tree.clone(),
Expand Down
84 changes: 53 additions & 31 deletions zebra-state/src/service/non_finalized_state/chain.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,9 +73,6 @@ pub struct Chain {

// Note commitment trees
//
/// The Sprout note commitment tree of the tip of this [`Chain`],
/// including all finalized notes, and the non-finalized notes in this chain.
pub(super) sprout_note_commitment_tree: Arc<sprout::tree::NoteCommitmentTree>,
/// The Sprout note commitment tree for each anchor.
/// This is required for interstitial states.
pub(crate) sprout_trees_by_anchor:
Expand Down Expand Up @@ -158,22 +155,22 @@ pub struct Chain {
}

impl Chain {
/// Create a new Chain with the given trees and network.
/// Create a new Chain with the given finalized tip trees and network.
pub(crate) fn new(
network: Network,
finalized_tip_height: Height,
sprout_note_commitment_tree: Arc<sprout::tree::NoteCommitmentTree>,
sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
orchard_note_commitment_tree: Arc<orchard::tree::NoteCommitmentTree>,
history_tree: Arc<HistoryTree>,
finalized_tip_chain_value_pools: ValueBalance<NonNegative>,
) -> Self {
Self {
let mut chain = Self {
network,
blocks: Default::default(),
height_by_hash: Default::default(),
tx_loc_by_hash: Default::default(),
created_utxos: Default::default(),
sprout_note_commitment_tree,
sapling_note_commitment_tree,
orchard_note_commitment_tree,
spent_utxos: Default::default(),
Expand All @@ -195,7 +192,11 @@ impl Chain {
history_tree,
history_trees_by_height: Default::default(),
chain_value_pools: finalized_tip_chain_value_pools,
}
};

chain.add_sprout_tree_and_anchor(finalized_tip_height, sprout_note_commitment_tree);

chain
}

/// Is the internal state of `self` the same as `other`?
Expand All @@ -220,7 +221,6 @@ impl Chain {
self.spent_utxos == other.spent_utxos &&

// note commitment trees
self.sprout_note_commitment_tree.root() == other.sprout_note_commitment_tree.root() &&
self.sprout_trees_by_anchor == other.sprout_trees_by_anchor &&
self.sprout_trees_by_height == other.sprout_trees_by_height &&
self.sapling_note_commitment_tree.root() == other.sapling_note_commitment_tree.root() &&
Expand Down Expand Up @@ -305,17 +305,14 @@ impl Chain {
.expect("only called while blocks is populated")
}

/// Fork a chain at the block with the given hash, if it is part of this
/// chain.
/// Fork and return a chain at the block with the given `fork_tip`, if it is part of this
/// chain. Otherwise, if this chain does not contain `fork_tip`, returns `Ok(None)`.
///
/// The passed trees must match the trees of the finalized tip. They are
/// extended by the commitments from the newly forked chain up to the passed
/// `fork_tip`.
/// If forking the chain fails, returns `Err(_)`.
#[allow(clippy::unwrap_in_result)]
pub fn fork(
&self,
fork_tip: block::Hash,
sprout_note_commitment_tree: Arc<sprout::tree::NoteCommitmentTree>,
sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
orchard_note_commitment_tree: Arc<orchard::tree::NoteCommitmentTree>,
history_tree: Arc<HistoryTree>,
Expand All @@ -325,7 +322,6 @@ impl Chain {
}

let mut forked = self.with_trees(
sprout_note_commitment_tree,
sapling_note_commitment_tree,
orchard_note_commitment_tree,
history_tree,
Expand Down Expand Up @@ -373,7 +369,9 @@ impl Chain {

// TODO: use NoteCommitmentTrees to store the trees as well?
let mut nct = NoteCommitmentTrees {
sprout: self.sprout_note_commitment_tree.clone(),
sprout: self
.sprout_tree(fork_height.into())
.expect("already checked chain contains fork height"),
sapling: self.sapling_note_commitment_tree.clone(),
orchard: self.orchard_note_commitment_tree.clone(),
};
Expand All @@ -398,8 +396,7 @@ impl Chain {
note_result.expect("scope has already finished")?;
history_result.expect("scope has already finished")?;

// Update the note commitment trees in the chain.
self.sprout_note_commitment_tree = nct.sprout;
// Update the note commitment trees and anchors in the chain.
self.sapling_note_commitment_tree = nct.sapling;
self.orchard_note_commitment_tree = nct.orchard;

Expand Down Expand Up @@ -535,6 +532,17 @@ impl Chain {
)
}

/// Returns the Sprout note commitment tree of the tip of this [`Chain`],
/// including all finalized notes, and the non-finalized notes in this chain.
///
/// # Panics
///
/// If this chain is empty.
pub fn sprout_note_commitment_tree(&self) -> Arc<sprout::tree::NoteCommitmentTree> {
self.sprout_tree(self.non_finalized_tip_height().into())
.expect("only called while sprout_trees_by_height is populated")
}

/// Returns the Sprout
/// [`NoteCommitmentTree`](sprout::tree::NoteCommitmentTree) specified by a
/// [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
Expand All @@ -548,6 +556,29 @@ impl Chain {
self.sprout_trees_by_height.get(&height).cloned()
}

/// Update the Sprout `tree` and anchor indexes at `height`.
pub(crate) fn add_sprout_tree_and_anchor(
&mut self,
height: Height,
tree: Arc<sprout::tree::NoteCommitmentTree>,
) {
// TODO:
// - should we assert that the height was empty before insertion?
// (the roots can have duplicates, but the heights should not)
// - should this be an UpdateWith impl?
self.sprout_trees_by_height.insert(height, tree.clone());

// Having updated all the note commitment trees and nullifier sets in
// this block, the roots of the note commitment trees as of the last
// transaction are the anchor treestates of this block.
//
// Use the previously cached root which was calculated in parallel.
let sprout_root = self.sprout_note_commitment_tree().root();
self.sprout_anchors.insert(sprout_root);
self.sprout_anchors_by_height.insert(height, sprout_root);
self.sprout_trees_by_anchor.insert(sprout_root, tree);
}

/// Returns the Sapling
/// [`NoteCommitmentTree`](sapling::tree::NoteCommitmentTree) specified by a
/// [`HashOrHeight`], if it exists in the non-finalized [`Chain`].
Expand Down Expand Up @@ -821,13 +852,12 @@ impl Chain {

// Cloning

/// Clone the Chain but not the history and note commitment trees, using
/// the specified trees instead.
/// Clone the Chain but not the history and note commitment trees,
/// starting with the specified trees at `fork_height` instead.
///
/// Useful when forking, where the trees are rebuilt anyway.
fn with_trees(
&self,
sprout_note_commitment_tree: Arc<sprout::tree::NoteCommitmentTree>,
sapling_note_commitment_tree: Arc<sapling::tree::NoteCommitmentTree>,
orchard_note_commitment_tree: Arc<orchard::tree::NoteCommitmentTree>,
history_tree: Arc<HistoryTree>,
Expand All @@ -839,7 +869,6 @@ impl Chain {
tx_loc_by_hash: self.tx_loc_by_hash.clone(),
created_utxos: self.created_utxos.clone(),
spent_utxos: self.spent_utxos.clone(),
sprout_note_commitment_tree,
sprout_trees_by_anchor: self.sprout_trees_by_anchor.clone(),
sprout_trees_by_height: self.sprout_trees_by_height.clone(),
sapling_note_commitment_tree,
Expand Down Expand Up @@ -879,7 +908,7 @@ impl Chain {
//
// TODO: use NoteCommitmentTrees to store the trees as well?
let mut nct = NoteCommitmentTrees {
sprout: self.sprout_note_commitment_tree.clone(),
sprout: self.sprout_note_commitment_tree(),
sapling: self.sapling_note_commitment_tree.clone(),
orchard: self.orchard_note_commitment_tree.clone(),
};
Expand All @@ -904,15 +933,13 @@ impl Chain {
partial_result.expect("scope has already finished")?;

// Update the note commitment trees in the chain.
self.sprout_note_commitment_tree = nct.sprout;
self.add_sprout_tree_and_anchor(height, nct.sprout);
self.sapling_note_commitment_tree = nct.sapling;
self.orchard_note_commitment_tree = nct.orchard;

// Do the Chain updates with data dependencies on note commitment tree updates

// Update the note commitment trees indexed by height.
self.sprout_trees_by_height
.insert(height, self.sprout_note_commitment_tree.clone());
self.sapling_trees_by_height
.insert(height, self.sapling_note_commitment_tree.clone());
self.orchard_trees_by_height
Expand All @@ -923,11 +950,6 @@ impl Chain {
// transaction are the treestates of this block.
//
// Use the previously cached roots, which were calculated in parallel.
let sprout_root = self.sprout_note_commitment_tree.root();
self.sprout_anchors.insert(sprout_root);
self.sprout_anchors_by_height.insert(height, sprout_root);
self.sprout_trees_by_anchor
.insert(sprout_root, self.sprout_note_commitment_tree.clone());
let sapling_root = self.sapling_note_commitment_tree.root();
self.sapling_anchors.insert(sapling_root);
self.sapling_anchors_by_height.insert(height, sapling_root);
Expand Down
32 changes: 17 additions & 15 deletions zebra-state/src/service/non_finalized_state/tests/prop.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use zebra_test::prelude::*;

use zebra_chain::{
amount::NonNegative,
block::{self, arbitrary::allow_all_transparent_coinbase_spends, Block},
block::{self, arbitrary::allow_all_transparent_coinbase_spends, Block, Height},
fmt::DisplayToDebug,
history_tree::{HistoryTree, NonEmptyHistoryTree},
parameters::NetworkUpgrade::*,
Expand Down Expand Up @@ -47,7 +47,7 @@ fn push_genesis_chain() -> Result<()> {
|((chain, count, network, empty_tree) in PreparedChain::default())| {
prop_assert!(empty_tree.is_none());

let mut only_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), empty_tree, ValueBalance::zero());
let mut only_chain = Chain::new(network, Height(0), Default::default(), Default::default(), Default::default(), empty_tree, ValueBalance::zero());
// contains the block value pool changes and chain value pool balances for each height
let mut chain_values = BTreeMap::new();

Expand Down Expand Up @@ -99,7 +99,7 @@ fn push_history_tree_chain() -> Result<()> {
let count = std::cmp::min(count, chain.len() - 1);
let chain = &chain[1..];

let mut only_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), finalized_tree, ValueBalance::zero());
let mut only_chain = Chain::new(network, Height(0), Default::default(), Default::default(), Default::default(), finalized_tree, ValueBalance::zero());

for block in chain
.iter()
Expand Down Expand Up @@ -143,6 +143,7 @@ fn forked_equals_pushed_genesis() -> Result<()> {
// correspond to the blocks in the original chain before the fork.
let mut partial_chain = Chain::new(
network,
Height(0),
Default::default(),
Default::default(),
Default::default(),
Expand All @@ -162,6 +163,7 @@ fn forked_equals_pushed_genesis() -> Result<()> {
// This chain will be forked.
let mut full_chain = Chain::new(
network,
Height(0),
Default::default(),
Default::default(),
Default::default(),
Expand Down Expand Up @@ -195,15 +197,15 @@ fn forked_equals_pushed_genesis() -> Result<()> {
}

// Use [`fork_at_count`] as the fork tip.
let fork_tip_hash = chain[fork_at_count - 1].hash;
let fork_tip_height = fork_at_count - 1;
let fork_tip_hash = chain[fork_tip_height].hash;

// Fork the chain.
let mut forked = full_chain
.fork(
fork_tip_hash,
Default::default(),
Default::default(),
Default::default(),
empty_tree,
)
.expect("fork works")
Expand Down Expand Up @@ -254,8 +256,8 @@ fn forked_equals_pushed_history_tree() -> Result<()> {
// use `fork_at_count` as the fork tip
let fork_tip_hash = chain[fork_at_count - 1].hash;

let mut full_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), finalized_tree.clone(), ValueBalance::zero());
let mut partial_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), finalized_tree.clone(), ValueBalance::zero());
let mut full_chain = Chain::new(network, Height(0), Default::default(), Default::default(), Default::default(), finalized_tree.clone(), ValueBalance::zero());
let mut partial_chain = Chain::new(network, Height(0), Default::default(), Default::default(), Default::default(), finalized_tree.clone(), ValueBalance::zero());

for block in chain
.iter()
Expand All @@ -275,7 +277,6 @@ fn forked_equals_pushed_history_tree() -> Result<()> {
fork_tip_hash,
Default::default(),
Default::default(),
Default::default(),
finalized_tree,
)
.expect("fork works")
Expand Down Expand Up @@ -323,7 +324,7 @@ fn finalized_equals_pushed_genesis() -> Result<()> {

let fake_value_pool = ValueBalance::<NonNegative>::fake_populated_pool();

let mut full_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), empty_tree, fake_value_pool);
let mut full_chain = Chain::new(network, Height(0), Default::default(), Default::default(), Default::default(), empty_tree, fake_value_pool);
for block in chain
.iter()
.take(finalized_count)
Expand All @@ -333,7 +334,8 @@ fn finalized_equals_pushed_genesis() -> Result<()> {

let mut partial_chain = Chain::new(
network,
full_chain.sprout_note_commitment_tree.clone(),
Height(finalized_count.try_into().unwrap()),
full_chain.sprout_note_commitment_tree(),
full_chain.sapling_note_commitment_tree.clone(),
full_chain.orchard_note_commitment_tree.clone(),
full_chain.history_tree.clone(),
Expand Down Expand Up @@ -393,7 +395,7 @@ fn finalized_equals_pushed_history_tree() -> Result<()> {

let fake_value_pool = ValueBalance::<NonNegative>::fake_populated_pool();

let mut full_chain = Chain::new(network, Default::default(), Default::default(), Default::default(), finalized_tree, fake_value_pool);
let mut full_chain = Chain::new(network, Height(0), Default::default(), Default::default(), Default::default(), finalized_tree, fake_value_pool);
for block in chain
.iter()
.take(finalized_count)
Expand All @@ -403,7 +405,8 @@ fn finalized_equals_pushed_history_tree() -> Result<()> {

let mut partial_chain = Chain::new(
network,
full_chain.sprout_note_commitment_tree.clone(),
Height(finalized_count.try_into().unwrap()),
full_chain.sprout_note_commitment_tree(),
full_chain.sapling_note_commitment_tree.clone(),
full_chain.orchard_note_commitment_tree.clone(),
full_chain.history_tree.clone(),
Expand Down Expand Up @@ -570,8 +573,8 @@ fn different_blocks_different_chains() -> Result<()> {
Default::default()
};

let chain1 = Chain::new(Network::Mainnet, Default::default(), Default::default(), Default::default(), finalized_tree1, ValueBalance::fake_populated_pool());
let chain2 = Chain::new(Network::Mainnet, Default::default(), Default::default(), Default::default(), finalized_tree2, ValueBalance::fake_populated_pool());
let chain1 = Chain::new(Network::Mainnet, Height(0), Default::default(), Default::default(), Default::default(), finalized_tree1, ValueBalance::fake_populated_pool());
let chain2 = Chain::new(Network::Mainnet, Height(0), Default::default(), Default::default(), Default::default(), finalized_tree2, ValueBalance::fake_populated_pool());

let block1 = vec1[1].clone().prepare().test_with_zero_spent_utxos();
let block2 = vec2[1].clone().prepare().test_with_zero_spent_utxos();
Expand Down Expand Up @@ -606,7 +609,6 @@ fn different_blocks_different_chains() -> Result<()> {
chain1.spent_utxos = chain2.spent_utxos.clone();

// note commitment trees
chain1.sprout_note_commitment_tree = chain2.sprout_note_commitment_tree.clone();
chain1.sprout_trees_by_anchor = chain2.sprout_trees_by_anchor.clone();
chain1.sprout_trees_by_height = chain2.sprout_trees_by_height.clone();
chain1.sapling_note_commitment_tree = chain2.sapling_note_commitment_tree.clone();
Expand Down
Loading

0 comments on commit 98dfa87

Please sign in to comment.