carry root with tree

This commit is contained in:
NikVolf 2019-09-06 18:52:27 +03:00
parent 6b36cb5a51
commit 5d2f84a154
2 changed files with 57 additions and 34 deletions

View File

@ -2,7 +2,7 @@ extern crate zcash_mmr as mmr;
use mmr::{NodeData, Tree, EntryLink, Entry}; use mmr::{NodeData, Tree, EntryLink, Entry};
fn prepare_tree(vec: Vec<NodeData>) -> (Tree, EntryLink) { fn prepare_tree(vec: Vec<NodeData>) -> Tree {
assert!(vec.len() > 0); assert!(vec.len() > 0);
@ -213,7 +213,7 @@ fn main() {
}, },
); );
let (_tree, root) = prepare_tree(initial_tree_vec); let tree = prepare_tree(initial_tree_vec);
println!("root: {}", root); println!("root: {}", tree.root());
} }

View File

@ -9,7 +9,6 @@ use crate::{Entry, EntryLink, NodeData, Error, EntryKind};
/// Exact amount of the loaded data can be calculated by the constructing party, /// Exact amount of the loaded data can be calculated by the constructing party,
/// depending on the length of the tree and maximum amount of operations that are going /// depending on the length of the tree and maximum amount of operations that are going
/// to happen after construction. /// to happen after construction.
#[derive(Default)]
pub struct Tree { pub struct Tree {
stored: HashMap<u32, Entry>, stored: HashMap<u32, Entry>,
@ -20,6 +19,8 @@ pub struct Tree {
// number of virtual nodes generated // number of virtual nodes generated
generated_count: u32, generated_count: u32,
root: EntryLink,
} }
/// Result of appending one or several leaves. /// Result of appending one or several leaves.
@ -76,22 +77,34 @@ impl Tree {
/// Populate tree with plain list of the leaves/nodes. Mostly for tests, /// Populate tree with plain list of the leaves/nodes. Mostly for tests,
/// since this `Tree` structure is for partially loaded tree. /// since this `Tree` structure is for partially loaded tree.
pub fn populate(loaded: Vec<Entry>) -> Self { pub fn populate(loaded: Vec<Entry>, root: EntryLink) -> Self {
let mut result = Tree::default(); let mut result = Tree::invalid();
result.stored_count = loaded.len() as u32; result.stored_count = loaded.len() as u32;
for (idx, item) in loaded.into_iter().enumerate() { for (idx, item) in loaded.into_iter().enumerate() {
result.stored.insert(idx as u32, item); result.stored.insert(idx as u32, item);
} }
result.root = root;
result result
} }
fn invalid() -> Self {
Tree {
root: EntryLink::Generated(0),
generated: Default::default(),
stored: Default::default(),
generated_count: 0,
stored_count: 0,
}
}
pub fn new( pub fn new(
length: u32, length: u32,
peaks: Vec<(u32, Entry)>, peaks: Vec<(u32, Entry)>,
extra: Vec<(u32, Entry)>, extra: Vec<(u32, Entry)>,
) -> (Self, EntryLink) { ) -> Self {
let mut result = Tree::default(); let mut result = Tree::invalid();
result.stored_count = length; result.stored_count = length;
let mut gen = 0; let mut gen = 0;
@ -109,7 +122,9 @@ impl Tree {
gen += 1; gen += 1;
} }
(result, root) result.root = root;
result
} }
fn get_peaks(&self, root: EntryLink, target: &mut Vec<EntryLink>) -> Result<(), Error> { fn get_peaks(&self, root: EntryLink, target: &mut Vec<EntryLink>) -> Result<(), Error> {
@ -132,7 +147,8 @@ impl Tree {
} }
/// Append one leaf to the tree. /// Append one leaf to the tree.
pub fn append_leaf(&mut self, root: EntryLink, new_leaf: NodeData) -> Result<AppendTransaction, Error> { pub fn append_leaf(&mut self, new_leaf: NodeData) -> Result<AppendTransaction, Error> {
let root = self.root;
let new_leaf_link = self.push(new_leaf.into()); let new_leaf_link = self.push(new_leaf.into());
let mut appended = Vec::new(); let mut appended = Vec::new();
appended.push(new_leaf_link); appended.push(new_leaf_link);
@ -172,6 +188,8 @@ impl Tree {
) )
} }
self.root = new_root;
Ok(AppendTransaction { Ok(AppendTransaction {
new_root, new_root,
appended, appended,
@ -196,10 +214,10 @@ impl Tree {
} }
/// Truncate one leaf from the end of the tree. /// Truncate one leaf from the end of the tree.
pub fn truncate_leaf(&mut self, root: EntryLink) -> Result<DeleteTransaction, Error> { pub fn truncate_leaf(&mut self) -> Result<DeleteTransaction, Error> {
let root = { let root = {
let (leaves, root_left_child) = { let (leaves, root_left_child) = {
let n = self.resolve_link(root)?; let n = self.resolve_link(self.root)?;
( (
n.node.leaf_count(), n.node.leaf_count(),
n.node.left()?, n.node.left()?,
@ -207,12 +225,13 @@ impl Tree {
}; };
if leaves & 1 != 0 { if leaves & 1 != 0 {
self.pop(); self.pop();
self.root = root_left_child;
return Ok(DeleteTransaction { return Ok(DeleteTransaction {
truncated: 1, truncated: 1,
new_root: root_left_child, new_root: root_left_child,
}) })
} else { } else {
self.resolve_link(root)? self.resolve_link(self.root)?
} }
}; };
@ -245,6 +264,8 @@ impl Tree {
for _ in 0..truncated { self.pop(); } for _ in 0..truncated { self.pop(); }
self.root = new_root;
Ok(DeleteTransaction { Ok(DeleteTransaction {
new_root, new_root,
truncated, truncated,
@ -255,6 +276,8 @@ impl Tree {
pub fn len(&self) -> u32 { pub fn len(&self) -> u32 {
self.stored_count self.stored_count
} }
pub fn root(&self) -> EntryLink { self.root }
} }
@ -331,7 +354,7 @@ mod tests {
kind: EntryKind::Leaf, kind: EntryKind::Leaf,
}; };
(EntryLink::Stored(2), Tree::populate(vec![node1, node2, node3])) (EntryLink::Stored(2), Tree::populate(vec![node1, node2, node3], EntryLink::Stored(2)))
} }
// returns tree with specified number of leafs and it's root // returns tree with specified number of leafs and it's root
@ -340,7 +363,7 @@ mod tests {
let (mut root, mut tree) = initial(); let (mut root, mut tree) = initial();
for i in 2..length { for i in 2..length {
root = tree root = tree
.append_leaf(root, leaf(i+1).into()) .append_leaf(leaf(i+1).into())
.expect("Failed to append") .expect("Failed to append")
.new_root; .new_root;
} }
@ -354,7 +377,7 @@ mod tests {
// ** APPEND 3 ** // ** APPEND 3 **
let append_tx = tree let append_tx = tree
.append_leaf(root, leaf(3)) .append_leaf(leaf(3))
.expect("Failed to append"); .expect("Failed to append");
let new_root_link = append_tx.new_root; let new_root_link = append_tx.new_root;
let new_root = tree.resolve_link(new_root_link).expect("Failed to resolve root").node; let new_root = tree.resolve_link(new_root_link).expect("Failed to resolve root").node;
@ -377,7 +400,7 @@ mod tests {
// ** APPEND 4 ** // ** APPEND 4 **
let append_tx = tree let append_tx = tree
.append_leaf(new_root_link, leaf(4)) .append_leaf(leaf(4))
.expect("Failed to append"); .expect("Failed to append");
let new_root_link = append_tx.new_root; let new_root_link = append_tx.new_root;
@ -406,7 +429,7 @@ mod tests {
// ** APPEND 5 ** // ** APPEND 5 **
let append_tx = tree let append_tx = tree
.append_leaf(new_root_link, leaf(5)) .append_leaf(leaf(5))
.expect("Failed to append"); .expect("Failed to append");
let new_root_link = append_tx.new_root; let new_root_link = append_tx.new_root;
let new_root = tree.resolve_link(new_root_link).expect("Failed to resolve root").node; let new_root = tree.resolve_link(new_root_link).expect("Failed to resolve root").node;
@ -439,7 +462,7 @@ mod tests {
// *** APPEND #6 *** // *** APPEND #6 ***
let append_tx = tree let append_tx = tree
.append_leaf(new_root_link, leaf(6)) .append_leaf(leaf(6))
.expect("Failed to append"); .expect("Failed to append");
let new_root_link = append_tx.new_root; let new_root_link = append_tx.new_root;
let new_root = tree.resolve_link(new_root_link).expect("Failed to resolve root").node; let new_root = tree.resolve_link(new_root_link).expect("Failed to resolve root").node;
@ -475,7 +498,7 @@ mod tests {
// *** APPEND #7 *** // *** APPEND #7 ***
let append_tx = tree let append_tx = tree
.append_leaf(new_root_link, leaf(7)) .append_leaf(leaf(7))
.expect("Failed to append"); .expect("Failed to append");
let new_root_link = append_tx.new_root; let new_root_link = append_tx.new_root;
let new_root = tree let new_root = tree
@ -518,7 +541,7 @@ mod tests {
fn truncate_simple() { fn truncate_simple() {
let (mut tree, root) = generated(9); let (mut tree, root) = generated(9);
let delete_tx = tree let delete_tx = tree
.truncate_leaf(root) .truncate_leaf()
.expect("Failed to truncate"); .expect("Failed to truncate");
// initial tree: // initial tree:
@ -553,7 +576,7 @@ mod tests {
fn truncate_generated() { fn truncate_generated() {
let (mut tree, root) = generated(10); let (mut tree, root) = generated(10);
let delete_tx = tree let delete_tx = tree
.truncate_leaf(root) .truncate_leaf()
.expect("Failed to truncate"); .expect("Failed to truncate");
// initial tree: // initial tree:
@ -611,13 +634,13 @@ mod tests {
for i in 0..2 { for i in 0..2 {
root = tree root = tree
.append_leaf(root, leaf(i+3)) .append_leaf(leaf(i+3))
.expect("Failed to append") .expect("Failed to append")
.new_root; .new_root;
} }
assert_eq!(tree.len(), 7); assert_eq!(tree.len(), 7);
tree.truncate_leaf(root).expect("Failed to truncate"); tree.truncate_leaf().expect("Failed to truncate");
assert_eq!(tree.len(), 4); assert_eq!(tree.len(), 4);
} }
@ -630,7 +653,7 @@ mod tests {
for i in 0..4094 { for i in 0..4094 {
root = tree root = tree
.append_leaf(root, leaf(i+3)) .append_leaf(leaf(i+3))
.expect("Failed to append") .expect("Failed to append")
.new_root; .new_root;
} }
@ -638,7 +661,7 @@ mod tests {
for _ in 0..2049 { for _ in 0..2049 {
root = tree root = tree
.truncate_leaf(root) .truncate_leaf()
.expect("Failed to truncate") .expect("Failed to truncate")
.new_root; .new_root;
} }
@ -655,12 +678,12 @@ mod tests {
let (mut root, mut tree) = initial(); let (mut root, mut tree) = initial();
for i in 0..number { for i in 0..number {
root = tree root = tree
.append_leaf(root, leaf(i+3)) .append_leaf(leaf(i+3))
.expect("Failed to append") .expect("Failed to append")
.new_root; .new_root;
} }
for _ in 0..number { for _ in 0..number {
root = tree.truncate_leaf(root).expect("Failed to truncate").new_root; root = tree.truncate_leaf().expect("Failed to truncate").new_root;
} }
TestResult::from_bool(if let EntryLink::Stored(2) = root { true } else { false }) TestResult::from_bool(if let EntryLink::Stored(2) = root { true } else { false })
@ -674,7 +697,7 @@ mod tests {
let (mut root, mut tree) = initial(); let (mut root, mut tree) = initial();
for i in 1..(number-1) { for i in 1..(number-1) {
root = tree root = tree
.append_leaf(root, leaf(i+2)) .append_leaf(leaf(i+2))
.expect("Failed to append") .expect("Failed to append")
.new_root; .new_root;
} }
@ -697,7 +720,7 @@ mod tests {
} else { } else {
let (mut root, mut tree) = initial(); let (mut root, mut tree) = initial();
for i in 1..(number-1) { for i in 1..(number-1) {
root = tree.append_leaf(root, leaf(i+2)).expect("Failed to append").new_root; root = tree.append_leaf(leaf(i+2)).expect("Failed to append").new_root;
} }
TestResult::from_bool( TestResult::from_bool(
@ -720,10 +743,10 @@ mod tests {
} else { } else {
let (mut root, mut tree) = initial(); let (mut root, mut tree) = initial();
for i in 0..add { for i in 0..add {
root = tree.append_leaf(root, leaf(i+3)).expect("Failed to append").new_root; root = tree.append_leaf(leaf(i+3)).expect("Failed to append").new_root;
} }
for _ in 0..delete { for _ in 0..delete {
root = tree.truncate_leaf(root).expect("Failed to truncate").new_root; root = tree.truncate_leaf().expect("Failed to truncate").new_root;
} }
let total = add - delete + 2; let total = add - delete + 2;
@ -747,10 +770,10 @@ mod tests {
} else { } else {
let (mut root, mut tree) = initial(); let (mut root, mut tree) = initial();
for i in 0..add { for i in 0..add {
root = tree.append_leaf(root, leaf(i+3)).expect("Failed to append").new_root; root = tree.append_leaf(leaf(i+3)).expect("Failed to append").new_root;
} }
for _ in 0..delete { for _ in 0..delete {
root = tree.truncate_leaf(root).expect("Failed to truncate").new_root; root = tree.truncate_leaf().expect("Failed to truncate").new_root;
} }
let total = add - delete + 2; let total = add - delete + 2;