tan/expr/
expr_iter.rs

1use std::mem;
2
3use super::Expr;
4
5impl Expr {
6    pub fn iter(&self) -> ExprIter<'_> {
7        ExprIter {
8            children: std::slice::from_ref(self),
9            parent: None,
10        }
11    }
12}
13
14// #insight
15// The iterator is implemented as a separate struct, for flexibility.
16
17// #todo support in-order, pre-order, post-order
18// #todo implement owned iterator
19// #todo implement mutable iterator
20// #todo https://aloso.github.io/2021/03/09/creating-an-iterator
21
22// #todo is this really DFS?
23/// A depth-first Expr iterator.
24#[derive(Default)]
25pub struct ExprIter<'a> {
26    children: &'a [Expr],
27    parent: Option<Box<ExprIter<'a>>>,
28}
29
30impl<'a> Iterator for ExprIter<'a> {
31    type Item = &'a Expr;
32
33    // #todo this does not traverse Array, Map, etc.
34    fn next(&mut self) -> Option<Self::Item> {
35        let Some(expr) = self.children.first() else {
36            return match self.parent.take() {
37                Some(parent) => {
38                    // continue with the parent expr
39                    *self = *parent;
40                    self.next()
41                }
42                None => None,
43            };
44        };
45
46        match expr.unpack() {
47            Expr::List(children) => {
48                self.children = &self.children[1..];
49                // iterate the sub-trees
50                *self = ExprIter {
51                    children: children.as_slice(),
52                    parent: Some(Box::new(mem::take(self))),
53                };
54                Some(expr)
55            }
56            _ => {
57                self.children = &self.children[1..];
58                Some(expr)
59            }
60        }
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use crate::{lexer::Lexer, parser::Parser};
67
68    #[test]
69    fn expr_iter_performs_depth_first_iteration() {
70        let input = "(quot (1 2 3 (4 5) (6 (+ 7 8)) 9 10))";
71
72        let mut lexer = Lexer::new(input);
73        let tokens = lexer.lex().unwrap();
74
75        let mut parser = Parser::new(&tokens);
76        let expr = parser.parse().unwrap();
77
78        let expr = &expr[0];
79
80        let terms: Vec<String> = expr.iter().map(|ax| ax.to_string()).collect();
81        let expected_terms = vec![
82            "(quot (1 2 3 (4 5) (6 (+ 7 8)) 9 10))",
83            "quot",
84            "(1 2 3 (4 5) (6 (+ 7 8)) 9 10)",
85            "1",
86            "2",
87            "3",
88            "(4 5)",
89            "4",
90            "5",
91            "(6 (+ 7 8))",
92            "6",
93            "(+ 7 8)",
94            "+",
95            "7",
96            "8",
97            "9",
98            "10",
99        ];
100        assert_eq!(terms, expected_terms);
101    }
102}