diff --git a/src/tree.rs b/src/tree.rs index 5495b49b4..119673eee 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -22,6 +22,11 @@ pub struct AppendTransaction { pub new_root: NodeLink, } +pub struct DeleteTransaction { + pub truncated: u32, + pub new_root: NodeLink, +} + impl Tree { fn resolve_link(&self, link: NodeLink) -> IndexedNode { match link { @@ -58,6 +63,7 @@ impl Tree { NodeLink::Generated(idx) } + // TODO: populate both stored and generated nodes? pub fn populate(loaded: Vec) -> Self { let mut result = Tree::default(); result.stored_count = loaded.len() as u32; @@ -69,7 +75,7 @@ impl Tree { } - pub fn append(&mut self, root: NodeLink, new_leaf: NodeData) -> AppendTransaction { + pub fn append_leaf(&mut self, root: NodeLink, new_leaf: NodeData) -> AppendTransaction { let is_complete= self.resolve_link(root).node.complete(); @@ -86,8 +92,6 @@ impl Tree { ); (new_root_node, appended) - - } else { let (root_left_child, root_right_child) = { let root = self.resolve_link(root).node; @@ -97,7 +101,7 @@ impl Tree { ) }; - let nested_append = self.append(root_right_child, new_leaf); + let nested_append = self.append_leaf(root_right_child, new_leaf); let appended = nested_append.appended; let subtree_root = nested_append.new_root; @@ -123,6 +127,61 @@ impl Tree { } } + fn pop(&mut self) { + self.stored.remove(&(self.stored_count-1)); + self.stored_count = self.stored_count - 1; + } + + pub fn truncate_leaf(&mut self, root: NodeLink) -> DeleteTransaction { + let root = { + let n = self.resolve_link(root); + let leaves = n.node.data.end_height - n.node.data.start_height + 1; + if leaves & 1 != 0 { + return DeleteTransaction { + truncated: 1, + new_root: n.node.left.expect("Root should have left child while deleting"), + } + } else { + n + } + }; + + let mut peaks = vec![root.node.left.expect("Root should have left child")]; + let mut subtree_root_link = root.node.right.expect("Root should have right child"); + let mut truncated = 1; + + loop { + let left_link = self.resolve_link(subtree_root_link).node.left; + if let Some(left_link) = left_link { + peaks.push(left_link); + subtree_root_link = self + .resolve_link(subtree_root_link).node.right + .expect("If left exists, right should exist as well"); + truncated += 1; + } else { + break; + } + } + + let root = peaks.drain(0..1).nth(0).expect("At lest 2 elements in peaks"); + let new_root = peaks.into_iter().fold( + root, + |root, next_peak| + self.push_generated( + combine_nodes( + self.resolve_link(root), + self.resolve_link(next_peak) + ) + ) + ); + + for _ in 0..truncated { self.pop(); } + + DeleteTransaction { + new_root, + truncated, + } + } } @@ -165,7 +224,6 @@ mod tests { fn leaf(height: u32) -> NodeData { NodeData { - // TODO: hash children subtree_commitment: [0u8; 32], start_time: 0, end_time: 0, @@ -173,8 +231,6 @@ mod tests { end_target: 0, start_sapling_root: [0u8; 32], end_sapling_root: [0u8; 32], - - // TODO: sum work? subtree_total_work: 0, start_height: height, end_height: height, @@ -184,7 +240,6 @@ mod tests { fn node(start_height: u32, end_height: u32) -> NodeData { NodeData { - // TODO: hash children subtree_commitment: [0u8; 32], start_time: 0, end_time: 0, @@ -192,8 +247,6 @@ mod tests { end_target: 0, start_sapling_root: [0u8; 32], end_sapling_root: [0u8; 32], - - // TODO: sum work? subtree_total_work: 0, start_height: start_height, end_height: end_height, @@ -214,10 +267,23 @@ mod tests { Tree::populate(vec![node1, node2, node3]) } + // returns tree with specified number of leafs and it's root + fn generated(length: u32) -> (Tree, NodeLink) { + assert!(length > 3); + let mut tree = initial(); + let mut root = NodeLink::Stored(2); + + for i in 2..length { + root = tree.append_leaf(root, leaf(i+1).into()).new_root; + } + + (tree, root) + } + #[test] fn discrete_append() { let mut tree = initial(); - let append_tx = tree.append(NodeLink::Stored(2), leaf(3)); + let append_tx = tree.append_leaf(NodeLink::Stored(2), leaf(3)); let new_root_link = append_tx.new_root; let new_root = tree.resolve_link(new_root_link).node; @@ -237,7 +303,7 @@ mod tests { assert_eq!(new_root.data.end_height, 3); assert_eq!(append_tx.appended.len(), 1); - let append_tx = tree.append(new_root_link, leaf(4)); + let append_tx = tree.append_leaf(new_root_link, leaf(4)); let new_root_link = append_tx.new_root; let new_root = tree.resolve_link(new_root_link).node; @@ -260,7 +326,7 @@ mod tests { assert_eq!(new_root.data.end_height, 4); assert_eq!(append_tx.appended.len(), 3); - let append_tx = tree.append(new_root_link, leaf(5)); + let append_tx = tree.append_leaf(new_root_link, leaf(5)); let new_root_link = append_tx.new_root; let new_root = tree.resolve_link(new_root_link).node; @@ -286,4 +352,16 @@ mod tests { assert_eq!(append_tx.appended.len(), 1); } + #[test] + fn truncate_simple() { + let (mut tree, root) = generated(9); + let delete_tx = tree.truncate_leaf(root); + + match delete_tx.new_root { + NodeLink::Stored(14) => { /* ok */ }, + _ => panic!("Root should be stored(14)") + } + + } + } \ No newline at end of file