Skip to main content

fhp_tree/
traverse.rs

1//! Allocation-free tree traversal iterators.
2//!
3//! All iterators borrow the arena and produce [`NodeId`] values.
4//! [`DepthFirst`] and [`Children`] are fully allocation-free.
5//! [`BreadthFirst`] uses a [`VecDeque`](std::collections::VecDeque) internally.
6
7use std::collections::VecDeque;
8
9use crate::arena::Arena;
10use crate::node::NodeId;
11
12/// Pre-order depth-first traversal, allocation-free.
13///
14/// Walks the subtree rooted at `root` using tree links (no separate stack).
15pub struct DepthFirst<'a> {
16    arena: &'a Arena,
17    current: NodeId,
18    root: NodeId,
19}
20
21impl<'a> DepthFirst<'a> {
22    /// Create a new depth-first iterator starting at `root`.
23    pub fn new(arena: &'a Arena, root: NodeId) -> Self {
24        Self {
25            arena,
26            current: root,
27            root,
28        }
29    }
30
31    /// Find the next node after `node` that is not a descendant.
32    fn next_non_child(&self, node: NodeId) -> NodeId {
33        let mut cur = node;
34        loop {
35            let n = self.arena.get(cur);
36            if !n.next_sibling.is_null() {
37                return n.next_sibling;
38            }
39            if n.parent.is_null() || n.parent == self.root {
40                return NodeId::NULL;
41            }
42            cur = n.parent;
43        }
44    }
45}
46
47impl<'a> Iterator for DepthFirst<'a> {
48    type Item = NodeId;
49
50    #[inline]
51    fn next(&mut self) -> Option<NodeId> {
52        if self.current.is_null() {
53            return None;
54        }
55
56        let result = self.current;
57        let node = self.arena.get(result);
58
59        // Try first child, then next sibling, then ancestor's sibling.
60        self.current = if !node.first_child.is_null() {
61            node.first_child
62        } else {
63            self.next_non_child(result)
64        };
65
66        Some(result)
67    }
68}
69
70/// Breadth-first (level-order) traversal.
71///
72/// Uses a [`VecDeque`] internally — allocation is unavoidable for BFS.
73pub struct BreadthFirst<'a> {
74    arena: &'a Arena,
75    queue: VecDeque<NodeId>,
76}
77
78impl<'a> BreadthFirst<'a> {
79    /// Create a new breadth-first iterator starting at `root`.
80    pub fn new(arena: &'a Arena, root: NodeId) -> Self {
81        let mut queue = VecDeque::new();
82        if !root.is_null() {
83            queue.push_back(root);
84        }
85        Self { arena, queue }
86    }
87}
88
89impl<'a> Iterator for BreadthFirst<'a> {
90    type Item = NodeId;
91
92    fn next(&mut self) -> Option<NodeId> {
93        let id = self.queue.pop_front()?;
94        let node = self.arena.get(id);
95
96        // Enqueue all children.
97        let mut child = node.first_child;
98        while !child.is_null() {
99            self.queue.push_back(child);
100            child = self.arena.get(child).next_sibling;
101        }
102
103        Some(id)
104    }
105}
106
107/// Iterate over the direct children of a node, allocation-free.
108pub struct Children<'a> {
109    arena: &'a Arena,
110    current: NodeId,
111}
112
113impl<'a> Children<'a> {
114    /// Create an iterator over children of `parent`.
115    pub fn new(arena: &'a Arena, parent: NodeId) -> Self {
116        let first = arena.get(parent).first_child;
117        Self {
118            arena,
119            current: first,
120        }
121    }
122}
123
124impl<'a> Iterator for Children<'a> {
125    type Item = NodeId;
126
127    #[inline]
128    fn next(&mut self) -> Option<NodeId> {
129        if self.current.is_null() {
130            return None;
131        }
132        let result = self.current;
133        self.current = self.arena.get(result).next_sibling;
134        Some(result)
135    }
136}
137
138/// Iterate over ancestors of a node (parent chain), allocation-free.
139///
140/// Does **not** include the starting node itself.
141pub struct Ancestors<'a> {
142    arena: &'a Arena,
143    current: NodeId,
144}
145
146impl<'a> Ancestors<'a> {
147    /// Create an iterator over ancestors of `node`.
148    pub fn new(arena: &'a Arena, node: NodeId) -> Self {
149        let parent = arena.get(node).parent;
150        Self {
151            arena,
152            current: parent,
153        }
154    }
155}
156
157impl<'a> Iterator for Ancestors<'a> {
158    type Item = NodeId;
159
160    #[inline]
161    fn next(&mut self) -> Option<NodeId> {
162        if self.current.is_null() {
163            return None;
164        }
165        let result = self.current;
166        self.current = self.arena.get(result).parent;
167        Some(result)
168    }
169}
170
171/// Iterate over next siblings of a node, allocation-free.
172///
173/// Does **not** include the starting node itself.
174pub struct Siblings<'a> {
175    arena: &'a Arena,
176    current: NodeId,
177}
178
179impl<'a> Siblings<'a> {
180    /// Create an iterator over next siblings of `node`.
181    pub fn new(arena: &'a Arena, node: NodeId) -> Self {
182        let next = arena.get(node).next_sibling;
183        Self {
184            arena,
185            current: next,
186        }
187    }
188}
189
190impl<'a> Iterator for Siblings<'a> {
191    type Item = NodeId;
192
193    #[inline]
194    fn next(&mut self) -> Option<NodeId> {
195        if self.current.is_null() {
196            return None;
197        }
198        let result = self.current;
199        self.current = self.arena.get(result).next_sibling;
200        Some(result)
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207    use fhp_core::tag::Tag;
208
209    fn build_test_tree() -> (Arena, NodeId) {
210        // <div>
211        //   <span>hello</span>
212        //   <p>world</p>
213        // </div>
214        let mut arena = Arena::new();
215        let root = arena.new_element(Tag::Unknown, 0);
216        let div = arena.new_element(Tag::Div, 1);
217        arena.append_child(root, div);
218
219        let span = arena.new_element(Tag::Span, 2);
220        arena.append_child(div, span);
221        let hello = arena.new_text(3, "hello");
222        arena.append_child(span, hello);
223
224        let p = arena.new_element(Tag::P, 2);
225        arena.append_child(div, p);
226        let world = arena.new_text(3, "world");
227        arena.append_child(p, world);
228
229        (arena, root)
230    }
231
232    #[test]
233    fn depth_first_order() {
234        let (arena, root) = build_test_tree();
235        let ids: Vec<NodeId> = DepthFirst::new(&arena, root).collect();
236
237        // root, div, span, hello, p, world
238        assert_eq!(ids.len(), 6);
239        assert_eq!(arena.get(ids[0]).tag, Tag::Unknown); // root
240        assert_eq!(arena.get(ids[1]).tag, Tag::Div);
241        assert_eq!(arena.get(ids[2]).tag, Tag::Span);
242        assert_eq!(arena.text(ids[3]), "hello");
243        assert_eq!(arena.get(ids[4]).tag, Tag::P);
244        assert_eq!(arena.text(ids[5]), "world");
245    }
246
247    #[test]
248    fn breadth_first_order() {
249        let (arena, root) = build_test_tree();
250        let ids: Vec<NodeId> = BreadthFirst::new(&arena, root).collect();
251
252        // root, div, span, p, hello, world
253        assert_eq!(ids.len(), 6);
254        assert_eq!(arena.get(ids[0]).tag, Tag::Unknown); // root
255        assert_eq!(arena.get(ids[1]).tag, Tag::Div);
256        assert_eq!(arena.get(ids[2]).tag, Tag::Span);
257        assert_eq!(arena.get(ids[3]).tag, Tag::P);
258        assert_eq!(arena.text(ids[4]), "hello");
259        assert_eq!(arena.text(ids[5]), "world");
260    }
261
262    #[test]
263    fn children_iterator() {
264        let (arena, root) = build_test_tree();
265        let div = arena.get(root).first_child;
266        let children: Vec<NodeId> = Children::new(&arena, div).collect();
267
268        assert_eq!(children.len(), 2);
269        assert_eq!(arena.get(children[0]).tag, Tag::Span);
270        assert_eq!(arena.get(children[1]).tag, Tag::P);
271    }
272
273    #[test]
274    fn ancestors_iterator() {
275        let (arena, root) = build_test_tree();
276        // hello text node is: root -> div -> span -> hello
277        let div = arena.get(root).first_child;
278        let span = arena.get(div).first_child;
279        let hello = arena.get(span).first_child;
280
281        let ancestors: Vec<NodeId> = Ancestors::new(&arena, hello).collect();
282        assert_eq!(ancestors.len(), 3); // span, div, root
283        assert_eq!(arena.get(ancestors[0]).tag, Tag::Span);
284        assert_eq!(arena.get(ancestors[1]).tag, Tag::Div);
285        assert_eq!(arena.get(ancestors[2]).tag, Tag::Unknown); // root
286    }
287
288    #[test]
289    fn siblings_iterator() {
290        let (arena, root) = build_test_tree();
291        let div = arena.get(root).first_child;
292        let span = arena.get(div).first_child;
293
294        let siblings: Vec<NodeId> = Siblings::new(&arena, span).collect();
295        assert_eq!(siblings.len(), 1);
296        assert_eq!(arena.get(siblings[0]).tag, Tag::P);
297    }
298
299    #[test]
300    fn empty_iterators() {
301        let (arena, root) = build_test_tree();
302        let div = arena.get(root).first_child;
303        let span = arena.get(div).first_child;
304        let hello = arena.get(span).first_child;
305
306        // hello has no children.
307        assert_eq!(Children::new(&arena, hello).count(), 0);
308        // hello has no siblings (it's the only child of span).
309        assert_eq!(Siblings::new(&arena, hello).count(), 0);
310    }
311}