nu_protocol/ast/
traverse.rs

1use crate::engine::StateWorkingSet;
2
3use super::{
4    Block, Expr, Expression, ListItem, MatchPattern, Pattern, PipelineRedirection, RecordItem,
5};
6
7/// Result of find_map closure
8#[derive(Default)]
9pub enum FindMapResult<T> {
10    Found(T),
11    #[default]
12    Continue,
13    Stop,
14}
15
16/// Trait for traversing the AST
17pub trait Traverse {
18    /// Generic function that do flat_map on an AST node.
19    /// Concatenates all recursive results on sub-expressions
20    /// into the `results` accumulator.
21    ///
22    /// # Arguments
23    /// * `f` - function that generates leaf elements
24    /// * `results` - accumulator
25    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
26    where
27        F: Fn(&'a Expression) -> Vec<T>;
28
29    /// Generic function that do find_map on an AST node.
30    /// Return the first result found by applying `f` on sub-expressions.
31    ///
32    /// # Arguments
33    /// * `f` - function that overrides the default behavior
34    fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
35    where
36        F: Fn(&'a Expression) -> FindMapResult<T>;
37}
38
39impl Traverse for Block {
40    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
41    where
42        F: Fn(&'a Expression) -> Vec<T>,
43    {
44        for pipeline in self.pipelines.iter() {
45            for element in pipeline.elements.iter() {
46                element.expr.flat_map(working_set, f, results);
47                if let Some(redir) = &element.redirection {
48                    redir.flat_map(working_set, f, results);
49                };
50            }
51        }
52    }
53
54    fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
55    where
56        F: Fn(&'a Expression) -> FindMapResult<T>,
57    {
58        self.pipelines.iter().find_map(|pipeline| {
59            pipeline.elements.iter().find_map(|element| {
60                element.expr.find_map(working_set, f).or(element
61                    .redirection
62                    .as_ref()
63                    .and_then(|redir| redir.find_map(working_set, f)))
64            })
65        })
66    }
67}
68
69impl Traverse for PipelineRedirection {
70    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
71    where
72        F: Fn(&'a Expression) -> Vec<T>,
73    {
74        let mut recur = |expr: &'a Expression| expr.flat_map(working_set, f, results);
75
76        match self {
77            PipelineRedirection::Single { target, .. } => target.expr().map(recur),
78            PipelineRedirection::Separate { out, err } => {
79                out.expr().map(&mut recur);
80                err.expr().map(&mut recur)
81            }
82        };
83    }
84
85    fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
86    where
87        F: Fn(&'a Expression) -> FindMapResult<T>,
88    {
89        let recur = |expr: &'a Expression| expr.find_map(working_set, f);
90        match self {
91            PipelineRedirection::Single { target, .. } => target.expr().and_then(recur),
92            PipelineRedirection::Separate { out, err } => {
93                [out, err].iter().filter_map(|t| t.expr()).find_map(recur)
94            }
95        }
96    }
97}
98
99impl Traverse for Expression {
100    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
101    where
102        F: Fn(&'a Expression) -> Vec<T>,
103    {
104        // leaf elements generated by `f` for this expression
105        results.extend(f(self));
106        let mut recur = |expr: &'a Expression| expr.flat_map(working_set, f, results);
107
108        match &self.expr {
109            Expr::RowCondition(block_id)
110            | Expr::Subexpression(block_id)
111            | Expr::Block(block_id)
112            | Expr::Closure(block_id) => {
113                let block = working_set.get_block(block_id.to_owned());
114                block.flat_map(working_set, f, results)
115            }
116            Expr::Range(range) => {
117                for sub_expr in [&range.from, &range.next, &range.to].into_iter().flatten() {
118                    recur(sub_expr);
119                }
120            }
121            Expr::Call(call) => {
122                for arg in &call.arguments {
123                    if let Some(sub_expr) = arg.expr() {
124                        recur(sub_expr);
125                    }
126                }
127            }
128            Expr::ExternalCall(head, args) => {
129                recur(head.as_ref());
130                for arg in args {
131                    recur(arg.expr());
132                }
133            }
134            Expr::UnaryNot(expr) | Expr::Collect(_, expr) => recur(expr.as_ref()),
135            Expr::BinaryOp(lhs, op, rhs) => {
136                recur(lhs);
137                recur(op);
138                recur(rhs);
139            }
140            Expr::MatchBlock(matches) => {
141                for (pattern, expr) in matches {
142                    pattern.flat_map(working_set, f, results);
143                    expr.flat_map(working_set, f, results);
144                }
145            }
146            Expr::List(items) => {
147                for item in items {
148                    match item {
149                        ListItem::Item(expr) | ListItem::Spread(_, expr) => recur(expr),
150                    }
151                }
152            }
153            Expr::Record(items) => {
154                for item in items {
155                    match item {
156                        RecordItem::Spread(_, expr) => recur(expr),
157                        RecordItem::Pair(key, val) => {
158                            recur(key);
159                            recur(val);
160                        }
161                    }
162                }
163            }
164            Expr::Table(table) => {
165                for column in &table.columns {
166                    recur(column);
167                }
168                for row in &table.rows {
169                    for item in row {
170                        recur(item);
171                    }
172                }
173            }
174            Expr::ValueWithUnit(vu) => recur(&vu.expr),
175            Expr::FullCellPath(fcp) => recur(&fcp.head),
176            Expr::Keyword(kw) => recur(&kw.expr),
177            Expr::StringInterpolation(vec) | Expr::GlobInterpolation(vec, _) => {
178                for item in vec {
179                    recur(item);
180                }
181            }
182            Expr::AttributeBlock(ab) => {
183                for attr in &ab.attributes {
184                    recur(&attr.expr);
185                }
186                recur(&ab.item);
187            }
188
189            _ => (),
190        };
191    }
192
193    fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
194    where
195        F: Fn(&'a Expression) -> FindMapResult<T>,
196    {
197        // behavior overridden by f
198        match f(self) {
199            FindMapResult::Found(t) => Some(t),
200            FindMapResult::Stop => None,
201            FindMapResult::Continue => {
202                let recur = |expr: &'a Expression| expr.find_map(working_set, f);
203                match &self.expr {
204                    Expr::RowCondition(block_id)
205                    | Expr::Subexpression(block_id)
206                    | Expr::Block(block_id)
207                    | Expr::Closure(block_id) => {
208                        let block = working_set.get_block(*block_id);
209                        block.find_map(working_set, f)
210                    }
211                    Expr::Range(range) => [&range.from, &range.next, &range.to]
212                        .iter()
213                        .find_map(|e| e.as_ref().and_then(recur)),
214                    Expr::Call(call) => call
215                        .arguments
216                        .iter()
217                        .find_map(|arg| arg.expr().and_then(recur)),
218                    Expr::ExternalCall(head, args) => {
219                        recur(head.as_ref()).or(args.iter().find_map(|arg| recur(arg.expr())))
220                    }
221                    Expr::UnaryNot(expr) | Expr::Collect(_, expr) => recur(expr.as_ref()),
222                    Expr::BinaryOp(lhs, op, rhs) => recur(lhs).or(recur(op)).or(recur(rhs)),
223                    Expr::MatchBlock(matches) => matches.iter().find_map(|(pattern, expr)| {
224                        pattern.find_map(working_set, f).or(recur(expr))
225                    }),
226                    Expr::List(items) => items.iter().find_map(|item| match item {
227                        ListItem::Item(expr) | ListItem::Spread(_, expr) => recur(expr),
228                    }),
229                    Expr::Record(items) => items.iter().find_map(|item| match item {
230                        RecordItem::Spread(_, expr) => recur(expr),
231                        RecordItem::Pair(key, val) => [key, val].into_iter().find_map(recur),
232                    }),
233                    Expr::Table(table) => table
234                        .columns
235                        .iter()
236                        .find_map(recur)
237                        .or(table.rows.iter().find_map(|row| row.iter().find_map(recur))),
238                    Expr::ValueWithUnit(vu) => recur(&vu.expr),
239                    Expr::FullCellPath(fcp) => recur(&fcp.head),
240                    Expr::Keyword(kw) => recur(&kw.expr),
241                    Expr::StringInterpolation(vec) | Expr::GlobInterpolation(vec, _) => {
242                        vec.iter().find_map(recur)
243                    }
244                    Expr::AttributeBlock(ab) => ab
245                        .attributes
246                        .iter()
247                        .find_map(|attr| recur(&attr.expr))
248                        .or_else(|| recur(&ab.item)),
249
250                    _ => None,
251                }
252            }
253        }
254    }
255}
256
257impl Traverse for MatchPattern {
258    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
259    where
260        F: Fn(&'a Expression) -> Vec<T>,
261    {
262        let mut recur_pattern =
263            |pattern: &'a MatchPattern| pattern.flat_map(working_set, f, results);
264
265        match &self.pattern {
266            Pattern::Expression(expr) => expr.flat_map(working_set, f, results),
267            Pattern::List(patterns) | Pattern::Or(patterns) => {
268                for pattern in patterns {
269                    recur_pattern(pattern);
270                }
271            }
272            Pattern::Record(entries) => {
273                for (_, p) in entries {
274                    recur_pattern(p);
275                }
276            }
277            _ => (),
278        };
279
280        if let Some(g) = self.guard.as_ref() {
281            g.flat_map(working_set, f, results);
282        }
283    }
284
285    fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
286    where
287        F: Fn(&'a Expression) -> FindMapResult<T>,
288    {
289        let recur = |expr: &'a Expression| expr.find_map(working_set, f);
290        let recur_pattern = |pattern: &'a MatchPattern| pattern.find_map(working_set, f);
291        match &self.pattern {
292            Pattern::Expression(expr) => recur(expr),
293            Pattern::List(patterns) | Pattern::Or(patterns) => {
294                patterns.iter().find_map(recur_pattern)
295            }
296            Pattern::Record(entries) => entries.iter().find_map(|(_, p)| recur_pattern(p)),
297            _ => None,
298        }
299        .or(self.guard.as_ref().and_then(|g| recur(g)))
300    }
301}
nu_protocol/ast/traverse.rs

nu_protocol/ast/
traverse.rs
