Merge pull request #182 from orxfun/leaves_into_with

Leaves into with

Author: orxfun

src/node_mut.rs

@@ -2142,7 +2142,7 @@ where
     /// * [`Dfs`] for (pre-order) depth-first ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Depth-first_search))
     /// * [`PostOrder`] for post-order ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Post-order,_LRN))
     ///
-    /// As opposed to [`walk_mut`], this method does require internal allocation.
+    /// As opposed to [`walk_mut`], this method does not require internal allocation.
     /// Furthermore, it allows to attach node depths or sibling indices to the yield values.
     /// Please see the examples below.
     ///
@@ -2371,7 +2371,7 @@ where
     /// assert!(tree.is_empty());
     /// assert_eq!(tree.get_root(), None);
     /// ```
-    pub fn into_walk<T>(self) -> impl Iterator<Item = V::Item> + use<'a, T, V, M, P, MO>
+    pub fn into_walk<T>(self) -> impl Iterator<Item = V::Item>
     where
         T: Traverser<OverData>,
     {
@@ -2390,7 +2390,7 @@ where
     /// * [`Dfs`] for (pre-order) depth-first ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Depth-first_search))
     /// * [`PostOrder`] for post-order ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Post-order,_LRN))
     ///
-    /// As opposed to [`into_walk`], this method does require internal allocation.
+    /// As opposed to [`into_walk`], this method does not require internal allocation.
     /// Furthermore, it allows to attach node depths or sibling indices to the yield values.
     /// Please see the examples below.
     ///
@@ -2623,7 +2623,7 @@ where
     /// * [`Dfs`] for (pre-order) depth-first ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Depth-first_search))
     /// * [`PostOrder`] for post-order ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Post-order,_LRN))
     ///
-    /// As opposed to [`leaves_mut`], this method does require internal allocation.
+    /// As opposed to [`leaves_mut`], this method does not require internal allocation.
     /// Furthermore, it allows to attach node depths or sibling indices to the yield values.
     /// Please see the examples below.
     ///
@@ -2755,6 +2755,197 @@ where
             })
     }
 
+    /// Creates an iterator that yields owned (removed) data of all leaves of the subtree rooted at this node.
+    ///
+    /// Note that once the returned iterator is dropped, regardless of whether it is completely used up or not,
+    /// the subtree rooted at this node will be **removed** from the tree it belongs to.
+    /// If this node is the root of the tree, the tree will be left empty.
+    ///
+    /// The order of the elements is determined by the generic [`Traverser`] parameter `T`.
+    /// Available implementations are:
+    /// * [`Bfs`] for breadth-first ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Breadth-first_search))
+    /// * [`Dfs`] for (pre-order) depth-first ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Depth-first_search))
+    /// * [`PostOrder`] for post-order ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Post-order,_LRN))
+    ///
+    /// See also [`leaves`] and [`leaves_mut`] for iterators over shared and mutable references, respectively.
+    ///
+    /// Note that tree traversing methods typically allocate a temporary data structure that is dropped once the
+    /// iterator is dropped.
+    /// In use cases where we repeatedly iterate using any of the **leaves** methods over different nodes or different
+    /// trees, we can avoid the allocation by creating the traverser only once and using [`leaves_with`], [`leaves_mut_with`]
+    /// and [`into_leaves_with`] methods instead.
+    /// These methods additionally allow for iterating over nodes rather than data; and yielding node depths and sibling
+    /// indices in addition to node data.
+    ///
+    /// > **(!)** As a method that removes nodes from the tree, this method might result in invalidating indices that are
+    /// > cached earlier in the [`Auto`] mode, but not in the [`Lazy`] mode. Please see the documentation of [MemoryPolicy]
+    /// > for details of node index validity. Specifically, the examples in the "Lazy Memory Claim: Preventing Invalid Indices"
+    /// > section presents a convenient way that allows us to make sure that the indices are valid.
+    ///
+    /// [`Auto`]: crate::Auto
+    /// [`Lazy`]: crate::Lazy
+    /// [`MemoryPolicy`]: crate::MemoryPolicy
+    ///
+    /// [`Bfs`]: crate::Bfs
+    /// [`Dfs`]: crate::Dfs
+    /// [`PostOrder`]: crate::PostOrder
+    /// [`leaves`]: crate::NodeRef::leaves
+    /// [`leaves_mut`]: crate::NodeMut::walk_mut
+    /// [`leaves_with`]: crate::NodeRef::leaves_with
+    /// [`leaves_mut_with`]: crate::NodeMut::leaves_mut_with
+    /// [`into_leaves_with`]: crate::NodeMut::into_leaves_with
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use orx_tree::*;
+    ///
+    /// //      1
+    /// //     ╱ ╲
+    /// //    ╱   ╲
+    /// //   2     3
+    /// //  ╱ ╲   ╱ ╲
+    /// // 4   5 6   7
+    /// // |     |  ╱ ╲
+    /// // 8     9 10  11
+    ///
+    /// let mut tree = DynTree::new(1);
+    /// let [id2, id3] = tree.root_mut().push_children([2, 3]);
+    /// let [id4, _] = tree.node_mut(id2).push_children([4, 5]);
+    /// tree.node_mut(id4).push_child(8);
+    /// let [id6, id7] = tree.node_mut(id3).push_children([6, 7]);
+    /// tree.node_mut(id6).push_child(9);
+    /// tree.node_mut(id7).push_children([10, 11]);
+    ///
+    /// // keep indices valid during removals
+    /// let mut tree = tree.into_lazy_reclaim();
+    ///
+    /// let n3 = tree.node_mut(id3);
+    /// let leaves: Vec<_> = n3.into_leaves::<Dfs>().collect();
+    /// assert_eq!(leaves, [9, 10, 11]);
+    ///
+    /// //      1
+    /// //     ╱
+    /// //    ╱
+    /// //   2
+    /// //  ╱ ╲
+    /// // 4   5
+    /// // |
+    /// // 8
+    /// let remaining: Vec<_> = tree.root().walk::<Dfs>().copied().collect();
+    /// assert_eq!(remaining, [1, 2, 4, 8, 5]);
+    ///
+    /// let leaves: Vec<_> = tree.root_mut().into_leaves::<Dfs>().collect();
+    /// assert_eq!(leaves, [8, 5]);
+    ///
+    /// assert!(tree.is_empty());
+    /// ```
+    pub fn into_leaves<T>(self) -> impl Iterator<Item = V::Item>
+    where
+        T: Traverser<OverData>,
+    {
+        let storage = T::Storage::<V>::default();
+        T::into_iter_with_storage_filtered(self, storage, |x| {
+            let ptr = <<OverData as Over>::Enumeration as Enumeration>::node_data(&x);
+            unsafe { &*ptr.ptr() }.next().is_empty()
+        })
+    }
+
+    /// Creates an iterator that yields owned (removed) data of all leaves of the subtree rooted at this node.
+    ///
+    /// Note that once the returned iterator is dropped, regardless of whether it is completely used up or not,
+    /// the subtree rooted at this node will be **removed** from the tree it belongs to.
+    /// If this node is the root of the tree, the tree will be left empty.
+    ///
+    /// The order of the elements is determined by the type of the `traverser` which implements [`Traverser`].
+    /// Available implementations are:
+    /// * [`Bfs`] for breadth-first ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Breadth-first_search))
+    /// * [`Dfs`] for (pre-order) depth-first ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Depth-first_search))
+    /// * [`PostOrder`] for post-order ([wikipedia](https://en.wikipedia.org/wiki/Tree_traversal#Post-order,_LRN))
+    ///
+    /// As opposed to [`into_leaves`], this method does not require internal allocation.
+    /// Furthermore, it allows to attach node depths or sibling indices to the yield values.
+    /// Please see the examples below.
+    ///
+    /// [`into_leaves`]: crate::NodeMut::into_leaves
+    /// [`Bfs`]: crate::Bfs
+    /// [`Dfs`]: crate::Dfs
+    /// [`PostOrder`]: crate::PostOrder
+    ///
+    /// > **(!)** As a method that removes nodes from the tree, this method might result in invalidating indices that are
+    /// > cached earlier in the [`Auto`] mode, but not in the [`Lazy`] mode. Please see the documentation of [MemoryPolicy]
+    /// > for details of node index validity. Specifically, the examples in the "Lazy Memory Claim: Preventing Invalid Indices"
+    /// > section presents a convenient way that allows us to make sure that the indices are valid.
+    ///
+    /// [`Auto`]: crate::Auto
+    /// [`Lazy`]: crate::Lazy
+    /// [`MemoryPolicy`]: crate::MemoryPolicy
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use orx_tree::*;
+    ///
+    /// //      1
+    /// //     ╱ ╲
+    /// //    ╱   ╲
+    /// //   2     3
+    /// //  ╱ ╲   ╱ ╲
+    /// // 4   5 6   7
+    /// // |     |  ╱ ╲
+    /// // 8     9 10  11
+    ///
+    /// let mut tree = DynTree::new(1);
+    /// let [id2, id3] = tree.root_mut().push_children([2, 3]);
+    /// let [id4, _] = tree.node_mut(id2).push_children([4, 5]);
+    /// tree.node_mut(id4).push_child(8);
+    /// let [id6, id7] = tree.node_mut(id3).push_children([6, 7]);
+    /// tree.node_mut(id6).push_child(9);
+    /// tree.node_mut(id7).push_children([10, 11]);
+    ///
+    /// // create the traverser 'dfs' only once, use it many times
+    /// // to walk over references, mutable references or removed values
+    /// // without additional allocation
+    ///
+    /// let mut dfs = Dfs::default();
+    ///
+    /// // keep indices valid during removals
+    /// let mut tree = tree.into_lazy_reclaim();
+    ///
+    /// let n3 = tree.node_mut(id3);
+    /// let leaves: Vec<_> = n3.into_leaves_with(&mut dfs).collect();
+    /// assert_eq!(leaves, [9, 10, 11]);
+    ///
+    /// //      1
+    /// //     ╱
+    /// //    ╱
+    /// //   2
+    /// //  ╱ ╲
+    /// // 4   5
+    /// // |
+    /// // 8
+    /// let remaining: Vec<_> = tree.root().walk_with(&mut dfs).copied().collect();
+    /// assert_eq!(remaining, [1, 2, 4, 8, 5]);
+    ///
+    /// let leaves: Vec<_> = tree.root_mut().into_leaves_with(&mut dfs).collect();
+    /// assert_eq!(leaves, [8, 5]);
+    ///
+    /// assert!(tree.is_empty());
+    /// ```
+    pub fn into_leaves_with<T, O>(
+        self,
+        traverser: &'a mut T,
+    ) -> impl Iterator<Item = OverItemInto<'a, V, O>>
+    where
+        O: OverMut,
+        T: Traverser<O>,
+    {
+        T::into_iter_with_storage_filtered(self, traverser.storage_mut(), |x| {
+            let ptr = <O::Enumeration as Enumeration>::node_data(x);
+            unsafe { &*ptr.ptr() }.next().is_empty()
+        })
+    }
+
     // recursive
 
     /// Recursively sets the data of all nodes belonging to the subtree rooted at this node using the `compute_data`

src/traversal/breadth_first/into_iter_filtered.rs

@@ -0,0 +1,126 @@
+use super::BreadthFirstEnumeration;
+use super::iter_ptr::BfsIterPtr;
+use super::queue::Item;
+use crate::TreeVariant;
+use crate::aliases::Col;
+use crate::memory::MemoryPolicy;
+use crate::pinned_storage::PinnedStorage;
+use alloc::collections::VecDeque;
+use orx_self_or::SoM;
+use orx_selfref_col::{NodePtr, Refs};
+
+pub struct BfsIterIntoFiltered<'a, V, M, P, E, S, F>
+where
+    V: TreeVariant,
+    M: MemoryPolicy,
+    P: PinnedStorage,
+    E: BreadthFirstEnumeration,
+    S: SoM<VecDeque<Item<V, E>>>,
+    F: Fn(&E::Item<NodePtr<V>>) -> bool,
+{
+    col: &'a mut Col<V, M, P>,
+    root_ptr: NodePtr<V>,
+    iter: BfsIterPtr<V, E, S>,
+    filter: F,
+}
+
+impl<'a, V, M, P, E, S, F> BfsIterIntoFiltered<'a, V, M, P, E, S, F>
+where
+    V: TreeVariant,
+    M: MemoryPolicy,
+    P: PinnedStorage,
+    E: BreadthFirstEnumeration,
+    S: SoM<VecDeque<Item<V, E>>>,
+    F: Fn(&E::Item<NodePtr<V>>) -> bool,
+{
+    /// # Safety
+    ///
+    /// We are creating a mutable iterator over nodes of the collection `col`.
+    /// This is safe only when the second argument `iter` makes sure that there exists only one mutable
+    /// reference to the collection.
+    ///
+    /// This is the case how this method is used, as follows:
+    /// * Mutable iterators are created through the `Dfs` traverser's `TraverserMut::iter_mut` method.
+    /// * This method requires a mutable reference to a mutable node `NodeMut` which is guaranteed to
+    ///   be the only reference to the collection.
+    /// * Finally, this iterator's lifetime is equal to the borrow duration of the mutable node.
+    #[allow(clippy::type_complexity)]
+    pub(crate) unsafe fn from(
+        (col, iter, root_ptr): (&'a mut Col<V, M, P>, BfsIterPtr<V, E, S>, NodePtr<V>),
+        filter: F,
+    ) -> Self {
+        let node = unsafe { &mut *root_ptr.ptr_mut() };
+
+        match node.prev().get() {
+            Some(parent) => {
+                let parent = unsafe { &mut *parent.ptr_mut() };
+                let sibling_idx = parent.next_mut().remove(unsafe { root_ptr.ptr() as usize });
+                debug_assert!(sibling_idx.is_some());
+
+                node.prev_mut().clear();
+            }
+            None => {
+                // node_ptr points to the root node
+                col.ends_mut().clear();
+            }
+        }
+
+        Self {
+            col,
+            root_ptr,
+            iter,
+            filter,
+        }
+    }
+
+    fn take_element(&mut self, element: E::Item<NodePtr<V>>) -> E::Item<V::Item> {
+        E::map_node_data(element, |ptr| {
+            let col = unsafe { &mut *(self.col as *mut Col<V, M, P>) };
+            col.close(ptr)
+        })
+    }
+}
+
+impl<V, M, P, E, S, F> Iterator for BfsIterIntoFiltered<'_, V, M, P, E, S, F>
+where
+    V: TreeVariant,
+    M: MemoryPolicy,
+    P: PinnedStorage,
+    E: BreadthFirstEnumeration,
+    S: SoM<VecDeque<Item<V, E>>>,
+    F: Fn(&E::Item<NodePtr<V>>) -> bool,
+{
+    type Item = E::Item<V::Item>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            match self.iter.next() {
+                Some(ptr) => {
+                    let included = (self.filter)(&ptr);
+                    let element = self.take_element(ptr);
+                    if included {
+                        return Some(element);
+                    }
+                }
+                None => return None,
+            }
+        }
+    }
+}
+
+impl<V, M, P, E, S, F> Drop for BfsIterIntoFiltered<'_, V, M, P, E, S, F>
+where
+    V: TreeVariant,
+    M: MemoryPolicy,
+    P: PinnedStorage,
+    E: BreadthFirstEnumeration,
+    S: SoM<VecDeque<Item<V, E>>>,
+    F: Fn(&E::Item<NodePtr<V>>) -> bool,
+{
+    fn drop(&mut self) {
+        while let Some(element) = self.iter.next() {
+            self.take_element(element);
+        }
+        self.col.reclaim_from_closed_node(self.root_ptr);
+    }
+}

src/traversal/breadth_first/mod.rs

@@ -3,6 +3,7 @@ mod tests;
 
 mod bfs_enumeration;
 pub(crate) mod into_iter;
+pub(crate) mod into_iter_filtered;
 pub(crate) mod iter_mut;
 pub(crate) mod iter_ptr;
 pub(crate) mod iter_ref;

src/traversal/breadth_first/traverser_core.rs

@@ -8,6 +8,7 @@ use crate::{
     pinned_storage::PinnedStorage,
     traversal::{
         Over, OverMut,
+        breadth_first::into_iter_filtered::BfsIterIntoFiltered,
         enumeration::Enumeration,
         over::OverItem,
         over_mut::{OverItemInto, OverItemMut},
@@ -126,4 +127,22 @@ impl<O: Over> TraverserCore<O> for Bfs<O> {
         let iter_ptr = BfsIterPtr::<V, O::Enumeration, _>::from((storage, root));
         unsafe { BfsIterInto::<V, M, P, _, _>::from((col, iter_ptr, root)) }
     }
+
+    fn into_iter_with_storage_filtered<'a, V, M, P, MO, F>(
+        node_mut: NodeMut<'a, V, M, P, MO>,
+        storage: impl SoM<Self::Storage<V>>,
+        filter: F,
+    ) -> impl Iterator<Item = OverItemInto<'a, V, O>>
+    where
+        V: TreeVariant + 'a,
+        M: MemoryPolicy,
+        P: PinnedStorage,
+        MO: NodeMutOrientation,
+        O: Over,
+        F: Fn(&<<O as Over>::Enumeration as Enumeration>::Item<NodePtr<V>>) -> bool,
+    {
+        let (col, root) = node_mut.into_inner();
+        let iter_ptr = BfsIterPtr::<V, O::Enumeration, _>::from((storage, root));
+        unsafe { BfsIterIntoFiltered::<V, M, P, _, _, _>::from((col, iter_ptr, root), filter) }
+    }
 }