nu_parser/
flatten.rs

1use nu_protocol::{
2    DeclId, GetSpan, Span, SyntaxShape, VarId,
3    ast::{
4        Argument, Block, Expr, Expression, ExternalArgument, ImportPatternMember, ListItem,
5        MatchPattern, PathMember, Pattern, Pipeline, PipelineElement, PipelineRedirection,
6        RecordItem,
7    },
8    engine::StateWorkingSet,
9};
10use std::fmt::{Display, Formatter, Result};
11
12#[derive(Debug, Eq, PartialEq, Ord, Clone, PartialOrd)]
13pub enum FlatShape {
14    Binary,
15    Block,
16    Bool,
17    Closure,
18    Custom(DeclId),
19    DateTime,
20    Directory,
21    // The stored span contains the call's head if this call is through an alias:
22    // This is only different from the name of the called external command,
23    // and is only useful for its location (not its contents).
24    External(Box<Span>),
25    ExternalArg,
26    ExternalResolved,
27    Filepath,
28    Flag,
29    Float,
30    Garbage,
31    GlobInterpolation,
32    GlobPattern,
33    Int,
34    InternalCall(DeclId),
35    Keyword,
36    List,
37    Literal,
38    MatchPattern,
39    Nothing,
40    Operator,
41    Pipe,
42    Range,
43    RawString,
44    Record,
45    Redirection,
46    Signature,
47    String,
48    StringInterpolation,
49    Table,
50    Variable(VarId),
51    VarDecl(VarId),
52}
53
54impl FlatShape {
55    pub fn as_str(&self) -> &str {
56        match self {
57            FlatShape::Binary => "shape_binary",
58            FlatShape::Block => "shape_block",
59            FlatShape::Bool => "shape_bool",
60            FlatShape::Closure => "shape_closure",
61            FlatShape::Custom(_) => "shape_custom",
62            FlatShape::DateTime => "shape_datetime",
63            FlatShape::Directory => "shape_directory",
64            FlatShape::External(_) => "shape_external",
65            FlatShape::ExternalArg => "shape_externalarg",
66            FlatShape::ExternalResolved => "shape_external_resolved",
67            FlatShape::Filepath => "shape_filepath",
68            FlatShape::Flag => "shape_flag",
69            FlatShape::Float => "shape_float",
70            FlatShape::Garbage => "shape_garbage",
71            FlatShape::GlobInterpolation => "shape_glob_interpolation",
72            FlatShape::GlobPattern => "shape_globpattern",
73            FlatShape::Int => "shape_int",
74            FlatShape::InternalCall(_) => "shape_internalcall",
75            FlatShape::Keyword => "shape_keyword",
76            FlatShape::List => "shape_list",
77            FlatShape::Literal => "shape_literal",
78            FlatShape::MatchPattern => "shape_match_pattern",
79            FlatShape::Nothing => "shape_nothing",
80            FlatShape::Operator => "shape_operator",
81            FlatShape::Pipe => "shape_pipe",
82            FlatShape::Range => "shape_range",
83            FlatShape::RawString => "shape_raw_string",
84            FlatShape::Record => "shape_record",
85            FlatShape::Redirection => "shape_redirection",
86            FlatShape::Signature => "shape_signature",
87            FlatShape::String => "shape_string",
88            FlatShape::StringInterpolation => "shape_string_interpolation",
89            FlatShape::Table => "shape_table",
90            FlatShape::Variable(_) => "shape_variable",
91            FlatShape::VarDecl(_) => "shape_vardecl",
92        }
93    }
94}
95
96impl Display for FlatShape {
97    fn fmt(&self, f: &mut Formatter) -> Result {
98        f.write_str(self.as_str())
99    }
100}
101
102/*
103The `_into` functions below (e.g., `flatten_block_into`) take an existing `output` `Vec`
104and append more data to it. This is to reduce the number of intermediate `Vec`s.
105The non-`into` functions (e.g., `flatten_block`) are part of the crate's public API
106and return a new `Vec` instead of modifying an existing one.
107*/
108
109fn flatten_block_into(
110    working_set: &StateWorkingSet,
111    block: &Block,
112    output: &mut Vec<(Span, FlatShape)>,
113) {
114    for pipeline in &block.pipelines {
115        flatten_pipeline_into(working_set, pipeline, output);
116    }
117}
118
119fn flatten_pipeline_into(
120    working_set: &StateWorkingSet,
121    pipeline: &Pipeline,
122    output: &mut Vec<(Span, FlatShape)>,
123) {
124    for expr in &pipeline.elements {
125        flatten_pipeline_element_into(working_set, expr, output)
126    }
127}
128
129fn flatten_pipeline_element_into(
130    working_set: &StateWorkingSet,
131    pipeline_element: &PipelineElement,
132    output: &mut Vec<(Span, FlatShape)>,
133) {
134    if let Some(span) = pipeline_element.pipe {
135        output.push((span, FlatShape::Pipe));
136    }
137
138    flatten_expression_into(working_set, &pipeline_element.expr, output);
139
140    if let Some(redirection) = pipeline_element.redirection.as_ref() {
141        match redirection {
142            PipelineRedirection::Single { target, .. } => {
143                output.push((target.span(), FlatShape::Redirection));
144                if let Some(expr) = target.expr() {
145                    flatten_expression_into(working_set, expr, output);
146                }
147            }
148            PipelineRedirection::Separate { out, err } => {
149                let (out, err) = if out.span() <= err.span() {
150                    (out, err)
151                } else {
152                    (err, out)
153                };
154
155                output.push((out.span(), FlatShape::Redirection));
156                if let Some(expr) = out.expr() {
157                    flatten_expression_into(working_set, expr, output);
158                }
159                output.push((err.span(), FlatShape::Redirection));
160                if let Some(expr) = err.expr() {
161                    flatten_expression_into(working_set, expr, output);
162                }
163            }
164        }
165    }
166}
167
168fn flatten_positional_arg_into(
169    working_set: &StateWorkingSet,
170    positional: &Expression,
171    shape: &SyntaxShape,
172    output: &mut Vec<(Span, FlatShape)>,
173) {
174    if matches!(shape, SyntaxShape::ExternalArgument)
175        && matches!(positional.expr, Expr::String(..) | Expr::GlobPattern(..))
176    {
177        // Make known external arguments look more like external arguments
178        output.push((positional.span, FlatShape::ExternalArg));
179    } else {
180        flatten_expression_into(working_set, positional, output)
181    }
182}
183
184fn flatten_expression_into(
185    working_set: &StateWorkingSet,
186    expr: &Expression,
187    output: &mut Vec<(Span, FlatShape)>,
188) {
189    match &expr.expr {
190        Expr::AttributeBlock(ab) => {
191            for attr in &ab.attributes {
192                flatten_expression_into(working_set, &attr.expr, output);
193            }
194            flatten_expression_into(working_set, &ab.item, output);
195        }
196        Expr::BinaryOp(lhs, op, rhs) => {
197            flatten_expression_into(working_set, lhs, output);
198            flatten_expression_into(working_set, op, output);
199            flatten_expression_into(working_set, rhs, output);
200        }
201        Expr::UnaryNot(not) => {
202            output.push((
203                Span::new(expr.span.start, expr.span.start + 3),
204                FlatShape::Operator,
205            ));
206            flatten_expression_into(working_set, not, output);
207        }
208        Expr::Collect(_, expr) => {
209            flatten_expression_into(working_set, expr, output);
210        }
211        Expr::Closure(block_id) => {
212            let outer_span = expr.span;
213
214            let block = working_set.get_block(*block_id);
215            let flattened = flatten_block(working_set, block);
216
217            if let Some(first) = flattened.first()
218                && first.0.start > outer_span.start
219            {
220                output.push((
221                    Span::new(outer_span.start, first.0.start),
222                    FlatShape::Closure,
223                ));
224            }
225
226            let last = if let Some(last) = flattened.last() {
227                if last.0.end < outer_span.end {
228                    Some((Span::new(last.0.end, outer_span.end), FlatShape::Closure))
229                } else {
230                    None
231                }
232            } else {
233                // for empty closures
234                Some((outer_span, FlatShape::Closure))
235            };
236
237            output.extend(flattened);
238            if let Some(last) = last {
239                output.push(last);
240            }
241        }
242        Expr::Block(block_id) | Expr::RowCondition(block_id) | Expr::Subexpression(block_id) => {
243            let outer_span = expr.span;
244
245            let flattened = flatten_block(working_set, working_set.get_block(*block_id));
246
247            if let Some(first) = flattened.first()
248                && first.0.start > outer_span.start
249            {
250                output.push((Span::new(outer_span.start, first.0.start), FlatShape::Block));
251            }
252
253            let last = if let Some(last) = flattened.last() {
254                if last.0.end < outer_span.end {
255                    Some((Span::new(last.0.end, outer_span.end), FlatShape::Block))
256                } else {
257                    None
258                }
259            } else {
260                None
261            };
262
263            output.extend(flattened);
264            if let Some(last) = last {
265                output.push(last);
266            }
267        }
268        Expr::Call(call) => {
269            let decl = working_set.get_decl(call.decl_id);
270
271            if call.head.end != 0 {
272                // Make sure we don't push synthetic calls
273                output.push((call.head, FlatShape::InternalCall(call.decl_id)));
274            }
275
276            // Follow positional arguments from the signature.
277            let signature = decl.signature();
278            let mut positional_args = signature
279                .required_positional
280                .iter()
281                .chain(&signature.optional_positional);
282
283            let arg_start = output.len();
284            for arg in &call.arguments {
285                match arg {
286                    Argument::Positional(positional) => {
287                        let positional_arg = positional_args.next();
288                        let shape = positional_arg
289                            .or(signature.rest_positional.as_ref())
290                            .map(|arg| &arg.shape)
291                            .unwrap_or(&SyntaxShape::Any);
292
293                        flatten_positional_arg_into(working_set, positional, shape, output)
294                    }
295                    Argument::Unknown(positional) => {
296                        let shape = signature
297                            .rest_positional
298                            .as_ref()
299                            .map(|arg| &arg.shape)
300                            .unwrap_or(&SyntaxShape::Any);
301
302                        flatten_positional_arg_into(working_set, positional, shape, output)
303                    }
304                    Argument::Named(named) => {
305                        if named.0.span.end != 0 {
306                            // Ignore synthetic flags
307                            output.push((named.0.span, FlatShape::Flag));
308                        }
309                        if let Some(expr) = &named.2 {
310                            flatten_expression_into(working_set, expr, output);
311                        }
312                    }
313                    Argument::Spread(expr) => {
314                        output.push((
315                            Span::new(expr.span.start - 3, expr.span.start),
316                            FlatShape::Operator,
317                        ));
318                        flatten_expression_into(working_set, expr, output);
319                    }
320                }
321            }
322            // sort these since flags and positional args can be intermixed
323            output[arg_start..].sort();
324        }
325        Expr::ExternalCall(head, args) => {
326            if let Expr::String(..) | Expr::GlobPattern(..) = &head.expr {
327                // If this external call is through an alias, then head.span contains the
328                // name of the alias (needed to highlight the right thing), but we also need
329                // the name of the aliased command (to decide *how* to highlight the call).
330                // The parser actually created this head by cloning from the alias's definition
331                // and then just overwriting the `span` field - but `span_id` still points to
332                // the original span, so we can recover it from there.
333                let span = working_set.get_span(head.span_id);
334                output.push((span, FlatShape::External(Box::new(head.span))));
335            } else {
336                flatten_expression_into(working_set, head, output);
337            }
338
339            for arg in args.as_ref() {
340                match arg {
341                    ExternalArgument::Regular(expr) => {
342                        if let Expr::String(..) | Expr::GlobPattern(..) = &expr.expr {
343                            output.push((expr.span, FlatShape::ExternalArg));
344                        } else {
345                            flatten_expression_into(working_set, expr, output);
346                        }
347                    }
348                    ExternalArgument::Spread(expr) => {
349                        output.push((
350                            Span::new(expr.span.start - 3, expr.span.start),
351                            FlatShape::Operator,
352                        ));
353                        flatten_expression_into(working_set, expr, output);
354                    }
355                }
356            }
357        }
358        Expr::Garbage => output.push((expr.span, FlatShape::Garbage)),
359        Expr::Nothing => output.push((expr.span, FlatShape::Nothing)),
360        Expr::DateTime(_) => output.push((expr.span, FlatShape::DateTime)),
361        Expr::Binary(_) => output.push((expr.span, FlatShape::Binary)),
362        Expr::Int(_) => output.push((expr.span, FlatShape::Int)),
363        Expr::Float(_) => output.push((expr.span, FlatShape::Float)),
364        Expr::MatchBlock(matches) => {
365            for (pattern, expr) in matches {
366                flatten_pattern_into(pattern, output);
367                flatten_expression_into(working_set, expr, output);
368            }
369        }
370        Expr::ValueWithUnit(value) => {
371            flatten_expression_into(working_set, &value.expr, output);
372            output.push((value.unit.span, FlatShape::String));
373        }
374        Expr::CellPath(cell_path) => {
375            output.extend(cell_path.members.iter().map(|member| match *member {
376                PathMember::String { span, .. } => (span, FlatShape::String),
377                PathMember::Int { span, .. } => (span, FlatShape::Int),
378            }));
379        }
380        Expr::FullCellPath(cell_path) => {
381            flatten_expression_into(working_set, &cell_path.head, output);
382            output.extend(cell_path.tail.iter().map(|member| match *member {
383                PathMember::String { span, .. } => (span, FlatShape::String),
384                PathMember::Int { span, .. } => (span, FlatShape::Int),
385            }));
386        }
387        Expr::ImportPattern(import_pattern) => {
388            output.push((import_pattern.head.span, FlatShape::String));
389
390            for member in &import_pattern.members {
391                match member {
392                    ImportPatternMember::Glob { span } => output.push((*span, FlatShape::String)),
393                    ImportPatternMember::Name { span, .. } => {
394                        output.push((*span, FlatShape::String))
395                    }
396                    ImportPatternMember::List { names } => {
397                        output.extend(names.iter().map(|&(_, span)| (span, FlatShape::String)))
398                    }
399                }
400            }
401        }
402        Expr::Overlay(_) => output.push((expr.span, FlatShape::String)),
403        Expr::Range(range) => {
404            if let Some(f) = &range.from {
405                flatten_expression_into(working_set, f, output);
406            }
407            if let Some(s) = &range.next {
408                output.push((range.operator.next_op_span, FlatShape::Operator));
409                flatten_expression_into(working_set, s, output);
410            }
411            output.push((range.operator.span, FlatShape::Operator));
412            if let Some(t) = &range.to {
413                flatten_expression_into(working_set, t, output);
414            }
415        }
416        Expr::Bool(_) => output.push((expr.span, FlatShape::Bool)),
417        Expr::Filepath(_, _) => output.push((expr.span, FlatShape::Filepath)),
418        Expr::Directory(_, _) => output.push((expr.span, FlatShape::Directory)),
419        Expr::GlobPattern(_, _) => output.push((expr.span, FlatShape::GlobPattern)),
420        Expr::List(list) => {
421            let outer_span = expr.span;
422            let mut last_end = outer_span.start;
423
424            for item in list {
425                match item {
426                    ListItem::Item(expr) => {
427                        let flattened = flatten_expression(working_set, expr);
428
429                        if let Some(first) = flattened.first()
430                            && first.0.start > last_end
431                        {
432                            output.push((Span::new(last_end, first.0.start), FlatShape::List));
433                        }
434
435                        if let Some(last) = flattened.last() {
436                            last_end = last.0.end;
437                        }
438
439                        output.extend(flattened);
440                    }
441                    ListItem::Spread(op_span, expr) => {
442                        if op_span.start > last_end {
443                            output.push((Span::new(last_end, op_span.start), FlatShape::List));
444                        }
445                        output.push((*op_span, FlatShape::Operator));
446                        last_end = op_span.end;
447
448                        let flattened_inner = flatten_expression(working_set, expr);
449                        if let Some(first) = flattened_inner.first()
450                            && first.0.start > last_end
451                        {
452                            output.push((Span::new(last_end, first.0.start), FlatShape::List));
453                        }
454                        if let Some(last) = flattened_inner.last() {
455                            last_end = last.0.end;
456                        }
457                        output.extend(flattened_inner);
458                    }
459                }
460            }
461
462            if last_end < outer_span.end {
463                output.push((Span::new(last_end, outer_span.end), FlatShape::List));
464            }
465        }
466        Expr::StringInterpolation(exprs) => {
467            let mut flattened = vec![];
468            for expr in exprs {
469                flatten_expression_into(working_set, expr, &mut flattened);
470            }
471
472            if let Some(first) = flattened.first()
473                && first.0.start != expr.span.start
474            {
475                // If we aren't a bare word interpolation, also highlight the outer quotes
476                output.push((
477                    Span::new(expr.span.start, expr.span.start + 2),
478                    FlatShape::StringInterpolation,
479                ));
480                flattened.push((
481                    Span::new(expr.span.end - 1, expr.span.end),
482                    FlatShape::StringInterpolation,
483                ));
484            }
485            output.extend(flattened);
486        }
487        Expr::GlobInterpolation(exprs, quoted) => {
488            let mut flattened = vec![];
489            for expr in exprs {
490                flatten_expression_into(working_set, expr, &mut flattened);
491            }
492
493            if *quoted {
494                // If we aren't a bare word interpolation, also highlight the outer quotes
495                output.push((
496                    Span::new(expr.span.start, expr.span.start + 2),
497                    FlatShape::GlobInterpolation,
498                ));
499                flattened.push((
500                    Span::new(expr.span.end - 1, expr.span.end),
501                    FlatShape::GlobInterpolation,
502                ));
503            }
504            output.extend(flattened);
505        }
506        Expr::Record(list) => {
507            let outer_span = expr.span;
508            let mut last_end = outer_span.start;
509
510            for l in list {
511                match l {
512                    RecordItem::Pair(key, val) => {
513                        let flattened_lhs = flatten_expression(working_set, key);
514                        let flattened_rhs = flatten_expression(working_set, val);
515
516                        if let Some(first) = flattened_lhs.first()
517                            && first.0.start > last_end
518                        {
519                            output.push((Span::new(last_end, first.0.start), FlatShape::Record));
520                        }
521                        if let Some(last) = flattened_lhs.last() {
522                            last_end = last.0.end;
523                        }
524                        output.extend(flattened_lhs);
525
526                        if let Some(first) = flattened_rhs.first()
527                            && first.0.start > last_end
528                        {
529                            output.push((Span::new(last_end, first.0.start), FlatShape::Record));
530                        }
531                        if let Some(last) = flattened_rhs.last() {
532                            last_end = last.0.end;
533                        }
534
535                        output.extend(flattened_rhs);
536                    }
537                    RecordItem::Spread(op_span, record) => {
538                        if op_span.start > last_end {
539                            output.push((Span::new(last_end, op_span.start), FlatShape::Record));
540                        }
541                        output.push((*op_span, FlatShape::Operator));
542                        last_end = op_span.end;
543
544                        let flattened = flatten_expression(working_set, record);
545                        if let Some(first) = flattened.first()
546                            && first.0.start > last_end
547                        {
548                            output.push((Span::new(last_end, first.0.start), FlatShape::Record));
549                        }
550                        if let Some(last) = flattened.last() {
551                            last_end = last.0.end;
552                        }
553                        output.extend(flattened);
554                    }
555                }
556            }
557            if last_end < outer_span.end {
558                output.push((Span::new(last_end, outer_span.end), FlatShape::Record));
559            }
560        }
561        Expr::Keyword(kw) => {
562            output.push((kw.span, FlatShape::Keyword));
563            flatten_expression_into(working_set, &kw.expr, output);
564        }
565        Expr::Operator(_) => output.push((expr.span, FlatShape::Operator)),
566        Expr::Signature(_) => output.push((expr.span, FlatShape::Signature)),
567        Expr::String(_) => output.push((expr.span, FlatShape::String)),
568        Expr::RawString(_) => output.push((expr.span, FlatShape::RawString)),
569        Expr::Table(table) => {
570            let outer_span = expr.span;
571            let mut last_end = outer_span.start;
572
573            for col in table.columns.as_ref() {
574                let flattened = flatten_expression(working_set, col);
575                if let Some(first) = flattened.first()
576                    && first.0.start > last_end
577                {
578                    output.push((Span::new(last_end, first.0.start), FlatShape::Table));
579                }
580
581                if let Some(last) = flattened.last() {
582                    last_end = last.0.end;
583                }
584
585                output.extend(flattened);
586            }
587            for row in table.rows.as_ref() {
588                for expr in row.as_ref() {
589                    let flattened = flatten_expression(working_set, expr);
590                    if let Some(first) = flattened.first()
591                        && first.0.start > last_end
592                    {
593                        output.push((Span::new(last_end, first.0.start), FlatShape::Table));
594                    }
595
596                    if let Some(last) = flattened.last() {
597                        last_end = last.0.end;
598                    }
599
600                    output.extend(flattened);
601                }
602            }
603
604            if last_end < outer_span.end {
605                output.push((Span::new(last_end, outer_span.end), FlatShape::Table));
606            }
607        }
608        Expr::Var(var_id) => output.push((expr.span, FlatShape::Variable(*var_id))),
609        Expr::VarDecl(var_id) => output.push((expr.span, FlatShape::VarDecl(*var_id))),
610    }
611}
612
613fn flatten_pattern_into(match_pattern: &MatchPattern, output: &mut Vec<(Span, FlatShape)>) {
614    match &match_pattern.pattern {
615        Pattern::Garbage => output.push((match_pattern.span, FlatShape::Garbage)),
616        Pattern::IgnoreValue => output.push((match_pattern.span, FlatShape::Nothing)),
617        Pattern::IgnoreRest => output.push((match_pattern.span, FlatShape::Nothing)),
618        Pattern::List(items) => {
619            if let Some(first) = items.first() {
620                if let Some(last) = items.last() {
621                    output.push((
622                        Span::new(match_pattern.span.start, first.span.start),
623                        FlatShape::MatchPattern,
624                    ));
625                    for item in items {
626                        flatten_pattern_into(item, output);
627                    }
628                    output.push((
629                        Span::new(last.span.end, match_pattern.span.end),
630                        FlatShape::MatchPattern,
631                    ))
632                }
633            } else {
634                output.push((match_pattern.span, FlatShape::MatchPattern));
635            }
636        }
637        Pattern::Record(items) => {
638            if let Some(first) = items.first() {
639                if let Some(last) = items.last() {
640                    output.push((
641                        Span::new(match_pattern.span.start, first.1.span.start),
642                        FlatShape::MatchPattern,
643                    ));
644                    for (_, pattern) in items {
645                        flatten_pattern_into(pattern, output);
646                    }
647                    output.push((
648                        Span::new(last.1.span.end, match_pattern.span.end),
649                        FlatShape::MatchPattern,
650                    ))
651                }
652            } else {
653                output.push((match_pattern.span, FlatShape::MatchPattern));
654            }
655        }
656        Pattern::Expression(_) | Pattern::Value(_) => {
657            output.push((match_pattern.span, FlatShape::MatchPattern))
658        }
659        Pattern::Variable(var_id) => output.push((match_pattern.span, FlatShape::VarDecl(*var_id))),
660        Pattern::Rest(var_id) => output.push((match_pattern.span, FlatShape::VarDecl(*var_id))),
661        Pattern::Or(patterns) => {
662            for pattern in patterns {
663                flatten_pattern_into(pattern, output);
664            }
665        }
666    }
667}
668
669pub fn flatten_block(working_set: &StateWorkingSet, block: &Block) -> Vec<(Span, FlatShape)> {
670    let mut output = Vec::new();
671    flatten_block_into(working_set, block, &mut output);
672    output
673}
674
675pub fn flatten_pipeline(
676    working_set: &StateWorkingSet,
677    pipeline: &Pipeline,
678) -> Vec<(Span, FlatShape)> {
679    let mut output = Vec::new();
680    flatten_pipeline_into(working_set, pipeline, &mut output);
681    output
682}
683
684pub fn flatten_pipeline_element(
685    working_set: &StateWorkingSet,
686    pipeline_element: &PipelineElement,
687) -> Vec<(Span, FlatShape)> {
688    let mut output = Vec::new();
689    flatten_pipeline_element_into(working_set, pipeline_element, &mut output);
690    output
691}
692
693pub fn flatten_expression(
694    working_set: &StateWorkingSet,
695    expr: &Expression,
696) -> Vec<(Span, FlatShape)> {
697    let mut output = Vec::new();
698    flatten_expression_into(working_set, expr, &mut output);
699    output
700}