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                        // Clone the block_id to create an owned value
209                        let block_id = block_id.to_owned();
210                        let block = working_set.get_block(block_id);
211                        block.find_map(working_set, f)
212                    }
213                    Expr::Range(range) => [&range.from, &range.next, &range.to]
214                        .iter()
215                        .find_map(|e| e.as_ref().and_then(recur)),
216                    Expr::Call(call) => call
217                        .arguments
218                        .iter()
219                        .find_map(|arg| arg.expr().and_then(recur)),
220                    Expr::ExternalCall(head, args) => {
221                        recur(head.as_ref()).or(args.iter().find_map(|arg| recur(arg.expr())))
222                    }
223                    Expr::UnaryNot(expr) | Expr::Collect(_, expr) => recur(expr.as_ref()),
224                    Expr::BinaryOp(lhs, op, rhs) => recur(lhs).or(recur(op)).or(recur(rhs)),
225                    Expr::MatchBlock(matches) => matches.iter().find_map(|(pattern, expr)| {
226                        pattern.find_map(working_set, f).or(recur(expr))
227                    }),
228                    Expr::List(items) => items.iter().find_map(|item| match item {
229                        ListItem::Item(expr) | ListItem::Spread(_, expr) => recur(expr),
230                    }),
231                    Expr::Record(items) => items.iter().find_map(|item| match item {
232                        RecordItem::Spread(_, expr) => recur(expr),
233                        RecordItem::Pair(key, val) => [key, val].into_iter().find_map(recur),
234                    }),
235                    Expr::Table(table) => table
236                        .columns
237                        .iter()
238                        .find_map(recur)
239                        .or(table.rows.iter().find_map(|row| row.iter().find_map(recur))),
240                    Expr::ValueWithUnit(vu) => recur(&vu.expr),
241                    Expr::FullCellPath(fcp) => recur(&fcp.head),
242                    Expr::Keyword(kw) => recur(&kw.expr),
243                    Expr::StringInterpolation(vec) | Expr::GlobInterpolation(vec, _) => {
244                        vec.iter().find_map(recur)
245                    }
246                    Expr::AttributeBlock(ab) => ab
247                        .attributes
248                        .iter()
249                        .find_map(|attr| recur(&attr.expr))
250                        .or_else(|| recur(&ab.item)),
251
252                    _ => None,
253                }
254            }
255        }
256    }
257}
258
259impl Traverse for MatchPattern {
260    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
261    where
262        F: Fn(&'a Expression) -> Vec<T>,
263    {
264        let mut recur_pattern =
265            |pattern: &'a MatchPattern| pattern.flat_map(working_set, f, results);
266
267        match &self.pattern {
268            Pattern::Expression(expr) => expr.flat_map(working_set, f, results),
269            Pattern::List(patterns) | Pattern::Or(patterns) => {
270                for pattern in patterns {
271                    recur_pattern(pattern);
272                }
273            }
274            Pattern::Record(entries) => {
275                for (_, p) in entries {
276                    recur_pattern(p);
277                }
278            }
279            _ => (),
280        };
281
282        if let Some(g) = self.guard.as_ref() {
283            g.flat_map(working_set, f, results);
284        }
285    }
286
287    fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
288    where
289        F: Fn(&'a Expression) -> FindMapResult<T>,
290    {
291        let recur = |expr: &'a Expression| expr.find_map(working_set, f);
292        let recur_pattern = |pattern: &'a MatchPattern| pattern.find_map(working_set, f);
293        match &self.pattern {
294            Pattern::Expression(expr) => recur(expr),
295            Pattern::List(patterns) | Pattern::Or(patterns) => {
296                patterns.iter().find_map(recur_pattern)
297            }
298            Pattern::Record(entries) => entries.iter().find_map(|(_, p)| recur_pattern(p)),
299            _ => None,
300        }
301        .or(self.guard.as_ref().and_then(|g| recur(g)))
302    }
303}
nu_protocol/ast/traverse.rs

nu_protocol/ast/
traverse.rs
