fix: let LSP suggest fields and methods in LValue chains

tooling/lsp/src/requests/completion.rs

@@ -17,7 +17,7 @@
     ast::{
         AsTraitPath, AttributeTarget, BlockExpression, CallExpression, ConstructorExpression,
         Expression, ExpressionKind, ForLoopStatement, GenericTypeArgs, Ident, IfExpression,
-        ItemVisibility, Lambda, LetStatement, MemberAccessExpression, MethodCallExpression,
+        ItemVisibility, LValue, Lambda, LetStatement, MemberAccessExpression, MethodCallExpression,
         NoirFunction, NoirStruct, NoirTraitImpl, Path, PathKind, Pattern, Statement,
         TraitImplItemKind, TypeImpl, UnresolvedGeneric, UnresolvedGenerics, UnresolvedType,
         UnresolvedTypeData, UseTree, UseTreeKind, Visitor,
@@ -41,7 +41,7 @@
 use super::process_request;
 
 mod auto_import;
 mod builtins;
 mod completion_items;
 mod kinds;
 mod sort_text;
@@ -233,7 +233,7 @@
         let mut idents: Vec<Ident> = Vec::new();
 
         // Find in which ident we are in, and in which part of it
         // (it could be that we are completting in the middle of an ident)
         for segment in &path.segments {
             let ident = &segment.ident;
 
@@ -387,7 +387,7 @@
                 let description = if let Some(ReferenceId::Local(definition_id)) =
                     self.interner.reference_at_location(location)
                 {
-                    let typ = self.interner.definition_type(definition_id);
+                    let typ = self.interner.definition_type(definition_id).follow_bindings();
                     Some(typ.to_string())
                 } else {
                     None
@@ -551,6 +551,7 @@
         function_completion_kind: FunctionCompletionKind,
         self_prefix: bool,
     ) {
+        let typ = &typ.follow_bindings();
         match typ {
             Type::Struct(struct_type, generics) => {
                 self.complete_struct_fields(&struct_type.borrow(), generics, prefix, self_prefix);
@@ -932,7 +933,8 @@
                     if let Some(ReferenceId::Local(definition_id)) =
                         self.interner.find_referenced(location)
                     {
-                        self.self_type = Some(self.interner.definition_type(definition_id));
+                        self.self_type =
+                            Some(self.interner.definition_type(definition_id).follow_bindings());
                     }
                 }
             }
@@ -941,6 +943,32 @@
         }
     }
 
+    fn get_lvalue_type(&self, lvalue: &LValue) -> Option<Type> {
+        match lvalue {
+            LValue::Ident(ident) => {
+                let location = Location::new(ident.span(), self.file);
+                if let Some(ReferenceId::Local(definition_id)) =
+                    self.interner.find_referenced(location)
+                {
+                    let typ = self.interner.definition_type(definition_id).follow_bindings();
+                    Some(typ)
+                } else {
+                    None
+                }
+            }
+            LValue::MemberAccess { object, field_name, .. } => {
+                let typ = self.get_lvalue_type(object)?;
+                get_field_type(&typ, &field_name.0.contents)
+            }
+            LValue::Index { array, .. } => {
+                let typ = self.get_lvalue_type(array)?;
+                get_array_element_type(typ)
+            }
+            LValue::Dereference(lvalue, ..) => self.get_lvalue_type(lvalue),
+            LValue::Interned(..) => None,
+        }
+    }
+
     fn includes_span(&self, span: Span) -> bool {
         span.start() as usize <= self.byte_index && self.byte_index <= span.end() as usize
     }
@@ -1258,11 +1286,75 @@
     }
 
     fn visit_lvalue_ident(&mut self, ident: &Ident) {
+        // If we have `foo.>|<` we suggest `foo`'s type fields and methods
         if self.byte == Some(b'.') && ident.span().end() as usize == self.byte_index - 1 {
             let location = Location::new(ident.span(), self.file);
             if let Some(ReferenceId::Local(definition_id)) = self.interner.find_referenced(location)
             {
-                let typ = self.interner.definition_type(definition_id);
+                let typ = self.interner.definition_type(definition_id).follow_bindings();
+                let prefix = "";
+                let self_prefix = false;
+                self.complete_type_fields_and_methods(
+                    &typ,
+                    prefix,
+                    FunctionCompletionKind::NameAndParameters,
+                    self_prefix,
+                );
+            }
+        }
+    }
+
+    fn visit_lvalue_member_access(
+        &mut self,
+        object: &LValue,
+        field_name: &Ident,
+        span: Span,
+    ) -> bool {
+        // If we have `foo.bar.>|<` we solve the type of `foo`, get the field `bar`,
+        // then suggest methods of the resulting type.
+        if self.byte == Some(b'.') && span.end() as usize == self.byte_index - 1 {
+            if let Some(typ) = self.get_lvalue_type(object) {
+                if let Some(typ) = get_field_type(&typ, &field_name.0.contents) {
+                    let prefix = "";
+                    let self_prefix = false;
+                    self.complete_type_fields_and_methods(
+                        &typ,
+                        prefix,
+                        FunctionCompletionKind::NameAndParameters,
+                        self_prefix,
+                    );
+                }
+            }
+
+            return false;
+        }
+        true
+    }
+
+    fn visit_lvalue_index(&mut self, array: &LValue, _index: &Expression, span: Span) -> bool {
+        // If we have `foo[index].>|<` we solve the type of `foo`, then get the array/slice element type,
+        // then suggest methods of that type.
+        if self.byte == Some(b'.') && span.end() as usize == self.byte_index - 1 {
+            if let Some(typ) = self.get_lvalue_type(array) {
+                if let Some(typ) = get_array_element_type(typ) {
+                    let prefix = "";
+                    let self_prefix = false;
+                    self.complete_type_fields_and_methods(
+                        &typ,
+                        prefix,
+                        FunctionCompletionKind::NameAndParameters,
+                        self_prefix,
+                    );
+                }
+            }
+            return false;
+        }
+        true
+    }
+
+    fn visit_lvalue_dereference(&mut self, lvalue: &LValue, span: Span) -> bool {
+        if self.byte == Some(b'.') && span.end() as usize == self.byte_index - 1 {
+            if let Some(typ) = self.get_lvalue_type(lvalue) {
                 let prefix = "";
                 let self_prefix = false;
                 self.complete_type_fields_and_methods(
@@ -1272,7 +1364,10 @@
                     self_prefix,
                 );
             }
+            return false;
         }
+
+        true
     }
 
     fn visit_variable(&mut self, path: &Path, _: Span) -> bool {
@@ -1443,6 +1538,34 @@
     }
 }
 
+fn get_field_type(typ: &Type, name: &str) -> Option<Type> {
+    match typ {
+        Type::Struct(struct_type, generics) => {
+            Some(struct_type.borrow().get_field(name, generics)?.0)
+        }
+        Type::Tuple(types) => {
+            if let Ok(index) = name.parse::<i32>() {
+                types.get(index as usize).cloned()
+            } else {
+                None
+            }
+        }
+        Type::Alias(alias_type, generics) => Some(alias_type.borrow().get_type(generics)),
+        _ => None,
+    }
+}
+
+fn get_array_element_type(typ: Type) -> Option<Type> {
+    match typ.follow_bindings() {
+        Type::Array(_, typ) | Type::Slice(typ) => Some(*typ),
+        Type::Alias(alias_type, generics) => {
+            let typ = alias_type.borrow().get_type(&generics);
+            get_array_element_type(typ)
+        }
+        _ => None,
+    }
+}
+
 /// Returns true if name matches a prefix written in code.
 /// `prefix` must already be in snake case.
 /// This method splits both name and prefix by underscore,
@@ -1451,8 +1574,8 @@
 ///
 /// For example:
 ///
 /// // "merk" and "ro" match "merkle" and "root" and are in order
 /// name_matches("compute_merkle_root", "merk_ro") == true
 ///
 /// // "ro" matches "root", but "merkle" comes before it, so no match
 /// name_matches("compute_merkle_root", "ro_mer") == false

tooling/lsp/src/requests/completion/tests.rs

@@ -133,14 +133,14 @@
     #[test]
     async fn test_use_first_segment() {
         let src = r#"
             mod foobaz {}
             mod foobar {}
             use foob>|<
         "#;
 
         assert_completion(
             src,
             vec![module_completion_item("foobaz"), module_completion_item("foobar")],
         )
         .await;
     }
@@ -303,7 +303,7 @@
             mod bar {
                 mod something {}
 
                 use super::foob>|<
             }
         "#;
 
@@ -1546,7 +1546,7 @@
     async fn test_auto_import_suggests_modules_too() {
         let src = r#"
             mod foo {
                 mod barbaz {
                     fn hello_world() {}
                 }
             }
@@ -1559,7 +1559,7 @@
         assert_eq!(items.len(), 1);
 
         let item = &items[0];
         assert_eq!(item.label, "barbaz");
         assert_eq!(
             item.label_details,
             Some(CompletionItemLabelDetails {
@@ -2015,4 +2015,108 @@
         )
         .await;
     }
+
+    #[test]
+    async fn test_suggests_when_assignment_follows_in_chain_1() {
+        let src = r#"
+        struct Foo {
+            bar: Bar
+        }
+
+        struct Bar {
+            baz: Field
+        }
+
+        fn f(foo: Foo) {
+            let mut x = 1;
+
+            foo.bar.>|<
+
+            x = 2;
+        }"#;
+
+        assert_completion(src, vec![field_completion_item("baz", "Field")]).await;
+    }
+
+    #[test]
+    async fn test_suggests_when_assignment_follows_in_chain_2() {
+        let src = r#"
+        struct Foo {
+            bar: Bar
+        }
+
+        struct Bar {
+            baz: Baz
+        }
+
+        struct Baz {
+            qux: Field
+        }
+
+        fn f(foo: Foo) {
+            let mut x = 1;
+
+            foo.bar.baz.>|<
+
+            x = 2;
+        }"#;
+
+        assert_completion(src, vec![field_completion_item("qux", "Field")]).await;
+    }
+
+    #[test]
+    async fn test_suggests_when_assignment_follows_in_chain_3() {
+        let src = r#"
+        struct Foo {
+            foo: Field
+        }
+
+        fn execute() {
+            let a = Foo { foo: 1 };
+            a.>|<
+
+            x = 1;
+        }"#;
+
+        assert_completion(src, vec![field_completion_item("foo", "Field")]).await;
+    }
+
+    #[test]
+    async fn test_suggests_when_assignment_follows_in_chain_4() {
+        let src = r#"
+        struct Foo {
+            bar: Bar
+        }
+
+        struct Bar {
+            baz: Field
+        }
+
+        fn execute() {
+            let foo = Foo { foo: 1 };
+            foo.bar.>|<
+
+            x = 1;
+        }"#;
+
+        assert_completion(src, vec![field_completion_item("baz", "Field")]).await;
+    }
+
+    #[test]
+    async fn test_suggests_when_assignment_follows_in_chain_with_index() {
+        let src = r#"
+        struct Foo {
+            bar: Field
+        }
+
+        fn f(foos: [Foo; 3]) {
+            let mut x = 1;
+
+            foos[0].>|<
+
+            x = 2;
+        }"#;
+
+        assert_completion(src, vec![field_completion_item("bar", "Field")]).await;
+    }
 }