nb_tree/tree/iter/
sortable.rs

1use itertools::Itertools;
2
3use crate::{
4    prelude::{Path, Tree},
5    tree::NodeIDX,
6};
7use std::{any::Any, hash::Hash, iter::FusedIterator, marker::PhantomData, mem};
8
9use super::depth::{Traversal, TreeIterTarget};
10
11pub struct IterSortable<'a, N, B, T, S>
12where
13    T: TreeIterTarget<'a, N, B>,
14    S: FnMut(&mut Vec<(&'a B, usize)>),
15{
16    tree: &'a Tree<N, B>,
17    root: Option<NodeIDX>,
18    stages: Vec<Vec<(&'a B, usize)>>,
19    up: usize,
20    sort: S,
21    _phantom: PhantomData<T>,
22}
23
24impl<'a, N, B, T, S> IterSortable<'a, N, B, T, S>
25where
26    T: TreeIterTarget<'a, N, B>,
27    S: FnMut(&mut Vec<(&'a B, usize)>),
28{
29    pub fn new(tree: &'a Tree<N, B>, sort: S) -> Self {
30        Self {
31            tree,
32            root: tree.get_root_idx(),
33            stages: vec![],
34            up: 0,
35            sort,
36            _phantom: PhantomData,
37        }
38    }
39}
40
41impl<'a, N, B, T> IterSortable<'a, N, B, T, fn(&mut Vec<(&'a B, usize)>)>
42where
43    T: TreeIterTarget<'a, N, B>,
44{
45    pub fn new_unsorted(tree: &'a Tree<N, B>) -> Self {
46        Self {
47            tree,
48            root: tree.get_root_idx(),
49            stages: vec![],
50            up: 0,
51            sort: |_| {},
52            _phantom: PhantomData,
53        }
54    }
55
56    pub fn up(&self) -> usize {
57        self.up
58    }
59}
60
61impl<'a, N, B, T, S> IterSortable<'a, N, B, T, S>
62where
63    B: Clone + Eq + Hash,
64    T: TreeIterTarget<'a, N, B>,
65    S: FnMut(&mut Vec<(&'a B, usize)>),
66{
67    pub fn new_sub_state(
68        tree: &'a Tree<N, B>,
69        path: Path<B>,
70        sort: S,
71    ) -> Result<Self, Option<Path<B>>> {
72        Ok(Self {
73            tree,
74            root: Some(
75                tree.get_idx(&path, None)
76                    .map_err(|r| r.map(|(_, idx_p)| path.path_to(idx_p)))?,
77            ),
78            stages: vec![],
79            up: 0,
80            sort,
81            _phantom: PhantomData,
82        })
83    }
84}
85
86impl<'a, N, B, T> IterSortable<'a, N, B, T, fn(&mut Vec<(&'a B, usize)>)>
87where
88    B: Clone + Eq + Hash,
89    T: TreeIterTarget<'a, N, B>,
90{
91    pub fn new_sub_state_unsorted(
92        tree: &'a Tree<N, B>,
93        path: Path<B>,
94    ) -> Result<Self, Option<Path<B>>> {
95        Ok(Self {
96            tree,
97            root: Some(
98                tree.get_idx(&path, None)
99                    .map_err(|r| r.map(|(_, idx_p)| path.path_to(idx_p)))?,
100            ),
101            stages: vec![],
102            up: 0,
103            sort: |_| {},
104            _phantom: PhantomData,
105        })
106    }
107
108    pub fn is_last_child(&mut self) -> bool {
109        if let Some(last) = self.stages.last_mut() {
110            last.is_empty()
111        } else {
112            true
113        }
114    }
115
116    /// Tries to visit the given branch next.
117    pub fn try_follow(&mut self, branch: B) -> bool {
118        if let Some(children) = self.stages.last_mut() {
119            if let Some((idx, _)) = children.iter().find_position(|&&x| *x.0 == branch) {
120                let last = children.len() - 1;
121                children.swap(idx, last);
122                true
123            } else {
124                false
125            }
126        } else {
127            false
128        }
129    }
130}
131
132impl<'a, N, B, T, S> Iterator for IterSortable<'a, N, B, T, S>
133where
134    T: TreeIterTarget<'a, N, B>,
135    S: FnMut(&mut Vec<(&'a B, usize)>),
136{
137    type Item = Traversal<T::Target, &'a B>;
138
139    fn next(&mut self) -> Option<Self::Item> {
140        if let Some(last) = self.stages.last_mut() {
141            // Iterate on a child (not the root)
142            if let Some((branch, node_idx)) = last.pop() {
143                // Go down to a child
144                let node = &self.tree.nodes[node_idx];
145                self.stages
146                    .push(node.children.iter().map(|(b, i)| (b, *i)).collect());
147                //self.path.push_leaf(branch.clone());
148                Some(Traversal::Step {
149                    up: std::mem::replace(&mut self.up, 0),
150                    branch,
151                    data: T::target(self.tree, node_idx),
152                })
153                //Some((self.path.clone(), &node.value))
154            } else {
155                // Move back up to the parent
156                self.stages.pop();
157                //self.path.pop_leaf();
158                self.up += 1;
159                self.next()
160            }
161        } else if let Some(root) = self.root.take() {
162            let node = &self.tree.nodes[root];
163            self.stages
164                .push(node.children.iter().map(|(b, i)| (b, *i)).collect());
165            Some(Traversal::Start(T::target(self.tree, root)))
166        } else {
167            None
168        }
169    }
170}
171
172impl<'a, N, B, T, S> FusedIterator for IterSortable<'a, N, B, T, S>
173where
174    T: TreeIterTarget<'a, N, B>,
175    S: FnMut(&mut Vec<(&'a B, usize)>),
176{
177}
nb_tree/tree/iter/sortable.rs

nb_tree/tree/iter/
sortable.rs
