Improve alias expansion, again (#4474)

crates/nu-parser/src/parser.rs

@@ -842,35 +842,21 @@ pub fn parse_call(
 
                 let expansion = working_set.get_alias(alias_id);
 
-                let orig_span = spans[pos];
+                let expansion_span = span(expansion);
+
+                let orig_span = span(&[spans[cmd_start], spans[pos]]);
                 let mut new_spans: Vec<Span> = vec![];
-                new_spans.extend(&spans[0..pos]);
+                new_spans.extend(&spans[0..cmd_start]);
                 new_spans.extend(expansion);
                 if spans.len() > pos {
                     new_spans.extend(&spans[(pos + 1)..]);
                 }
 
-                let (result, err) = parse_expression(working_set, &new_spans, false);
+                let (mut result, err) = parse_expression(working_set, &new_spans, false);
 
-                let expression = match result {
-                    Expression {
-                        expr: Expr::Call(mut call),
-                        span,
-                        ty,
-                        custom_completion,
-                    } => {
-                        call.head = orig_span;
-                        Expression {
-                            expr: Expr::Call(call),
-                            span,
-                            ty,
-                            custom_completion,
-                        }
-                    }
-                    x => x,
-                };
+                result.replace_span(working_set, expansion_span, orig_span);
 
-                return (expression, err);
+                return (result, err);
             }
         }
 

crates/nu-protocol/src/ast/expression.rs

@@ -402,4 +402,126 @@ impl Expression {
             Expr::VarDecl(_) => {}
         }
     }
+
+    pub fn replace_span(
+        &mut self,
+        working_set: &mut StateWorkingSet,
+        replaced: Span,
+        new_span: Span,
+    ) {
+        if replaced.contains_span(self.span) {
+            self.span = new_span;
+        }
+        match &mut self.expr {
+            Expr::BinaryOp(left, _, right) => {
+                left.replace_span(working_set, replaced, new_span);
+                right.replace_span(working_set, replaced, new_span);
+            }
+            Expr::Block(block_id) => {
+                let mut block = working_set.get_block(*block_id).clone();
+
+                for stmt in block.stmts.iter_mut() {
+                    if let Statement::Pipeline(pipeline) = stmt {
+                        for expr in pipeline.expressions.iter_mut() {
+                            expr.replace_span(working_set, replaced, new_span)
+                        }
+                    }
+                }
+
+                *block_id = working_set.add_block(block);
+            }
+            Expr::Bool(_) => {}
+            Expr::Call(call) => {
+                if replaced.contains_span(call.head) {
+                    call.head = new_span;
+                }
+                for positional in &mut call.positional {
+                    positional.replace_span(working_set, replaced, new_span);
+                }
+                for named in &mut call.named {
+                    if let Some(expr) = &mut named.1 {
+                        expr.replace_span(working_set, replaced, new_span)
+                    }
+                }
+            }
+            Expr::CellPath(_) => {}
+            Expr::ExternalCall(head, args) => {
+                head.replace_span(working_set, replaced, new_span);
+                for arg in args {
+                    arg.replace_span(working_set, replaced, new_span)
+                }
+            }
+            Expr::Filepath(_) => {}
+            Expr::Float(_) => {}
+            Expr::FullCellPath(full_cell_path) => {
+                full_cell_path
+                    .head
+                    .replace_span(working_set, replaced, new_span);
+            }
+            Expr::ImportPattern(_) => {}
+            Expr::Garbage => {}
+            Expr::Nothing => {}
+            Expr::GlobPattern(_) => {}
+            Expr::Int(_) => {}
+            Expr::Keyword(_, _, expr) => expr.replace_span(working_set, replaced, new_span),
+            Expr::List(list) => {
+                for l in list {
+                    l.replace_span(working_set, replaced, new_span)
+                }
+            }
+            Expr::Operator(_) => {}
+            Expr::Range(left, middle, right, ..) => {
+                if let Some(left) = left {
+                    left.replace_span(working_set, replaced, new_span)
+                }
+                if let Some(middle) = middle {
+                    middle.replace_span(working_set, replaced, new_span)
+                }
+                if let Some(right) = right {
+                    right.replace_span(working_set, replaced, new_span)
+                }
+            }
+            Expr::Record(fields) => {
+                for (field_name, field_value) in fields {
+                    field_name.replace_span(working_set, replaced, new_span);
+                    field_value.replace_span(working_set, replaced, new_span);
+                }
+            }
+            Expr::Signature(_) => {}
+            Expr::String(_) => {}
+            Expr::StringInterpolation(items) => {
+                for i in items {
+                    i.replace_span(working_set, replaced, new_span)
+                }
+            }
+            Expr::RowCondition(block_id) | Expr::Subexpression(block_id) => {
+                let mut block = working_set.get_block(*block_id).clone();
+
+                for stmt in block.stmts.iter_mut() {
+                    if let Statement::Pipeline(pipeline) = stmt {
+                        for expr in pipeline.expressions.iter_mut() {
+                            expr.replace_span(working_set, replaced, new_span)
+                        }
+                    }
+                }
+
+                *block_id = working_set.add_block(block);
+            }
+            Expr::Table(headers, cells) => {
+                for header in headers {
+                    header.replace_span(working_set, replaced, new_span)
+                }
+
+                for row in cells {
+                    for cell in row.iter_mut() {
+                        cell.replace_span(working_set, replaced, new_span)
+                    }
+                }
+            }
+
+            Expr::ValueWithUnit(expr, _) => expr.replace_span(working_set, replaced, new_span),
+            Expr::Var(_) => {}
+            Expr::VarDecl(_) => {}
+        }
+    }
 }

crates/nu-protocol/src/span.rs

@@ -48,6 +48,10 @@ impl Span {
         pos >= self.start && pos < self.end
     }
 
+    pub fn contains_span(&self, span: Span) -> bool {
+        span.start >= self.start && span.end <= self.end
+    }
+
     /// Point to the space just past this span, useful for missing
     /// values
     pub fn past(&self) -> Span {