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#[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 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 *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 *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}