diff --git a/src/lib.rs b/src/lib.rs index e398029..e55c308 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -48,9 +48,13 @@ pub struct MMRNode { impl MMRNode { fn complete(&self) -> bool { - let leaves = self.data.end_height - self.data.start_height + 1; + let leaves = self.leaf_count(); leaves & (leaves - 1) == 0 } + + fn leaf_count(&self) -> u32 { + self.data.end_height - self.data.start_height + 1 + } } impl From for MMRNode { diff --git a/src/tree.rs b/src/tree.rs index f199461..c26696b 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -141,7 +141,7 @@ impl Tree { pub fn append_leaf(&mut self, root: NodeLink, new_leaf: NodeData) -> AppendTransaction { let is_complete= self.resolve_link(root).node.complete(); - let (new_root_node, mut appended) = if is_complete { + let (new_root, mut appended) = if is_complete { let new_leaf_link = self.push(new_leaf.into()); let mut appended = Vec::new(); @@ -153,35 +153,77 @@ impl Tree { self.resolve_link(new_leaf_link), ); - (new_root_node, appended) + (self.push_generated(new_root_node), appended) } else { - let (root_left_child, root_right_child) = { - let root = self.resolve_link(root).node; - ( - root.left.expect("Root should always have left child"), - root.right.expect("Root should always have right child"), +// let (root_left_child, root_right_child) = { +// let root = self.resolve_link(root).node; +// ( +// root.left.expect("Root should always have left child"), +// root.right.expect("Root should always have right child"), +// ) +// }; +// +// let nested_append = self.append_leaf(root_right_child, new_leaf); +// let appended = nested_append.appended; +// let subtree_root = nested_append.new_root; +// +// let new_root_node = combine_nodes( +// self.resolve_link(root_left_child), +// self.resolve_link(subtree_root), +// ); +// +// (new_root_node, appended) + + let new_leaf_link = self.push(new_leaf.into()); + let mut appended = Vec::new(); + appended.push(new_leaf_link); + + let mut peaks = self.get_peaks(root); + + let mut merge_stack = Vec::new(); + merge_stack.push(new_leaf_link); + + while let Some(next_peak) = peaks.pop() { + dbg!(next_peak); + let next_merge = merge_stack.pop().expect("there should be at least one, checked below"); + + if let Some(stored) = { + let peak = self.resolve_link(next_peak); + let m = self.resolve_link(next_merge); + if peak.node.leaf_count() == m.node.leaf_count() { + Some(combine_nodes(peak, m)) + } else { None } + } { + let link = self.push(stored); + merge_stack.push(link); + appended.push(link); + dbg!(link); + continue; + } + merge_stack.push(next_merge); + merge_stack.push(next_peak); + } + + let mut root = merge_stack.pop().expect("There should be at least one node in stack"); + while let Some(next_child) = merge_stack.pop() { + root = self.push_generated( + combine_nodes( + self.resolve_link(root), + self.resolve_link(next_child), + ) ) - }; + } - let nested_append = self.append_leaf(root_right_child, new_leaf); - let appended = nested_append.appended; - let subtree_root = nested_append.new_root; - - let new_root_node = combine_nodes( - self.resolve_link(root_left_child), - self.resolve_link(subtree_root), - ); - - (new_root_node, appended) + (root, appended) }; - let new_root = if new_root_node.complete() { - let new_root= self.push(new_root_node); - appended.push(new_root); - new_root - } else { - self.push_generated(new_root_node) - }; +// let new_root = if new_root_node.complete() { +// let new_root= self.push(new_root_node); +// appended.push(new_root); +// new_root +// } else { +// self.push_generated(new_root_node) +// }; AppendTransaction { new_root, @@ -448,7 +490,7 @@ mod tests { }); // *** APPEND #6 *** - + dbg!("go"); let append_tx = tree.append_leaf(new_root_link, leaf(6)); let new_root_link = append_tx.new_root; let new_root = tree.resolve_link(new_root_link).node; @@ -620,5 +662,18 @@ mod tests { TestResult::from_bool(if let NodeLink::Stored(2) = root { true } else { false }) } } + + fn leaf_count(number: u32) -> TestResult { + if (number > 1024 * 1024 || number < 3) { + TestResult::discard() + } else { + let (mut root, mut tree) = initial(); + for i in 0..(number-2) { + root = tree.append_leaf(root, leaf(i+3)).new_root; + } + + TestResult::from_bool(tree.resolve_link(root).node.leaf_count() == number) + } + } } } \ No newline at end of file