nb_tree/tree/iter/
depth.rs

1use std::collections::hash_map::Iter as HMIter;
2use std::{hash::Hash, iter::FusedIterator, marker::PhantomData};
3
4use crate::{path::Path, prelude::DiffNode, tree::node::TreeNode};
5
6use super::super::{NodeIDX, Tree};
7
8///////////////////////////////////////////////////////////////////////////////
9//                                   Iter                                    //
10///////////////////////////////////////////////////////////////////////////////
11
12/// Describes a node's position relatively to the previous one and contains some data.
13#[derive(Debug, PartialEq)]
14pub enum Traversal<N, B> {
15    Start(N),
16    Step {
17        /// How much higher the parent of the described node is compared to the previous node
18        up: usize,
19        /// The branch leading from the parent to the described node
20        branch: B,
21        /// Node data
22        data: N,
23    },
24}
25
26impl<N, B> Traversal<N, B> {
27    pub fn take(self) -> N {
28        match self {
29            Self::Start(data) => data,
30            Self::Step {
31                up: _,
32                branch: _,
33                data,
34            } => data,
35        }
36    }
37
38    pub fn map<M>(self, op: impl Fn(N) -> M) -> Traversal<M, B> {
39        match self {
40            Self::Start(data) => Traversal::Start(op(data)),
41            Self::Step { up, branch, data } => Traversal::Step {
42                up,
43                branch,
44                data: op(data),
45            },
46        }
47    }
48    pub fn data(&self) -> &N {
49        match self {
50            Self::Start(data) => data,
51            Self::Step {
52                up: _,
53                branch: _,
54                data,
55            } => data,
56        }
57    }
58    pub fn branch(&self) -> Option<&B> {
59        match self {
60            Self::Start(_) => None,
61            Self::Step {
62                up: _,
63                branch,
64                data: _,
65            } => Some(branch),
66        }
67    }
68    pub fn up(&self) -> usize {
69        match self {
70            Self::Start(_) => 0,
71            Self::Step {
72                up,
73                branch: _,
74                data: _,
75            } => *up,
76        }
77    }
78    pub(crate) fn truncate_up(&mut self, cut: usize) {
79        if let Self::Step { up, branch, data } = self {
80            *up -= cut;
81        }
82    }
83
84    pub fn is_root(&self) -> bool {
85        match self {
86            Traversal::Start(_) => true,
87            Traversal::Step { .. } => false,
88        }
89    }
90}
91
92///////////////////////////////////////////////////////////////////////////////
93
94pub trait TreeIterTarget<'a, N, B> {
95    type Target;
96    fn target(tree: &'a Tree<N, B>, idx: NodeIDX) -> Self::Target;
97}
98
99impl<'a, N, B> TreeIterTarget<'a, N, B> for &'a TreeNode<N, B> {
100    type Target = Self;
101    fn target(tree: &'a Tree<N, B>, idx: NodeIDX) -> Self::Target {
102        &tree.nodes[idx]
103    }
104}
105
106impl<'a, N, B> TreeIterTarget<'a, N, B> for &'a N {
107    type Target = Self;
108    fn target(tree: &'a Tree<N, B>, idx: NodeIDX) -> Self::Target {
109        &tree.nodes[idx].value
110    }
111}
112
113impl<'a, N, B> TreeIterTarget<'a, N, B> for () {
114    type Target = Self;
115    fn target(_: &'a Tree<N, B>, _: NodeIDX) -> Self::Target {}
116}
117
118impl<'a, N, B> TreeIterTarget<'a, N, B> for NodeIDX {
119    type Target = Self;
120    fn target(_: &'a Tree<N, B>, idx: NodeIDX) -> Self::Target {
121        idx
122    }
123}
124
125pub struct Iter<'a, N, B, T>
126where
127    T: TreeIterTarget<'a, N, B>,
128{
129    tree: &'a Tree<N, B>,
130    root: Option<NodeIDX>,
131    stages: Vec<HMIter<'a, B, usize>>,
132    up: usize,
133    _phantom: PhantomData<T>,
134    //path: Path<B>,
135}
136
137impl<'a, N, B, T> Iter<'a, N, B, T>
138where
139    T: TreeIterTarget<'a, N, B>,
140{
141    pub fn new(tree: &'a Tree<N, B>) -> Self {
142        Self {
143            tree,
144            root: tree.get_root_idx(),
145            stages: vec![],
146            up: 0,
147            _phantom: PhantomData,
148            //path: Path::new(),
149        }
150    }
151
152    pub fn up(&self) -> usize {
153        self.up
154    }
155}
156
157impl<'a, N, B, T> Iter<'a, N, B, T>
158where
159    B: Clone + Eq + Hash,
160    T: TreeIterTarget<'a, N, B>,
161{
162    pub fn new_sub_state(tree: &'a Tree<N, B>, path: Path<B>) -> Result<Self, Option<Path<B>>> {
163        Ok(Self {
164            tree,
165            root: Some(
166                tree.get_idx(&path, None)
167                    .map_err(|r| r.map(|(_, idx_p)| path.path_to(idx_p)))?,
168            ),
169            stages: vec![],
170            up: 0,
171            _phantom: PhantomData,
172            //path,
173        })
174    }
175
176    pub fn is_last_child(&mut self) -> bool {
177        if let Some(last) = self.stages.last_mut() {
178            last.peekable().peek().is_none()
179        } else {
180            true
181        }
182    }
183}
184//Vec<HMIterator<B, TreeNode<N, B>>>,
185//);
186
187impl<'a, N, B, T> Iterator for Iter<'a, N, B, T>
188where
189    T: TreeIterTarget<'a, N, B>,
190{
191    type Item = Traversal<T::Target, &'a B>;
192
193    fn next(&mut self) -> Option<Self::Item> {
194        if let Some(last) = self.stages.last_mut() {
195            // Iterate on a child (not the root)
196            if let Some((branch, node_idx)) = last.next() {
197                // Go down to a child
198                let node = &self.tree.nodes[*node_idx];
199                self.stages.push(node.children.iter());
200                //self.path.push_leaf(branch.clone());
201                Some(Traversal::Step {
202                    up: std::mem::replace(&mut self.up, 0),
203                    branch,
204                    data: T::target(self.tree, *node_idx),
205                })
206                //Some((self.path.clone(), &node.value))
207            } else {
208                // Move back up to the parent
209                self.stages.pop();
210                //self.path.pop_leaf();
211                self.up += 1;
212                self.next()
213            }
214        } else if let Some(root) = self.root.take() {
215            let node = &self.tree.nodes[root];
216            self.stages.push(node.children.iter());
217            Some(Traversal::Start(T::target(self.tree, root)))
218        } else {
219            None
220        }
221    }
222}
223
224impl<'a, N, B, T> FusedIterator for Iter<'a, N, B, T> where T: TreeIterTarget<'a, N, B> {}
225
226impl<'a, N, B> IntoIterator for &'a Tree<N, B> {
227    type Item = Traversal<&'a N, &'a B>;
228
229    type IntoIter = Iter<'a, N, B, &'a N>;
230
231    fn into_iter(self) -> Self::IntoIter {
232        Self::IntoIter::new(self)
233    }
234}
235
236///////////////////////////////////////////////////////////////////////////////
237//                                  IterMut                                  //
238///////////////////////////////////////////////////////////////////////////////
239
240//TODO
241
242///////////////////////////////////////////////////////////////////////////////
243//                                 IntoIter                                  //
244///////////////////////////////////////////////////////////////////////////////
245
246pub trait TreeIntoIterTarget<N, B> {
247    type Target;
248    fn target(node: TreeNode<N, B>) -> Self::Target;
249}
250
251impl<N, B> TreeIntoIterTarget<N, B> for N {
252    type Target = Self;
253    fn target(node: TreeNode<N, B>) -> Self::Target {
254        node.value
255    }
256}
257
258impl<N, B> TreeIntoIterTarget<N, B> for TreeNode<N, B> {
259    type Target = Self;
260    fn target(node: TreeNode<N, B>) -> Self::Target {
261        node
262    }
263}
264
265/// Tree iterator
266pub struct IntoIter<N, B, T>
267where
268    B: Clone,
269    T: TreeIntoIterTarget<N, B>,
270{
271    idxs: Vec<Traversal<NodeIDX, B>>,
272    nodes: Vec<Option<TreeNode<N, B>>>,
273    _phantom: PhantomData<T>,
274}
275
276impl<N, B, T> IntoIter<N, B, T>
277where
278    B: Clone,
279    T: TreeIntoIterTarget<N, B>,
280{
281    pub fn new(tree: Tree<N, B>) -> Self {
282        let mut s = Self {
283            //TODO: Remove B: Clone
284            idxs: Iter::<N, B, NodeIDX>::new(&tree)
285                .map(|n| match n {
286                    Traversal::Start(idx) => Traversal::Start(idx),
287                    Traversal::Step { up, branch, data } => Traversal::Step {
288                        up,
289                        branch: branch.clone(),
290                        data,
291                    },
292                })
293                .collect(),
294            nodes: tree.nodes.into_iter().map(Some).collect(),
295            _phantom: PhantomData,
296        };
297        // The vector's last value will be popped on each iteration
298        s.idxs.reverse();
299        s
300    }
301}
302
303impl<N, B, T> Iterator for IntoIter<N, B, T>
304where
305    B: Clone,
306    T: TreeIntoIterTarget<N, B>,
307{
308    type Item = Traversal<T::Target, B>;
309
310    fn next(&mut self) -> Option<Self::Item> {
311        match self.idxs.pop() {
312            Some(Traversal::Start(idx)) => {
313                Some(Traversal::Start(T::target(self.nodes[idx].take().unwrap())))
314            }
315            Some(Traversal::Step { up, branch, data }) => Some(Traversal::Step {
316                up,
317                branch,
318                data: T::target(self.nodes[data].take().unwrap()),
319            }),
320            None => None,
321        }
322    }
323}
324
325impl<N, B> IntoIterator for Tree<N, B>
326where
327    B: Clone,
328{
329    type Item = Traversal<N, B>;
330
331    type IntoIter = IntoIter<N, B, N>;
332
333    fn into_iter(self) -> Self::IntoIter {
334        Self::IntoIter::new(self)
335    }
336}
nb_tree/tree/iter/depth.rs

nb_tree/tree/iter/
depth.rs
