indextree/
traverse.rs

1//! Iterators.
2
3#![allow(clippy::redundant_closure_call)]
4
5use crate::{Arena, Node, NodeId};
6
7#[derive(Clone)]
8struct Iter<'a, T> {
9    arena: &'a Arena<T>,
10    node: Option<NodeId>,
11}
12
13impl<'a, T> Iter<'a, T> {
14    fn new(arena: &'a Arena<T>, node: impl Into<Option<NodeId>>) -> Self {
15        let node = node.into();
16
17        Self { arena, node }
18    }
19}
20
21#[derive(Clone)]
22struct DoubleEndedIter<'a, T> {
23    arena: &'a Arena<T>,
24    head: Option<NodeId>,
25    tail: Option<NodeId>,
26}
27
28impl<'a, T> DoubleEndedIter<'a, T> {
29    fn new(
30        arena: &'a Arena<T>,
31        head: impl Into<Option<NodeId>>,
32        tail: impl Into<Option<NodeId>>,
33    ) -> Self {
34        let head = head.into();
35        let tail = tail.into();
36
37        Self { arena, head, tail }
38    }
39}
40
41macro_rules! new_iterator {
42    ($(#[$attr:meta])* $name:ident, inner = $inner:ident, new = $new:expr $(,)?) => {
43        $(#[$attr])*
44        #[repr(transparent)]
45        #[derive(Clone)]
46        pub struct $name<'a, T>($inner<'a, T>);
47
48        #[allow(deprecated)]
49        impl<'a, T> $name<'a, T> {
50            pub(crate) fn new(arena: &'a Arena<T>, node: NodeId) -> Self {
51                let new: fn(&'a Arena<T>, NodeId) -> $inner<'a, T> = $new;
52                Self(new(arena, node))
53            }
54        }
55    };
56    ($(#[$attr:meta])* $name:ident, new = $new:expr, next = $next:expr $(,)?) => {
57        new_iterator!(
58            $(#[$attr])*
59            $name,
60            inner = Iter,
61            new = $new,
62        );
63
64        #[allow(deprecated)]
65        impl<'a, T> Iterator for $name<'a, T> {
66            type Item = NodeId;
67
68            fn next(&mut self) -> Option<NodeId> {
69                let next: fn(&Node<T>) -> Option<NodeId> = $next;
70
71                let node = self.0.node.take()?;
72                self.0.node = next(&self.0.arena[node]);
73                Some(node)
74            }
75        }
76
77        #[allow(deprecated)]
78        impl<'a, T> core::iter::FusedIterator for $name<'a, T> {}
79    };
80    ($(#[$attr:meta])* $name:ident, new = $new:expr, next = $next:expr, next_back = $next_back:expr $(,)?) => {
81        new_iterator!(
82            $(#[$attr])*
83            $name,
84            inner = DoubleEndedIter,
85            new = $new,
86        );
87
88        #[allow(deprecated)]
89        impl<'a, T> Iterator for $name<'a, T> {
90            type Item = NodeId;
91
92            fn next(&mut self) -> Option<NodeId> {
93                match (self.0.head, self.0.tail) {
94                    (Some(head), Some(tail)) if head == tail => {
95                        let result = head;
96                        self.0.head = None;
97                        self.0.tail = None;
98                        Some(result)
99                    }
100                    (Some(head), None) | (Some(head), Some(_)) => {
101                        let next: fn(&Node<T>) -> Option<NodeId> = $next;
102
103                        self.0.head = next(&self.0.arena[head]);
104                        Some(head)
105                    }
106                    (None, Some(_)) | (None, None) => None,
107                }
108            }
109        }
110
111        #[allow(deprecated)]
112        impl<'a, T> ::core::iter::DoubleEndedIterator for $name<'a, T> {
113            fn next_back(&mut self) -> Option<Self::Item> {
114                match (self.0.head, self.0.tail) {
115                    (Some(head), Some(tail)) if head == tail => {
116                        let result = head;
117                        self.0.head = None;
118                        self.0.tail = None;
119                        Some(result)
120                    }
121                    (None, Some(tail)) | (Some(_), Some(tail)) => {
122                        let next_back: fn(&Node<T>) -> Option<NodeId> = $next_back;
123
124                        self.0.tail = next_back(&self.0.arena[tail]);
125                        Some(tail)
126                    }
127                    (Some(_), None)| (None, None) => None,
128                }
129            }
130        }
131    };
132    ($(#[$attr:meta])* $name:ident, next = $next:expr $(,)?) => {
133        new_iterator!(
134            $(#[$attr])*
135            $name,
136            new = |arena, node| Iter::new(arena, node),
137            next = $next,
138        );
139    };
140    ($(#[$attr:meta])* $name:ident, next = $next:expr, next_back = $next_back:expr $(,)?) => {
141        new_iterator!(
142            $(#[$attr])*
143            $name,
144            new = |arena, node| DoubleEndedIter::new(arena, node, None),
145            next = $next,
146            next_back = $next_back,
147        );
148    };
149}
150
151new_iterator!(
152    /// An iterator of the IDs of the ancestors of a given node.
153    Ancestors,
154    next = |node| node.parent,
155);
156
157new_iterator!(
158    /// An iterator of the IDs of the predecessors of a given node.
159    Predecessors,
160    next = |node| node.previous_sibling.or(node.parent),
161);
162
163new_iterator!(
164    /// An iterator of the IDs of the siblings before a given node.
165    PrecedingSiblings,
166    new = |arena, node| {
167        let first = arena
168            .get(node)
169            .unwrap()
170            .parent
171            .and_then(|parent_id| arena.get(parent_id))
172            .and_then(|parent| parent.first_child);
173
174        DoubleEndedIter::new(arena, node, first)
175    },
176    next = |head| head.previous_sibling,
177    next_back = |tail| tail.next_sibling,
178);
179
180new_iterator!(
181    /// An iterator of the IDs of the siblings after a given node.
182    FollowingSiblings,
183    new = |arena, node| {
184        let last = arena
185            .get(node)
186            .unwrap()
187            .parent
188            .and_then(|parent_id| arena.get(parent_id))
189            .and_then(|parent| parent.last_child);
190
191        DoubleEndedIter::new(arena, node, last)
192    },
193    next = |head| head.next_sibling,
194    next_back = |tail| tail.previous_sibling,
195);
196
197new_iterator!(
198    /// An iterator of the IDs of the children of a given node, in insertion order.
199    Children,
200    new = |arena, node| DoubleEndedIter::new(arena, arena[node].first_child, arena[node].last_child),
201    next = |node| node.next_sibling,
202    next_back = |tail| tail.previous_sibling,
203);
204
205new_iterator!(
206    #[deprecated(
207        since = "4.7.0",
208        note = "please, use Children::rev() instead if you want to iterate in reverse"
209    )]
210    /// An iterator of the IDs of the children of a given node, in reverse insertion order.
211    ReverseChildren,
212    new = |arena, node| Iter::new(arena, arena[node].last_child),
213    next = |node| node.previous_sibling,
214);
215
216#[derive(Clone)]
217/// An iterator of the IDs of a given node and its descendants, as a pre-order depth-first search where children are visited in insertion order.
218///
219/// i.e. node -> first child -> second child
220pub struct Descendants<'a, T>(Traverse<'a, T>);
221
222impl<'a, T> Descendants<'a, T> {
223    pub(crate) fn new(arena: &'a Arena<T>, current: NodeId) -> Self {
224        Self(Traverse::new(arena, current))
225    }
226}
227
228impl<T> Iterator for Descendants<'_, T> {
229    type Item = NodeId;
230
231    fn next(&mut self) -> Option<NodeId> {
232        self.0.find_map(|edge| match edge {
233            NodeEdge::Start(node) => Some(node),
234            NodeEdge::End(_) => None,
235        })
236    }
237}
238
239impl<T> core::iter::FusedIterator for Descendants<'_, T> {}
240
241#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
242/// Indicator if the node is at a start or endpoint of the tree
243pub enum NodeEdge {
244    /// Indicates that start of a node that has children.
245    ///
246    /// Yielded by `Traverse::next()` before the node’s descendants. In HTML or
247    /// XML, this corresponds to an opening tag like `<div>`.
248    Start(NodeId),
249
250    /// Indicates that end of a node that has children.
251    ///
252    /// Yielded by `Traverse::next()` after the node’s descendants. In HTML or
253    /// XML, this corresponds to a closing tag like `</div>`
254    End(NodeId),
255}
256
257impl NodeEdge {
258    /// Returns the next `NodeEdge` to be returned by forward depth-first traversal.
259    ///
260    /// # Examples
261    ///
262    /// ```
263    /// # use indextree::{Arena, NodeEdge};
264    /// # let mut arena = Arena::new();
265    /// # let n1 = arena.new_node("1");
266    /// # let n1_1 = arena.new_node("1_1");
267    /// # n1.append(n1_1, &mut arena);
268    /// # let n1_1_1 = arena.new_node("1_1_1");
269    /// # n1_1.append(n1_1_1, &mut arena);
270    /// # let n1_2 = arena.new_node("1_2");
271    /// # n1.append(n1_2, &mut arena);
272    /// # let n1_3 = arena.new_node("1_3");
273    /// # n1.append(n1_3, &mut arena);
274    /// #
275    /// // arena
276    /// // `-- 1
277    /// //     |-- 1_1
278    /// //     |   `-- 1_1_1
279    /// //     |-- 1_2
280    /// //     `-- 1_3
281    ///
282    /// let steps = std::iter::successors(
283    ///     Some(NodeEdge::Start(n1)),
284    ///     |current| current.next_traverse(&arena)
285    /// )
286    ///     .collect::<Vec<_>>();
287    /// let traversed_by_iter = n1.traverse(&arena).collect::<Vec<_>>();
288    /// assert_eq!(
289    ///     steps,
290    ///     traversed_by_iter,
291    ///     "repeated `.next_traverse()`s emit same events as `NodeId::traverse()` iterator"
292    /// );
293    /// ```
294    ///
295    /// `NodeEdge` itself does not borrow an arena, so you can modify the nodes
296    /// being traversed.
297    ///
298    /// ```
299    /// # use indextree::{Arena, NodeEdge};
300    /// # let mut arena = Arena::new();
301    /// # let n1 = arena.new_node("1".to_owned());
302    /// # let n1_1 = arena.new_node("1_1".to_owned());
303    /// # n1.append(n1_1, &mut arena);
304    /// # let n1_1_1 = arena.new_node("1_1_1".to_owned());
305    /// # n1_1.append(n1_1_1, &mut arena);
306    /// # let n1_2 = arena.new_node("1_2".to_owned());
307    /// # n1.append(n1_2, &mut arena);
308    /// # let n1_3 = arena.new_node("1_3".to_owned());
309    /// # n1.append(n1_3, &mut arena);
310    /// #
311    /// // arena: Arena<String>
312    /// // `-- 1
313    /// //     |-- 1_1
314    /// //     |   `-- 1_1_1
315    /// //     |-- 1_2
316    /// //     `-- 1_3
317    ///
318    /// assert_eq!(*arena[n1].get(), "1");
319    /// assert_eq!(*arena[n1_1_1].get(), "1_1_1");
320    /// assert_eq!(*arena[n1_3].get(), "1_3");
321    ///
322    /// let mut next = Some(NodeEdge::Start(n1));
323    /// let mut count = 0;
324    /// while let Some(current) = next {
325    ///     next = current.next_traverse(&arena);
326    ///     let current = match current {
327    ///         NodeEdge::Start(id) => id,
328    ///         NodeEdge::End(_) => continue,
329    ///     };
330    ///
331    ///     arena[current].get_mut().push_str(&format!(" (count={})", count));
332    ///     count += 1;
333    /// }
334    ///
335    /// assert_eq!(*arena[n1].get(), "1 (count=0)");
336    /// assert_eq!(*arena[n1_1_1].get(), "1_1_1 (count=2)");
337    /// assert_eq!(*arena[n1_3].get(), "1_3 (count=4)");
338    /// ```
339    #[must_use]
340    pub fn next_traverse<T>(self, arena: &Arena<T>) -> Option<Self> {
341        match self {
342            NodeEdge::Start(node) => match arena[node].first_child {
343                Some(first_child) => Some(NodeEdge::Start(first_child)),
344                None => Some(NodeEdge::End(node)),
345            },
346            NodeEdge::End(node) => {
347                let node = &arena[node];
348                match node.next_sibling {
349                    Some(next_sibling) => Some(NodeEdge::Start(next_sibling)),
350                    // `node.parent()` here can only be `None` if the tree has
351                    // been modified during iteration, but silently stoping
352                    // iteration seems a more sensible behavior than panicking.
353                    None => node.parent.map(NodeEdge::End),
354                }
355            }
356        }
357    }
358
359    /// Returns the previous `NodeEdge` to be returned by forward depth-first traversal.
360    ///
361    /// # Examples
362    ///
363    /// ```
364    /// # use indextree::{Arena, NodeEdge};
365    /// # let mut arena = Arena::new();
366    /// # let n1 = arena.new_node("1");
367    /// # let n1_1 = arena.new_node("1_1");
368    /// # n1.append(n1_1, &mut arena);
369    /// # let n1_1_1 = arena.new_node("1_1_1");
370    /// # n1_1.append(n1_1_1, &mut arena);
371    /// # let n1_2 = arena.new_node("1_2");
372    /// # n1.append(n1_2, &mut arena);
373    /// # let n1_3 = arena.new_node("1_3");
374    /// # n1.append(n1_3, &mut arena);
375    /// #
376    /// // arena
377    /// // `-- 1
378    /// //     |-- 1_1
379    /// //     |   `-- 1_1_1
380    /// //     |-- 1_2
381    /// //     `-- 1_3
382    ///
383    /// let steps = std::iter::successors(
384    ///     Some(NodeEdge::End(n1)),
385    ///     |current| current.prev_traverse(&arena)
386    /// )
387    ///     .collect::<Vec<_>>();
388    /// let traversed_by_iter = n1.reverse_traverse(&arena).collect::<Vec<_>>();
389    /// assert_eq!(
390    ///     steps,
391    ///     traversed_by_iter,
392    ///     "repeated `.prev_traverse()`s emit same events as \
393    ///      `NodeId::reverse_traverse()` iterator"
394    /// );
395    /// ```
396    ///
397    /// `NodeEdge` itself does not borrow an arena, so you can modify the nodes
398    /// being traversed.
399    ///
400    /// ```
401    /// use indextree::{Arena, NodeEdge};
402    ///
403    /// # let mut arena = Arena::new();
404    /// # let n1 = arena.new_node("1".to_owned());
405    /// # let n1_1 = arena.new_node("1_1".to_owned());
406    /// # n1.append(n1_1, &mut arena);
407    /// # let n1_1_1 = arena.new_node("1_1_1".to_owned());
408    /// # n1_1.append(n1_1_1, &mut arena);
409    /// # let n1_2 = arena.new_node("1_2".to_owned());
410    /// # n1.append(n1_2, &mut arena);
411    /// # let n1_3 = arena.new_node("1_3".to_owned());
412    /// # n1.append(n1_3, &mut arena);
413    /// #
414    /// // arena: Arena<String>
415    /// // `-- 1
416    /// //     |-- 1_1
417    /// //     |   `-- 1_1_1
418    /// //     |-- 1_2
419    /// //     `-- 1_3
420    ///
421    /// assert_eq!(*arena[n1_3].get(), "1_3");
422    /// assert_eq!(*arena[n1_1_1].get(), "1_1_1");
423    /// assert_eq!(*arena[n1].get(), "1");
424    ///
425    /// let mut next = Some(NodeEdge::End(n1_3));
426    /// let mut count = 0;
427    /// while let Some(current) = next {
428    ///     next = current.prev_traverse(&arena);
429    ///     let current = match current {
430    ///         NodeEdge::Start(id) => id,
431    ///         NodeEdge::End(_) => continue,
432    ///     };
433    ///
434    ///     arena[current].get_mut().push_str(&format!(" (count={})", count));
435    ///     count += 1;
436    /// }
437    ///
438    /// assert_eq!(*arena[n1_3].get(), "1_3 (count=0)");
439    /// assert_eq!(*arena[n1_1_1].get(), "1_1_1 (count=2)");
440    /// assert_eq!(*arena[n1].get(), "1 (count=4)");
441    /// ```
442    #[must_use]
443    pub fn prev_traverse<T>(self, arena: &Arena<T>) -> Option<Self> {
444        match self {
445            NodeEdge::End(node) => match arena[node].last_child {
446                Some(last_child) => Some(NodeEdge::End(last_child)),
447                None => Some(NodeEdge::Start(node)),
448            },
449            NodeEdge::Start(node) => {
450                let node = &arena[node];
451                match node.previous_sibling {
452                    Some(previous_sibling) => Some(NodeEdge::End(previous_sibling)),
453                    // `node.parent()` here can only be `None` if the tree has
454                    // been modified during iteration, but silently stopping
455                    // iteration seems a more sensible behavior than panicking.
456                    None => node.parent.map(NodeEdge::Start),
457                }
458            }
459        }
460    }
461}
462
463#[derive(Clone)]
464/// An iterator of the "sides" of a node visited during a depth-first pre-order traversal,
465/// where node sides are visited start to end and children are visited in insertion order.
466///
467/// i.e. node.start -> first child -> second child -> node.end
468pub struct Traverse<'a, T> {
469    arena: &'a Arena<T>,
470    root: NodeId,
471    next: Option<NodeEdge>,
472}
473
474impl<'a, T> Traverse<'a, T> {
475    pub(crate) fn new(arena: &'a Arena<T>, current: NodeId) -> Self {
476        Self {
477            arena,
478            root: current,
479            next: Some(NodeEdge::Start(current)),
480        }
481    }
482
483    /// Calculates the next node.
484    fn next_of_next(&self, next: NodeEdge) -> Option<NodeEdge> {
485        if next == NodeEdge::End(self.root) {
486            return None;
487        }
488        next.next_traverse(self.arena)
489    }
490
491    /// Returns a reference to the arena.
492    #[inline]
493    #[must_use]
494    pub(crate) fn arena(&self) -> &Arena<T> {
495        self.arena
496    }
497}
498
499impl<T> Iterator for Traverse<'_, T> {
500    type Item = NodeEdge;
501
502    fn next(&mut self) -> Option<NodeEdge> {
503        let next = self.next.take()?;
504        self.next = self.next_of_next(next);
505        Some(next)
506    }
507}
508
509impl<T> core::iter::FusedIterator for Traverse<'_, T> {}
510
511#[derive(Clone)]
512/// An iterator of the "sides" of a node visited during a depth-first pre-order traversal,
513/// where nodes are visited end to start and children are visited in reverse insertion order.
514///
515/// i.e. node.end -> second child -> first child -> node.start
516pub struct ReverseTraverse<'a, T> {
517    arena: &'a Arena<T>,
518    root: NodeId,
519    next: Option<NodeEdge>,
520}
521
522impl<'a, T> ReverseTraverse<'a, T> {
523    pub(crate) fn new(arena: &'a Arena<T>, current: NodeId) -> Self {
524        Self {
525            arena,
526            root: current,
527            next: Some(NodeEdge::End(current)),
528        }
529    }
530
531    /// Calculates the next node.
532    fn next_of_next(&self, next: NodeEdge) -> Option<NodeEdge> {
533        if next == NodeEdge::Start(self.root) {
534            return None;
535        }
536        next.prev_traverse(self.arena)
537    }
538}
539
540impl<T> Iterator for ReverseTraverse<'_, T> {
541    type Item = NodeEdge;
542
543    fn next(&mut self) -> Option<NodeEdge> {
544        let next = self.next.take()?;
545        self.next = self.next_of_next(next);
546        Some(next)
547    }
548}
549
550impl<T> core::iter::FusedIterator for ReverseTraverse<'_, T> {}
indextree/traverse.rs

indextree/
traverse.rs
