mirror of
https://github.com/0glabs/0g-storage-node.git
synced 2025-01-26 15:05:17 +00:00
f4d5228234
Some checks are pending
abi-consistent-check / build-and-compare (push) Waiting to run
code-coverage / unittest-cov (push) Waiting to run
rust / check (push) Waiting to run
rust / test (push) Waiting to run
rust / lints (push) Waiting to run
functional-test / test (push) Waiting to run
* add detailed metrics for storage layer
909 lines
33 KiB
Rust
909 lines
33 KiB
Rust
mod merkle_tree;
|
|
mod metrics;
|
|
mod node_manager;
|
|
mod proof;
|
|
mod sha3;
|
|
|
|
use anyhow::{anyhow, bail, Result};
|
|
use itertools::Itertools;
|
|
use std::cmp::Ordering;
|
|
use std::collections::{BTreeMap, HashMap};
|
|
use std::fmt::Debug;
|
|
use std::marker::PhantomData;
|
|
use std::sync::Arc;
|
|
use std::time::Instant;
|
|
use tracing::{trace, warn};
|
|
|
|
use crate::merkle_tree::MerkleTreeWrite;
|
|
pub use crate::merkle_tree::{
|
|
Algorithm, HashElement, MerkleTreeInitialData, MerkleTreeRead, ZERO_HASHES,
|
|
};
|
|
pub use crate::node_manager::{EmptyNodeDatabase, NodeDatabase, NodeManager, NodeTransaction};
|
|
pub use proof::{Proof, RangeProof};
|
|
pub use sha3::Sha3Algorithm;
|
|
|
|
pub struct AppendMerkleTree<E: HashElement, A: Algorithm<E>> {
|
|
/// Keep all the nodes in the latest version. `layers[0]` is the layer of leaves.
|
|
node_manager: NodeManager<E>,
|
|
/// Keep the delta nodes that can be used to construct a history tree.
|
|
/// The key is the root node of that version.
|
|
delta_nodes_map: BTreeMap<u64, DeltaNodes<E>>,
|
|
root_to_tx_seq_map: HashMap<E, u64>,
|
|
|
|
/// For `last_chunk_merkle` after the first chunk, this is set to `Some(10)` so that
|
|
/// `revert_to` can reset the state correctly when needed.
|
|
min_depth: Option<usize>,
|
|
/// Used to compute the correct padding hash.
|
|
/// 0 for `pora_chunk_merkle` and 10 for not-first `last_chunk_merkle`.
|
|
leaf_height: usize,
|
|
_a: PhantomData<A>,
|
|
}
|
|
|
|
impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
|
|
pub fn new(leaves: Vec<E>, leaf_height: usize, start_tx_seq: Option<u64>) -> Self {
|
|
let mut merkle = Self {
|
|
node_manager: NodeManager::new_dummy(),
|
|
delta_nodes_map: BTreeMap::new(),
|
|
root_to_tx_seq_map: HashMap::new(),
|
|
min_depth: None,
|
|
leaf_height,
|
|
_a: Default::default(),
|
|
};
|
|
merkle.node_manager.start_transaction();
|
|
merkle.node_manager.add_layer();
|
|
merkle.node_manager.append_nodes(0, &leaves);
|
|
if merkle.leaves() == 0 {
|
|
if let Some(seq) = start_tx_seq {
|
|
merkle.delta_nodes_map.insert(
|
|
seq,
|
|
DeltaNodes {
|
|
right_most_nodes: vec![],
|
|
},
|
|
);
|
|
}
|
|
merkle.node_manager.commit();
|
|
return merkle;
|
|
}
|
|
// Reconstruct the whole tree.
|
|
merkle.recompute(0, 0, None);
|
|
merkle.node_manager.commit();
|
|
// Commit the first version in memory.
|
|
// TODO(zz): Check when the roots become available.
|
|
merkle.commit(start_tx_seq);
|
|
merkle
|
|
}
|
|
|
|
pub fn new_with_subtrees(
|
|
node_db: Arc<dyn NodeDatabase<E>>,
|
|
node_cache_capacity: usize,
|
|
leaf_height: usize,
|
|
) -> Result<Self> {
|
|
let mut merkle = Self {
|
|
node_manager: NodeManager::new(node_db, node_cache_capacity)?,
|
|
delta_nodes_map: BTreeMap::new(),
|
|
root_to_tx_seq_map: HashMap::new(),
|
|
min_depth: None,
|
|
leaf_height,
|
|
_a: Default::default(),
|
|
};
|
|
if merkle.height() == 0 {
|
|
merkle.node_manager.start_transaction();
|
|
merkle.node_manager.add_layer();
|
|
merkle.node_manager.commit();
|
|
}
|
|
Ok(merkle)
|
|
}
|
|
|
|
/// This is only used for the last chunk, so `leaf_height` is always 0 so far.
|
|
pub fn new_with_depth(leaves: Vec<E>, depth: usize, start_tx_seq: Option<u64>) -> Self {
|
|
let mut node_manager = NodeManager::new_dummy();
|
|
node_manager.start_transaction();
|
|
if leaves.is_empty() {
|
|
// Create an empty merkle tree with `depth`.
|
|
let mut merkle = Self {
|
|
// dummy node manager for the last chunk.
|
|
node_manager,
|
|
delta_nodes_map: BTreeMap::new(),
|
|
root_to_tx_seq_map: HashMap::new(),
|
|
min_depth: Some(depth),
|
|
leaf_height: 0,
|
|
_a: Default::default(),
|
|
};
|
|
for _ in 0..depth {
|
|
merkle.node_manager.add_layer();
|
|
}
|
|
if let Some(seq) = start_tx_seq {
|
|
merkle.delta_nodes_map.insert(
|
|
seq,
|
|
DeltaNodes {
|
|
right_most_nodes: vec![],
|
|
},
|
|
);
|
|
}
|
|
merkle.node_manager.commit();
|
|
merkle
|
|
} else {
|
|
let mut merkle = Self {
|
|
// dummy node manager for the last chunk.
|
|
node_manager,
|
|
delta_nodes_map: BTreeMap::new(),
|
|
root_to_tx_seq_map: HashMap::new(),
|
|
min_depth: Some(depth),
|
|
leaf_height: 0,
|
|
_a: Default::default(),
|
|
};
|
|
merkle.node_manager.add_layer();
|
|
merkle.append_nodes(0, &leaves);
|
|
for _ in 1..depth {
|
|
merkle.node_manager.add_layer();
|
|
}
|
|
// Reconstruct the whole tree.
|
|
merkle.recompute(0, 0, None);
|
|
merkle.node_manager.commit();
|
|
// Commit the first version in memory.
|
|
merkle.commit(start_tx_seq);
|
|
merkle
|
|
}
|
|
}
|
|
|
|
pub fn append(&mut self, new_leaf: E) {
|
|
let start_time = Instant::now();
|
|
if new_leaf == E::null() {
|
|
// appending null is not allowed.
|
|
return;
|
|
}
|
|
self.node_manager.start_transaction();
|
|
self.node_manager.push_node(0, new_leaf);
|
|
self.recompute_after_append_leaves(self.leaves() - 1);
|
|
|
|
self.node_manager.commit();
|
|
metrics::APPEND.update_since(start_time);
|
|
}
|
|
|
|
pub fn append_list(&mut self, leaf_list: Vec<E>) {
|
|
let start_time = Instant::now();
|
|
if leaf_list.contains(&E::null()) {
|
|
// appending null is not allowed.
|
|
return;
|
|
}
|
|
self.node_manager.start_transaction();
|
|
let start_index = self.leaves();
|
|
self.node_manager.append_nodes(0, &leaf_list);
|
|
self.recompute_after_append_leaves(start_index);
|
|
self.node_manager.commit();
|
|
metrics::APPEND_LIST.update_since(start_time);
|
|
}
|
|
|
|
/// Append a leaf list by providing their intermediate node hash.
|
|
/// The appended subtree must be aligned. And it's up to the caller to
|
|
/// append the padding nodes for alignment.
|
|
/// Other nodes in the subtree will be set to `null` nodes.
|
|
/// TODO: Optimize to avoid storing the `null` nodes?
|
|
pub fn append_subtree(&mut self, subtree_depth: usize, subtree_root: E) -> Result<()> {
|
|
let start_time = Instant::now();
|
|
if subtree_root == E::null() {
|
|
// appending null is not allowed.
|
|
bail!("subtree_root is null");
|
|
}
|
|
self.node_manager.start_transaction();
|
|
let start_index = self.leaves();
|
|
self.append_subtree_inner(subtree_depth, subtree_root)?;
|
|
self.recompute_after_append_subtree(start_index, subtree_depth - 1);
|
|
self.node_manager.commit();
|
|
metrics::APPEND_SUBTREE.update_since(start_time);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub fn append_subtree_list(&mut self, subtree_list: Vec<(usize, E)>) -> Result<()> {
|
|
let start_time = Instant::now();
|
|
if subtree_list.iter().any(|(_, root)| root == &E::null()) {
|
|
// appending null is not allowed.
|
|
bail!("subtree_list contains null");
|
|
}
|
|
self.node_manager.start_transaction();
|
|
for (subtree_depth, subtree_root) in subtree_list {
|
|
let start_index = self.leaves();
|
|
self.append_subtree_inner(subtree_depth, subtree_root)?;
|
|
self.recompute_after_append_subtree(start_index, subtree_depth - 1);
|
|
}
|
|
self.node_manager.commit();
|
|
metrics::APPEND_SUBTREE_LIST.update_since(start_time);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Change the value of the last leaf and return the new merkle root.
|
|
/// This is needed if our merkle-tree in memory only keeps intermediate nodes instead of real leaves.
|
|
pub fn update_last(&mut self, updated_leaf: E) {
|
|
let start_time = Instant::now();
|
|
if updated_leaf == E::null() {
|
|
// updating to null is not allowed.
|
|
return;
|
|
}
|
|
self.node_manager.start_transaction();
|
|
if self.layer_len(0) == 0 {
|
|
// Special case for the first data.
|
|
self.push_node(0, updated_leaf);
|
|
} else {
|
|
self.update_node(0, self.layer_len(0) - 1, updated_leaf);
|
|
}
|
|
self.recompute_after_append_leaves(self.leaves() - 1);
|
|
self.node_manager.commit();
|
|
metrics::UPDATE_LAST.update_since(start_time);
|
|
}
|
|
|
|
/// Fill an unknown `null` leaf with its real value.
|
|
/// Panics if the leaf is already set and different or the index is out of range.
|
|
/// TODO: Batch computing intermediate nodes.
|
|
pub fn fill_leaf(&mut self, index: usize, leaf: E) {
|
|
if leaf == E::null() {
|
|
// fill leaf with null is not allowed.
|
|
} else if self.node(0, index) == E::null() {
|
|
self.node_manager.start_transaction();
|
|
self.update_node(0, index, leaf);
|
|
self.recompute_after_fill_leaves(index, index + 1);
|
|
self.node_manager.commit();
|
|
} else if self.node(0, index) != leaf {
|
|
panic!(
|
|
"Fill with invalid leaf, index={} was={:?} get={:?}",
|
|
index,
|
|
self.node(0, index),
|
|
leaf
|
|
);
|
|
}
|
|
}
|
|
|
|
/// Fill nodes with a valid proof data.
|
|
/// This requires that the proof is built against this tree.
|
|
/// This should only be called after validating the proof (including checking root existence).
|
|
/// Returns `Error` if the data is conflict with existing ones.
|
|
pub fn fill_with_range_proof(
|
|
&mut self,
|
|
proof: RangeProof<E>,
|
|
) -> Result<Vec<(usize, usize, E)>> {
|
|
self.node_manager.start_transaction();
|
|
let mut updated_nodes = Vec::new();
|
|
let mut left_nodes = proof.left_proof.proof_nodes_in_tree();
|
|
if left_nodes.len() >= self.leaf_height {
|
|
updated_nodes
|
|
.append(&mut self.fill_with_proof(left_nodes.split_off(self.leaf_height))?);
|
|
}
|
|
let mut right_nodes = proof.right_proof.proof_nodes_in_tree();
|
|
if right_nodes.len() >= self.leaf_height {
|
|
updated_nodes
|
|
.append(&mut self.fill_with_proof(right_nodes.split_off(self.leaf_height))?);
|
|
}
|
|
self.node_manager.commit();
|
|
Ok(updated_nodes)
|
|
}
|
|
|
|
pub fn fill_with_file_proof(
|
|
&mut self,
|
|
proof: Proof<E>,
|
|
mut tx_merkle_nodes: Vec<(usize, E)>,
|
|
start_index: u64,
|
|
) -> Result<Vec<(usize, usize, E)>> {
|
|
let tx_merkle_nodes_size = tx_merkle_nodes.len();
|
|
if self.leaf_height != 0 {
|
|
tx_merkle_nodes = tx_merkle_nodes
|
|
.into_iter()
|
|
.filter_map(|(height, data)| {
|
|
if height > self.leaf_height {
|
|
Some((height - self.leaf_height - 1, data))
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect();
|
|
}
|
|
if tx_merkle_nodes.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
self.node_manager.start_transaction();
|
|
let mut position_and_data =
|
|
proof.file_proof_nodes_in_tree(tx_merkle_nodes, tx_merkle_nodes_size);
|
|
let start_index = (start_index >> self.leaf_height) as usize;
|
|
for (i, (position, _)) in position_and_data.iter_mut().enumerate() {
|
|
*position += start_index >> i;
|
|
}
|
|
let updated_nodes = self.fill_with_proof(position_and_data)?;
|
|
self.node_manager.commit();
|
|
Ok(updated_nodes)
|
|
}
|
|
|
|
/// This assumes that the proof leaf is no lower than the tree leaf. It holds for both SegmentProof and ChunkProof.
|
|
/// Return the inserted nodes and position.
|
|
fn fill_with_proof(
|
|
&mut self,
|
|
position_and_data: Vec<(usize, E)>,
|
|
) -> Result<Vec<(usize, usize, E)>> {
|
|
let mut updated_nodes = Vec::new();
|
|
// A valid proof should not fail the following checks.
|
|
for (i, (position, data)) in position_and_data.into_iter().enumerate() {
|
|
if position > self.layer_len(i) {
|
|
bail!(
|
|
"proof position out of range, position={} layer.len()={}",
|
|
position,
|
|
self.layer_len(i)
|
|
);
|
|
}
|
|
if position == self.layer_len(i) {
|
|
// skip padding node.
|
|
continue;
|
|
}
|
|
if self.node(i, position) == E::null() {
|
|
self.update_node(i, position, data.clone());
|
|
updated_nodes.push((i, position, data))
|
|
} else if self.node(i, position) != data {
|
|
// The last node in each layer may have changed in the tree.
|
|
trace!(
|
|
"conflict data layer={} position={} tree_data={:?} proof_data={:?}",
|
|
i,
|
|
position,
|
|
self.node(i, position),
|
|
data
|
|
);
|
|
}
|
|
}
|
|
Ok(updated_nodes)
|
|
}
|
|
|
|
pub fn check_root(&self, root: &E) -> bool {
|
|
self.root_to_tx_seq_map.contains_key(root)
|
|
}
|
|
|
|
pub fn leaf_at(&self, position: usize) -> Result<Option<E>> {
|
|
if position >= self.leaves() {
|
|
bail!("Out of bound: position={} end={}", position, self.leaves());
|
|
}
|
|
if self.node(0, position) != E::null() {
|
|
Ok(Some(self.node(0, position)))
|
|
} else {
|
|
// The leaf hash is unknown.
|
|
Ok(None)
|
|
}
|
|
}
|
|
|
|
/// Return a list of subtrees that can be used to rebuild the tree.
|
|
pub fn get_subtrees(&self) -> Vec<(usize, E)> {
|
|
let mut next_index = 0;
|
|
let mut subtree_list: Vec<(usize, E)> = Vec::new();
|
|
while next_index < self.leaves() {
|
|
let root_tuple = self.first_known_root_at(next_index);
|
|
let subtree_size = 1 << (root_tuple.0 - 1);
|
|
let root_start_index = next_index / subtree_size * subtree_size;
|
|
|
|
// Previous subtrees are included within the new subtree.
|
|
// Pop them out and replace with the new one.
|
|
if root_start_index < next_index {
|
|
while let Some(last) = subtree_list.pop() {
|
|
next_index -= 1 << (last.0 - 1);
|
|
if next_index == root_start_index {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
next_index += subtree_size;
|
|
subtree_list.push(root_tuple);
|
|
}
|
|
subtree_list
|
|
}
|
|
}
|
|
|
|
impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
|
|
pub fn commit(&mut self, tx_seq: Option<u64>) {
|
|
if let Some(tx_seq) = tx_seq {
|
|
if self.leaves() == 0 {
|
|
// The state is empty, so we just save the root as `null`.
|
|
// Note that this root should not be used.
|
|
self.delta_nodes_map.insert(
|
|
tx_seq,
|
|
DeltaNodes {
|
|
right_most_nodes: vec![],
|
|
},
|
|
);
|
|
return;
|
|
}
|
|
let mut right_most_nodes = Vec::new();
|
|
for height in 0..self.height() {
|
|
let pos = self.layer_len(height) - 1;
|
|
right_most_nodes.push((pos, self.node(height, pos)));
|
|
}
|
|
let root = self.root();
|
|
self.delta_nodes_map
|
|
.insert(tx_seq, DeltaNodes::new(right_most_nodes));
|
|
self.root_to_tx_seq_map.insert(root, tx_seq);
|
|
}
|
|
}
|
|
|
|
fn before_extend_layer(&mut self, height: usize) {
|
|
if height == self.height() {
|
|
self.node_manager.add_layer()
|
|
}
|
|
}
|
|
|
|
fn recompute_after_append_leaves(&mut self, start_index: usize) {
|
|
self.recompute(start_index, 0, None)
|
|
}
|
|
|
|
fn recompute_after_append_subtree(&mut self, start_index: usize, height: usize) {
|
|
self.recompute(start_index, height, None)
|
|
}
|
|
|
|
fn recompute_after_fill_leaves(&mut self, start_index: usize, end_index: usize) {
|
|
self.recompute(start_index, 0, Some(end_index))
|
|
}
|
|
|
|
/// Given a range of changed leaf nodes and recompute the tree.
|
|
fn recompute(
|
|
&mut self,
|
|
mut start_index: usize,
|
|
mut height: usize,
|
|
mut maybe_end_index: Option<usize>,
|
|
) {
|
|
start_index >>= height;
|
|
maybe_end_index = maybe_end_index.map(|end| end >> height);
|
|
// Loop until we compute the new root and reach `tree_depth`.
|
|
while self.layer_len(height) > 1 || height < self.height() - 1 {
|
|
let next_layer_start_index = start_index >> 1;
|
|
if start_index % 2 == 1 {
|
|
start_index -= 1;
|
|
}
|
|
|
|
let mut end_index = maybe_end_index.unwrap_or(self.layer_len(height));
|
|
if end_index % 2 == 1 && end_index != self.layer_len(height) {
|
|
end_index += 1;
|
|
}
|
|
let mut i = 0;
|
|
let iter = self
|
|
.node_manager
|
|
.get_nodes(height, start_index, end_index)
|
|
.chunks(2);
|
|
// We cannot modify the parent layer while iterating the child layer,
|
|
// so just keep the changes and update them later.
|
|
let mut parent_update = Vec::new();
|
|
for chunk_iter in &iter {
|
|
let chunk: Vec<_> = chunk_iter.collect();
|
|
if chunk.len() == 2 {
|
|
let left = &chunk[0];
|
|
let right = &chunk[1];
|
|
// If either left or right is null (unknown), we cannot compute the parent hash.
|
|
// Note that if we are recompute a range of an existing tree,
|
|
// we do not need to keep these possibly null parent. This is only saved
|
|
// for the case of constructing a new tree from the leaves.
|
|
let parent = if *left == E::null() || *right == E::null() {
|
|
E::null()
|
|
} else {
|
|
A::parent(left, right)
|
|
};
|
|
parent_update.push((next_layer_start_index + i, parent));
|
|
i += 1;
|
|
} else {
|
|
assert_eq!(chunk.len(), 1);
|
|
let r = &chunk[0];
|
|
// Same as above.
|
|
let parent = if *r == E::null() {
|
|
E::null()
|
|
} else {
|
|
A::parent_single(r, height + self.leaf_height)
|
|
};
|
|
parent_update.push((next_layer_start_index + i, parent));
|
|
}
|
|
}
|
|
if !parent_update.is_empty() {
|
|
self.before_extend_layer(height + 1);
|
|
}
|
|
// `parent_update` is in increasing order by `parent_index`, so
|
|
// we can just overwrite `last_changed_parent_index` with new values.
|
|
let mut last_changed_parent_index = None;
|
|
for (parent_index, parent) in parent_update {
|
|
match parent_index.cmp(&self.layer_len(height + 1)) {
|
|
Ordering::Less => {
|
|
// We do not overwrite with null.
|
|
if parent != E::null() {
|
|
if self.node(height + 1, parent_index) == E::null()
|
|
// The last node in a layer can be updated.
|
|
|| (self.node(height + 1, parent_index) != parent
|
|
&& parent_index == self.layer_len(height + 1) - 1)
|
|
{
|
|
self.update_node(height + 1, parent_index, parent);
|
|
last_changed_parent_index = Some(parent_index);
|
|
} else if self.node(height + 1, parent_index) != parent {
|
|
// Recompute changes a node in the middle. This should be impossible
|
|
// if the inputs are valid.
|
|
panic!("Invalid append merkle tree! height={} index={} expected={:?} get={:?}",
|
|
height + 1, parent_index, self.node(height + 1, parent_index), parent);
|
|
}
|
|
}
|
|
}
|
|
Ordering::Equal => {
|
|
self.push_node(height + 1, parent);
|
|
last_changed_parent_index = Some(parent_index);
|
|
}
|
|
Ordering::Greater => {
|
|
unreachable!("depth={}, parent_index={}", height, parent_index);
|
|
}
|
|
}
|
|
}
|
|
if last_changed_parent_index.is_none() {
|
|
break;
|
|
}
|
|
maybe_end_index = last_changed_parent_index.map(|i| i + 1);
|
|
height += 1;
|
|
start_index = next_layer_start_index;
|
|
}
|
|
}
|
|
|
|
fn append_subtree_inner(&mut self, subtree_depth: usize, subtree_root: E) -> Result<()> {
|
|
if subtree_depth == 0 {
|
|
bail!("Subtree depth should not be zero!");
|
|
}
|
|
if self.leaves() % (1 << (subtree_depth - 1)) != 0 {
|
|
warn!(
|
|
"The current leaves count is not aligned with the merged subtree, \
|
|
this is only possible during recovery, leaves={}",
|
|
self.leaves()
|
|
);
|
|
}
|
|
for height in 0..(subtree_depth - 1) {
|
|
self.before_extend_layer(height);
|
|
let subtree_layer_size = 1 << (subtree_depth - 1 - height);
|
|
self.append_nodes(height, &vec![E::null(); subtree_layer_size]);
|
|
}
|
|
self.before_extend_layer(subtree_depth - 1);
|
|
self.push_node(subtree_depth - 1, subtree_root);
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub fn validate(&self, proof: &Proof<E>, leaf: &E, position: usize) -> Result<bool> {
|
|
proof.validate::<A>(leaf, position)?;
|
|
Ok(self.root_to_tx_seq_map.contains_key(&proof.root()))
|
|
}
|
|
|
|
pub fn revert_to(&mut self, tx_seq: u64) -> Result<()> {
|
|
if self.layer_len(0) == 0 {
|
|
// Any previous state of an empty tree is always empty.
|
|
return Ok(());
|
|
}
|
|
self.node_manager.start_transaction();
|
|
let delta_nodes = self
|
|
.delta_nodes_map
|
|
.get(&tx_seq)
|
|
.ok_or_else(|| anyhow!("tx_seq unavailable, root={:?}", tx_seq))?
|
|
.clone();
|
|
// Dropping the upper layers that are not in the old merkle tree.
|
|
for height in (delta_nodes.right_most_nodes.len()..self.height()).rev() {
|
|
self.node_manager.truncate_layer(height);
|
|
}
|
|
for (height, (last_index, right_most_node)) in
|
|
delta_nodes.right_most_nodes.iter().enumerate()
|
|
{
|
|
self.node_manager.truncate_nodes(height, *last_index + 1);
|
|
self.update_node(height, *last_index, right_most_node.clone())
|
|
}
|
|
self.clear_after(tx_seq);
|
|
self.node_manager.commit();
|
|
Ok(())
|
|
}
|
|
|
|
// Revert to a tx_seq not in `delta_nodes_map`.
|
|
// This is needed to revert the last unfinished tx after restart.
|
|
pub fn revert_to_leaves(&mut self, leaves: usize) -> Result<()> {
|
|
self.node_manager.start_transaction();
|
|
for height in (0..self.height()).rev() {
|
|
let kept_nodes = leaves >> height;
|
|
if kept_nodes == 0 {
|
|
self.node_manager.truncate_layer(height);
|
|
} else {
|
|
self.node_manager.truncate_nodes(height, kept_nodes + 1);
|
|
}
|
|
}
|
|
self.recompute_after_append_leaves(leaves);
|
|
self.node_manager.commit();
|
|
Ok(())
|
|
}
|
|
|
|
pub fn tx_seq_at_root(&self, root_hash: &E) -> Result<u64> {
|
|
self.root_to_tx_seq_map
|
|
.get(root_hash)
|
|
.cloned()
|
|
.ok_or_else(|| anyhow!("old root unavailable, root={:?}", root_hash))
|
|
}
|
|
|
|
pub fn at_version(&self, tx_seq: u64) -> Result<HistoryTree<E>> {
|
|
let delta_nodes = self
|
|
.delta_nodes_map
|
|
.get(&tx_seq)
|
|
.ok_or_else(|| anyhow!("tx_seq unavailable, tx_seq={:?}", tx_seq))?;
|
|
if delta_nodes.height() == 0 {
|
|
bail!("empty tree");
|
|
}
|
|
Ok(HistoryTree {
|
|
node_manager: &self.node_manager,
|
|
delta_nodes,
|
|
leaf_height: self.leaf_height,
|
|
})
|
|
}
|
|
|
|
pub fn reset(&mut self) {
|
|
self.node_manager.start_transaction();
|
|
for height in (0..self.height()).rev() {
|
|
self.node_manager.truncate_layer(height);
|
|
}
|
|
if let Some(depth) = self.min_depth {
|
|
for _ in 0..depth {
|
|
self.node_manager.add_layer();
|
|
}
|
|
} else {
|
|
self.node_manager.add_layer();
|
|
}
|
|
self.node_manager.commit();
|
|
}
|
|
|
|
fn clear_after(&mut self, tx_seq: u64) {
|
|
let mut tx_seq = tx_seq + 1;
|
|
while self.delta_nodes_map.contains_key(&tx_seq) {
|
|
if let Some(nodes) = self.delta_nodes_map.remove(&tx_seq) {
|
|
if nodes.height() != 0 {
|
|
self.root_to_tx_seq_map.remove(nodes.root());
|
|
}
|
|
}
|
|
tx_seq += 1;
|
|
}
|
|
}
|
|
|
|
/// Return the height and the root hash of the first available node from the leaf to the root.
|
|
/// The caller should ensure that `index` is within range.
|
|
fn first_known_root_at(&self, index: usize) -> (usize, E) {
|
|
let mut height = 0;
|
|
let mut index_in_layer = index;
|
|
while height < self.height() {
|
|
let node = self.node(height, index_in_layer);
|
|
if !node.is_null() {
|
|
return (height + 1, node);
|
|
}
|
|
height += 1;
|
|
index_in_layer /= 2;
|
|
}
|
|
unreachable!("root is always available")
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
struct DeltaNodes<E: HashElement> {
|
|
/// The right most nodes in a layer and its position.
|
|
right_most_nodes: Vec<(usize, E)>,
|
|
}
|
|
|
|
impl<E: HashElement> DeltaNodes<E> {
|
|
fn new(right_most_nodes: Vec<(usize, E)>) -> Self {
|
|
Self { right_most_nodes }
|
|
}
|
|
|
|
fn get(&self, height: usize, position: usize) -> Result<Option<&E>> {
|
|
if height >= self.right_most_nodes.len() || position > self.right_most_nodes[height].0 {
|
|
Err(anyhow!("position out of tree range"))
|
|
} else if position == self.right_most_nodes[height].0 {
|
|
Ok(Some(&self.right_most_nodes[height].1))
|
|
} else {
|
|
Ok(None)
|
|
}
|
|
}
|
|
|
|
fn layer_len(&self, height: usize) -> usize {
|
|
self.right_most_nodes[height].0 + 1
|
|
}
|
|
|
|
fn height(&self) -> usize {
|
|
self.right_most_nodes.len()
|
|
}
|
|
|
|
fn root(&self) -> &E {
|
|
&self.right_most_nodes.last().unwrap().1
|
|
}
|
|
}
|
|
|
|
pub struct HistoryTree<'m, E: HashElement> {
|
|
/// A reference to the global tree nodes.
|
|
node_manager: &'m NodeManager<E>,
|
|
/// The delta nodes that are difference from `layers`.
|
|
/// This could be a reference, we just take ownership for convenience.
|
|
delta_nodes: &'m DeltaNodes<E>,
|
|
|
|
leaf_height: usize,
|
|
}
|
|
|
|
impl<E: HashElement, A: Algorithm<E>> MerkleTreeRead for AppendMerkleTree<E, A> {
|
|
type E = E;
|
|
|
|
fn node(&self, layer: usize, index: usize) -> Self::E {
|
|
self.node_manager
|
|
.get_node(layer, index)
|
|
.expect("index checked")
|
|
}
|
|
|
|
fn height(&self) -> usize {
|
|
self.node_manager.num_layers()
|
|
}
|
|
|
|
fn layer_len(&self, layer_height: usize) -> usize {
|
|
self.node_manager.layer_size(layer_height)
|
|
}
|
|
|
|
fn padding_node(&self, height: usize) -> Self::E {
|
|
E::end_pad(height + self.leaf_height)
|
|
}
|
|
}
|
|
|
|
impl<'a, E: HashElement> MerkleTreeRead for HistoryTree<'a, E> {
|
|
type E = E;
|
|
fn node(&self, layer: usize, index: usize) -> Self::E {
|
|
match self.delta_nodes.get(layer, index).expect("range checked") {
|
|
Some(node) if *node != E::null() => node.clone(),
|
|
_ => self
|
|
.node_manager
|
|
.get_node(layer, index)
|
|
.expect("index checked"),
|
|
}
|
|
}
|
|
|
|
fn height(&self) -> usize {
|
|
self.delta_nodes.height()
|
|
}
|
|
|
|
fn layer_len(&self, layer_height: usize) -> usize {
|
|
self.delta_nodes.layer_len(layer_height)
|
|
}
|
|
|
|
fn padding_node(&self, height: usize) -> Self::E {
|
|
E::end_pad(height + self.leaf_height)
|
|
}
|
|
}
|
|
|
|
impl<E: HashElement, A: Algorithm<E>> MerkleTreeWrite for AppendMerkleTree<E, A> {
|
|
type E = E;
|
|
|
|
fn push_node(&mut self, layer: usize, node: Self::E) {
|
|
self.node_manager.push_node(layer, node);
|
|
}
|
|
|
|
fn append_nodes(&mut self, layer: usize, nodes: &[Self::E]) {
|
|
self.node_manager.append_nodes(layer, nodes);
|
|
}
|
|
|
|
fn update_node(&mut self, layer: usize, pos: usize, node: Self::E) {
|
|
self.node_manager.add_node(layer, pos, node);
|
|
}
|
|
}
|
|
|
|
#[macro_export]
|
|
macro_rules! ensure_eq {
|
|
($given:expr, $expected:expr) => {
|
|
ensure!(
|
|
$given == $expected,
|
|
format!(
|
|
"equal check fails! {}:{}: {}={:?}, {}={:?}",
|
|
file!(),
|
|
line!(),
|
|
stringify!($given),
|
|
$given,
|
|
stringify!($expected),
|
|
$expected,
|
|
)
|
|
);
|
|
};
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use crate::merkle_tree::MerkleTreeRead;
|
|
|
|
use crate::sha3::Sha3Algorithm;
|
|
use crate::AppendMerkleTree;
|
|
use ethereum_types::H256;
|
|
|
|
#[test]
|
|
fn test_proof() {
|
|
let n = [1, 2, 6, 1025];
|
|
for entry_len in n {
|
|
let mut data = Vec::new();
|
|
for _ in 0..entry_len {
|
|
data.push(H256::random());
|
|
}
|
|
let mut merkle =
|
|
AppendMerkleTree::<H256, Sha3Algorithm>::new(vec![H256::zero()], 0, None);
|
|
merkle.append_list(data.clone());
|
|
merkle.commit(Some(0));
|
|
verify(&data, &mut merkle);
|
|
|
|
data.push(H256::random());
|
|
merkle.append(*data.last().unwrap());
|
|
merkle.commit(Some(1));
|
|
verify(&data, &mut merkle);
|
|
|
|
for _ in 0..6 {
|
|
data.push(H256::random());
|
|
}
|
|
merkle.append_list(data[data.len() - 6..].to_vec());
|
|
merkle.commit(Some(2));
|
|
verify(&data, &mut merkle);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_proof_against_modified_merkle() {
|
|
let n = [1, 2, 6, 1025];
|
|
for entry_len in n {
|
|
let mut data = Vec::new();
|
|
for _ in 0..entry_len {
|
|
data.push(H256::random());
|
|
}
|
|
let mut merkle =
|
|
AppendMerkleTree::<H256, Sha3Algorithm>::new(vec![H256::zero()], 0, None);
|
|
merkle.append_list(data.clone());
|
|
merkle.commit(Some(0));
|
|
|
|
for i in (0..data.len()).step_by(6) {
|
|
let end = std::cmp::min(i + 3, data.len());
|
|
let range_proof = merkle.gen_range_proof(i + 1, end + 1).unwrap();
|
|
let mut new_data = Vec::new();
|
|
for _ in 0..3 {
|
|
new_data.push(H256::random());
|
|
}
|
|
merkle.append_list(new_data);
|
|
let seq = i as u64 / 6 + 1;
|
|
merkle.commit(Some(seq));
|
|
let r = range_proof.validate::<Sha3Algorithm>(&data[i..end], i + 1);
|
|
assert!(r.is_ok(), "{:?}", r);
|
|
merkle.fill_with_range_proof(range_proof).unwrap();
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_proof_at_version() {
|
|
let n = [2, 255, 256, 257];
|
|
let mut merkle = AppendMerkleTree::<H256, Sha3Algorithm>::new(vec![H256::zero()], 0, None);
|
|
let mut start_pos = 0;
|
|
|
|
for (tx_seq, &entry_len) in n.iter().enumerate() {
|
|
let mut data = Vec::new();
|
|
for _ in 0..entry_len {
|
|
data.push(H256::random());
|
|
}
|
|
merkle.append_list(data.clone());
|
|
merkle.commit(Some(tx_seq as u64));
|
|
for i in (0..data.len()).step_by(6) {
|
|
let end = std::cmp::min(start_pos + i + 3, data.len());
|
|
let range_proof = merkle
|
|
.at_version(tx_seq as u64)
|
|
.unwrap()
|
|
.gen_range_proof(start_pos + i + 1, start_pos + end + 1)
|
|
.unwrap();
|
|
let r = range_proof.validate::<Sha3Algorithm>(&data[i..end], start_pos + i + 1);
|
|
assert!(r.is_ok(), "{:?}", r);
|
|
merkle.fill_with_range_proof(range_proof).unwrap();
|
|
}
|
|
|
|
start_pos += entry_len;
|
|
}
|
|
}
|
|
|
|
fn verify(data: &[H256], merkle: &mut AppendMerkleTree<H256, Sha3Algorithm>) {
|
|
for (i, item) in data.iter().enumerate() {
|
|
let proof = merkle.gen_proof(i + 1).unwrap();
|
|
let r = merkle.validate(&proof, item, i + 1);
|
|
assert!(matches!(r, Ok(true)), "{:?}", r);
|
|
}
|
|
for i in (0..data.len()).step_by(6) {
|
|
let end = std::cmp::min(i + 3, data.len());
|
|
let range_proof = merkle.gen_range_proof(i + 1, end + 1).unwrap();
|
|
let r = range_proof.validate::<Sha3Algorithm>(&data[i..end], i + 1);
|
|
assert!(r.is_ok(), "{:?}", r);
|
|
merkle.fill_with_range_proof(range_proof).unwrap();
|
|
}
|
|
}
|
|
}
|