From 739425431a502df303023942f8edd84f03b3f4cc Mon Sep 17 00:00:00 2001
From: JT <jonathan.d.turner@gmail.com>
Date: Tue, 17 Aug 2021 12:26:05 +1200
Subject: [PATCH] improve type inference

---
 crates/nu-engine/src/eval.rs   |  4 +---
 crates/nu-parser/src/parser.rs | 26 +++++++++++++++++++++++++-
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/crates/nu-engine/src/eval.rs b/crates/nu-engine/src/eval.rs
index 5af947525..089cc68af 100644
--- a/crates/nu-engine/src/eval.rs
+++ b/crates/nu-engine/src/eval.rs
@@ -182,9 +182,7 @@ fn eval_call(state: &State, call: &Call) -> Result<Value, ShellError> {
         })
     } else if decl.signature.name == "vars" {
         state.parser_state.borrow().print_vars();
-        Ok(Value::Nothing {
-            span: call.positional[0].span,
-        })
+        Ok(Value::Nothing { span: call.head })
     } else if decl.signature.name == "decls" {
         state.parser_state.borrow().print_decls();
         Ok(Value::Nothing { span: call.head })
diff --git a/crates/nu-parser/src/parser.rs b/crates/nu-parser/src/parser.rs
index 18ca7b6e7..deab3f13b 100644
--- a/crates/nu-parser/src/parser.rs
+++ b/crates/nu-parser/src/parser.rs
@@ -1784,6 +1784,8 @@ impl<'a> ParserWorkingSet<'a> {
 
         let mut args = vec![];
 
+        let mut contained_type: Option<Type> = None;
+
         if !output.block.is_empty() {
             for arg in &output.block[0].commands {
                 let mut spans_idx = 0;
@@ -1793,6 +1795,14 @@ impl<'a> ParserWorkingSet<'a> {
                         self.parse_multispan_value(&arg.parts, &mut spans_idx, element_shape);
                     error = error.or(err);
 
+                    if let Some(ref ctype) = contained_type {
+                        if *ctype != arg.ty {
+                            contained_type = Some(Type::Unknown);
+                        }
+                    } else {
+                        contained_type = Some(arg.ty.clone());
+                    }
+
                     args.push(arg);
 
                     spans_idx += 1;
@@ -1804,7 +1814,11 @@ impl<'a> ParserWorkingSet<'a> {
             Expression {
                 expr: Expr::List(args),
                 span,
-                ty: Type::List(Box::new(Type::Unknown)), // FIXME
+                ty: Type::List(Box::new(if let Some(ty) = contained_type {
+                    ty.clone()
+                } else {
+                    Type::Unknown
+                })),
             },
             error,
         )
@@ -2416,6 +2430,16 @@ impl<'a> ParserWorkingSet<'a> {
                 let (call, call_span, err) =
                     self.parse_internal_call(spans[0], &spans[1..], decl_id);
 
+                // Update the variable to the known type if we can.
+                if err.is_none() {
+                    let var_id = call.positional[0]
+                        .as_var()
+                        .expect("internal error: expected variable");
+                    let rhs_type = call.positional[1].ty.clone();
+
+                    self.set_variable_type(var_id, rhs_type);
+                }
+
                 return (
                     Statement::Expression(Expression {
                         expr: Expr::Call(call),