Merge pull request #30 from rust-scraper/shouya/master

Add support for appending and prepending subtree

Author: cfvescovo

src/lib.rs

@@ -205,6 +205,30 @@ impl<T> Tree<T> {
         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<T>) -> NodeMut<T> {
+        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<T>) -> NodeMut<T> {
+        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<T>) -> NodeMut<T> {
+        let root_id = self.tree.extend_tree(subtree).id;
+        self.prepend_id(root_id)
+    }
+
     /// Inserts a new sibling before this node.
     ///
     /// # Panics

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));
+    }
+}