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 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
102fn 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 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 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 output.push((call.head, FlatShape::InternalCall(call.decl_id)));
274 }
275
276 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 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 output[arg_start..].sort();
324 }
325 Expr::ExternalCall(head, args) => {
326 if let Expr::String(..) | Expr::GlobPattern(..) = &head.expr {
327 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 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 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}