feat!(parser): type inferrence following cellpath

Author: blindFS

crates/nu-cli/src/completions/cell_path_completions.rs

@@ -72,8 +72,8 @@ impl Completer for CellPathCompletion<'_> {
             for suggestion in get_suggestions_by_value(&value, current_span) {
                 matcher.add_semantic_suggestion(suggestion);
             }
-        } else if let Some(ty) = type_follow_cell_path(&self.full_cell_path.head.ty, path_members) {
-            for suggestion in get_suggestions_by_type(ty, current_span) {
+        } else if let Some(ty) = self.full_cell_path.head.ty.follow_cell_path(path_members) {
+            for suggestion in get_suggestions_by_type(&ty, current_span) {
                 matcher.add_semantic_suggestion(suggestion);
             }
         }
@@ -155,32 +155,6 @@ fn get_suggestions_by_value(
     }
 }
 
-fn type_follow_cell_path<'a>(
-    parent_type: &'a Type,
-    path_members: &[PathMember],
-) -> Option<&'a Type> {
-    match (parent_type, path_members.first()) {
-        (Type::Table(fields), Some(PathMember::String { val, .. }))
-        | (Type::Record(fields), Some(PathMember::String { val, .. })) => {
-            let sub_type = fields
-                .iter()
-                .find_map(|(name, ty)| (name == val).then_some(ty))?;
-            type_follow_cell_path(sub_type, &path_members[1..])
-        }
-        (Type::Table(_), Some(PathMember::Int { .. })) => {
-            type_follow_cell_path(parent_type, &path_members[1..])
-        }
-        (Type::List(inner), Some(PathMember::Int { .. })) => {
-            type_follow_cell_path(inner, &path_members[1..])
-        }
-        (Type::List(inner), Some(PathMember::String { .. })) => {
-            type_follow_cell_path(inner, path_members)
-        }
-        (_, None) => Some(parent_type),
-        _ => None,
-    }
-}
-
 fn get_suggestions_by_type(ty: &Type, current_span: reedline::Span) -> Vec<SemanticSuggestion> {
     match ty {
         Type::Record(fields) | Type::Table(fields) => fields

crates/nu-cli/src/completions/operator_completions.rs

@@ -228,7 +228,7 @@ impl Completer for OperatorCompletion<'_> {
             Type::Record(_) | Type::Range => collection_comparison_ops(),
             Type::List(_) | Type::Table(_) => valid_list_ops(),
             // Unknown type, resort to evaluated values
-            Type::Any => match &self.left_hand_side.expr {
+            Type::Any | Type::OneOf(_) => match &self.left_hand_side.expr {
                 Expr::FullCellPath(path) => {
                     // for `$ <tab>`
                     if let Expr::Garbage = path.head.expr {

crates/nu-parser/src/parser.rs

@@ -2774,11 +2774,11 @@ pub fn parse_full_cell_path(
         };
 
         let tail = parse_cell_path(working_set, tokens, expect_dot);
-        // FIXME: Get the type of the data at the tail using follow_cell_path() (or something)
         let ty = if !tail.is_empty() {
-            // Until the aforementioned fix is implemented, this is necessary to allow mutable list upserts
-            // such as $a.1 = 2 to work correctly.
-            Type::Any
+            head.ty
+                .follow_cell_path(&tail)
+                .map(|ty| ty.into_owned())
+                .unwrap_or(Type::Any)
         } else {
             head.ty.clone()
         };

crates/nu-protocol/src/ty.rs

@@ -1,6 +1,6 @@
-use crate::SyntaxShape;
+use crate::{SyntaxShape, ast::PathMember};
 use serde::{Deserialize, Serialize};
-use std::fmt::Display;
+use std::{borrow::Cow, fmt::Display};
 #[cfg(test)]
 use strum_macros::EnumIter;
 
@@ -35,6 +35,76 @@ pub enum Type {
     Table(Box<[(String, Type)]>),
 }
 
+fn follow_cell_path_recursive<'a>(
+    current: Cow<'a, Type>,
+    path_members: &mut dyn Iterator<Item = &'a PathMember>,
+) -> Option<Cow<'a, Type>> {
+    let Some(first) = path_members.next() else {
+        return Some(current);
+    };
+    match (current.as_ref(), first) {
+        (Type::Record(fields), PathMember::String { val, .. }) => {
+            let idx = fields.iter().position(|(name, _)| name == val)?;
+            let next = match current {
+                Cow::Borrowed(Type::Record(f)) => Cow::Borrowed(&f[idx].1),
+                Cow::Owned(Type::Record(f)) => Cow::Owned(f[idx].1.to_owned()),
+                _ => unreachable!(),
+            };
+            follow_cell_path_recursive(next, path_members)
+        }
+
+        // Table to Record (Int)
+        (Type::Table(f), PathMember::Int { .. }) => {
+            follow_cell_path_recursive(Cow::Owned(Type::Record(f.clone())), path_members)
+        }
+
+        // Table to List (String)
+        (Type::Table(fields), PathMember::String { val, .. }) => {
+            let (_, sub_type) = fields.iter().find(|(name, _)| name == val)?;
+            let list_type = Type::List(Box::new(sub_type.clone()));
+            follow_cell_path_recursive(Cow::Owned(list_type), path_members)
+        }
+
+        (Type::List(_), PathMember::Int { .. }) => {
+            let next = match current {
+                Cow::Borrowed(Type::List(i)) => Cow::Borrowed(i.as_ref()),
+                Cow::Owned(Type::List(i)) => Cow::Owned(*i),
+                _ => unreachable!(),
+            };
+            follow_cell_path_recursive(next, path_members)
+        }
+
+        // List of Records indexed by key names
+        (Type::List(_), PathMember::String { .. }) => {
+            let next = match current {
+                Cow::Borrowed(Type::List(i)) => Cow::Borrowed(i.as_ref()),
+                Cow::Owned(Type::List(i)) => Cow::Owned(*i),
+                _ => unreachable!(),
+            };
+
+            let mut found_int_member = false;
+            let mut new_iter = std::iter::once(first).chain(path_members).filter(|pm| {
+                let first_int = !found_int_member && matches!(pm, PathMember::Int { .. });
+                if first_int {
+                    found_int_member = true;
+                }
+                !first_int
+            });
+            let inner_ty = follow_cell_path_recursive(next, &mut new_iter);
+
+            // If there's no int path member, need to wrap in a List type
+            // e.g. [{foo: bar}].foo -> [bar], list<record<foo: string>> -> list<string>
+            if found_int_member {
+                inner_ty
+            } else {
+                inner_ty.map(|inner_ty| Cow::Owned(Type::List(Box::new(inner_ty.into_owned()))))
+            }
+        }
+
+        _ => None,
+    }
+}
+
 impl Type {
     pub fn list(inner: Type) -> Self {
         Self::List(Box::new(inner))
@@ -294,6 +364,10 @@ impl Type {
             Type::Glob => String::from("glob"),
         }
     }
+
+    pub fn follow_cell_path<'a>(&'a self, path_members: &'a [PathMember]) -> Option<Cow<'a, Self>> {
+        follow_cell_path_recursive(Cow::Borrowed(self), &mut path_members.iter())
+    }
 }
 
 impl Display for Type {