diff --git a/src/lib.rs b/src/lib.rs index d84d563..b4e7fec 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -205,6 +205,30 @@ 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. + // Allowing this for compactness. + #[allow(clippy::option_map_unit_fn)] + 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> { @@ -341,6 +365,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 diff --git a/tests/subtree.rs b/tests/subtree.rs new file mode 100644 index 0000000..6af769d --- /dev/null +++ b/tests/subtree.rs @@ -0,0 +1,39 @@ +#[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)); + } +}