nu_protocol/ast/
traverse.rs

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

nu_protocol/ast/
traverse.rs
