diff --git a/src/lib.rs b/src/lib.rs index 1a6585c..935053e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -189,6 +189,28 @@ impl Tree { self.vec.push(Node::new(value)); unsafe { self.get_unchecked_mut(id) } } + + /// Merge with another tree as orphan, returning the new root of tree being merged. + pub fn extend_tree(&mut self, mut other_tree: Tree) -> NodeMut { + let offset = self.vec.len(); + let offset_id = |id: NodeId| -> NodeId { + let old_index = id.to_index(); + let new_index = old_index + offset; + unsafe { NodeId::from_index(new_index) } + }; + let other_tree_root_id = offset_id(other_tree.root().id); + for node in other_tree.vec.iter_mut() { + node.parent.as_mut().map(|id| *id = offset_id(*id)); + node.prev_sibling.as_mut().map(|id| *id = offset_id(*id)); + node.next_sibling.as_mut().map(|id| *id = offset_id(*id)); + node.children.as_mut().map(|(id1, id2)| { + *id1 = offset_id(*id1); + *id2 = offset_id(*id2); + }); + } + self.vec.extend(other_tree.vec); + unsafe { self.get_unchecked_mut(other_tree_root_id) } + } } impl<'a, T: 'a> NodeRef<'a, T> { @@ -315,6 +337,18 @@ impl<'a, T: 'a> NodeMut<'a, T> { self.prepend_id(id) } + /// Appends a subtree, return the root of the merged subtree. + pub fn append_subtree(&mut self, subtree: Tree) -> NodeMut { + let root_id = self.tree.extend_tree(subtree).id; + self.append_id(root_id) + } + + /// Prepends a subtree, return the root of the merged subtree. + pub fn prepend_subtree(&mut self, subtree: Tree) -> NodeMut { + let root_id = self.tree.extend_tree(subtree).id; + self.prepend_id(root_id) + } + /// Inserts a new sibling before this node. /// /// # Panics @@ -438,6 +472,7 @@ impl<'a, T: 'a> NodeMut<'a, T> { { let mut new_sibling = self.tree.get_mut(new_sibling_id).unwrap(); + new_sibling.detach(); new_sibling.node().parent = Some(parent_id); new_sibling.node().prev_sibling = prev_sibling_id; new_sibling.node().next_sibling = Some(self.id); @@ -472,6 +507,7 @@ impl<'a, T: 'a> NodeMut<'a, T> { { let mut new_sibling = self.tree.get_mut(new_sibling_id).unwrap(); + new_sibling.detach(); new_sibling.node().parent = Some(parent_id); new_sibling.node().prev_sibling = Some(self.id); new_sibling.node().next_sibling = next_sibling_id; diff --git a/tests/subtree.rs b/tests/subtree.rs new file mode 100644 index 0000000..e540a28 --- /dev/null +++ b/tests/subtree.rs @@ -0,0 +1,38 @@ +#[macro_use] +extern crate ego_tree; + +#[cfg(test)] +mod test { + #[test] + fn prepend_subtree() { + let mut tree = tree!('a' => { 'b', 'c' => { 'd', 'e' } }); + let node_id = tree.root().first_child().unwrap().id(); + let mut node = tree.get_mut(node_id).unwrap(); + assert_eq!(node.value(), &'b'); + + let subtree = tree!('f' => { 'g', 'h' => { 'i', 'j' } }); + let mut root_subtree = node.prepend_subtree(subtree); + assert_eq!(root_subtree.parent().unwrap().value(), &'b'); + assert_eq!( + root_subtree.parent().unwrap().parent().unwrap().value(), + &'a' + ); + + let new_tree = tree!('a' => { 'b' => { 'f' => { 'g', 'h' => { 'i', 'j' } } }, 'c' => { 'd', 'e' } }); + assert_eq!(format!("{:#?}", tree), format!("{:#?}", new_tree)); + } + + #[test] + fn append_subtree() { + let mut tree = tree!('a' => { 'b', 'c' }); + let mut node = tree.root_mut(); + assert_eq!(node.value(), &'a'); + + let subtree = tree!('d' => { 'e', 'f' }); + let mut root_subtree = node.append_subtree(subtree); + assert_eq!(root_subtree.parent().unwrap().value(), &'a'); + + let new_tree = tree!('a' => { 'b', 'c', 'd' => { 'e', 'f' } }); + assert_eq!(format!("{:#?}", tree), format!("{:#?}", new_tree)); + } +} diff --git a/tests/tree.rs b/tests/tree.rs index a27f296..ac22561 100644 --- a/tests/tree.rs +++ b/tests/tree.rs @@ -1,6 +1,6 @@ extern crate ego_tree; -use ego_tree::Tree; +use ego_tree::{tree, Tree}; #[test] fn new() { @@ -69,3 +69,90 @@ fn neq() { let two = Tree::new('b'); assert_eq!(one, two); } + +#[test] +fn insert_id_after() { + let mut tree = tree! { + "root" => { + "a" => { + "child 1", + }, + "b" => { + "child 2", + }, + } + }; + + let a = tree.root().first_child().unwrap().id(); + let b = tree.root().last_child().unwrap().id(); + + assert_eq!(2, tree.root().children().count()); + assert_eq!(1, tree.get(a).unwrap().children().count()); + assert_eq!(1, tree.get(b).unwrap().children().count()); + + let child_1 = tree.get(a).unwrap().first_child().unwrap().id(); + tree.get_mut(b).unwrap().insert_id_after(child_1); + + assert_eq!( + 0, + tree.get(a).unwrap().children().count(), + "child 1 should be moved from a" + ); + assert_eq!( + 1, + tree.get(b).unwrap().children().count(), + "b should be unchanged" + ); + assert_eq!( + child_1, + tree.root().last_child().unwrap().id(), + "child 1 should be last child of root" + ); + assert_eq!(3, tree.root().children().count()); +} + +#[test] +fn insert_id_before() { + let mut tree = tree! { + "root" => { + "a" => { + "child 1", + }, + "b" => { + "child 2", + }, + } + }; + + let a = tree.root().first_child().unwrap().id(); + let b = tree.root().last_child().unwrap().id(); + + assert_eq!(2, tree.root().children().count()); + assert_eq!(1, tree.get(a).unwrap().children().count()); + assert_eq!(1, tree.get(b).unwrap().children().count()); + + let child_1 = tree.get(a).unwrap().first_child().unwrap().id(); + tree.get_mut(b).unwrap().insert_id_before(child_1); + + assert_eq!( + 0, + tree.get(a).unwrap().children().count(), + "child 1 should be moved from a" + ); + assert_eq!( + 1, + tree.get(b).unwrap().children().count(), + "b should be unchanged" + ); + assert_eq!( + b, + tree.root().last_child().unwrap().id(), + "b should be last child of root" + ); + assert_eq!( + child_1, + tree.get(b).unwrap().prev_sibling().unwrap().id(), + "child 1 should be between a and b" + ); + assert_eq!(3, tree.root().children().count()); +}