nu_parser/
flatten.rs

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

nu_parser/
flatten.rs
