feat(LSP): suggest enum variants without parameters

tooling/lsp/src/requests/completion.rs

@@ -1,6 +1,7 @@
 use std::{
     collections::{BTreeMap, HashMap, HashSet},
     future::{self, Future},
+    ops::Deref,
 };
 
 use async_lsp::ResponseError;
@@ -47,7 +48,7 @@
 use super::{process_request, TraitReexport};
 
 mod auto_import;
 mod builtins;
 mod completion_items;
 mod kinds;
 mod sort_text;
@@ -199,15 +200,15 @@
         };
 
         let location = Location::new(span, self.file);
-        let Some(ReferenceId::Type(struct_id)) = self.interner.find_referenced(location) else {
+        let Some(ReferenceId::Type(type_id)) = self.interner.find_referenced(location) else {
             return;
         };
 
-        let struct_type = self.interner.get_type(struct_id);
-        let struct_type = struct_type.borrow();
+        let data_type = self.interner.get_type(type_id);
+        let data_type = data_type.borrow();
 
         // First get all of the struct's fields
-        let Some(fields) = struct_type.get_fields_as_written() else {
+        let Some(fields) = data_type.get_fields_as_written() else {
             return;
         };
 
@@ -223,7 +224,7 @@
             self.completion_items.push(self.struct_field_completion_item(
                 &field.name.0.contents,
                 &field.typ,
-                struct_type.id,
+                data_type.id,
                 *field_index,
                 self_prefix,
             ));
@@ -320,10 +321,11 @@
 
             match module_def_id {
                 ModuleDefId::ModuleId(id) => module_id = id,
-                ModuleDefId::TypeId(struct_id) => {
-                    let struct_type = self.interner.get_type(struct_id);
+                ModuleDefId::TypeId(type_id) => {
+                    let data_type = self.interner.get_type(type_id);
+                    self.complete_enum_variants_without_parameters(&data_type.borrow(), &prefix);
                     self.complete_type_methods(
-                        &Type::DataType(struct_type, vec![]),
+                        &Type::DataType(data_type, vec![]),
                         &prefix,
                         FunctionKind::Any,
                         function_completion_kind,
@@ -657,7 +659,7 @@
             return;
         };
 
-        let struct_id = get_type_struct_id(typ);
+        let type_id = get_type_type_id(typ);
         let is_primitive = typ.is_primitive();
         let has_self_param = matches!(function_kind, FunctionKind::SelfType(..));
 
@@ -669,15 +671,11 @@
             for (func_id, trait_id) in
                 methods.find_matching_methods(typ, has_self_param, self.interner)
             {
-                if let Some(struct_id) = struct_id {
+                if let Some(type_id) = type_id {
                     let modifiers = self.interner.function_modifiers(&func_id);
                     let visibility = modifiers.visibility;
-                    if !struct_member_is_visible(
-                        struct_id,
-                        visibility,
-                        self.module_id,
-                        self.def_maps,
-                    ) {
+                    if !struct_member_is_visible(type_id, visibility, self.module_id, self.def_maps)
+                    {
                         continue;
                     }
                 }
@@ -801,6 +799,23 @@
         }
     }
 
+    fn complete_enum_variants_without_parameters(&mut self, data_type: &DataType, prefix: &str) {
+        let Some(variants) = data_type.get_variants_as_written() else {
+            return;
+        };
+
+        for (index, variant) in variants.iter().enumerate() {
+            // Variants with parameters are represented as functions and are suggested in `complete_type_methods`
+            if variant.is_function || !name_matches(&variant.name.0.contents, prefix) {
+                continue;
+            }
+
+            let item =
+                self.enum_variant_completion_item(variant.name.to_string(), data_type.id, index);
+            self.completion_items.push(item);
+        }
+    }
+
     fn complete_struct_fields(
         &mut self,
         struct_type: &DataType,
@@ -1118,7 +1133,7 @@
 
     /// Try to suggest the name of a module to declare based on which
     /// files exist in the filesystem, excluding modules that are already declared.
     fn complete_module_delcaration(&mut self, module: &ModuleDeclaration) -> Option<()> {
         let filename = self.files.get_absolute_name(self.file).ok()?.into_path_buf();
 
         let is_main_lib_or_mod = filename.ends_with("main.nr")
@@ -1854,7 +1869,7 @@
             return;
         }
 
         self.complete_module_delcaration(module);
     }
 }
 
@@ -1900,13 +1915,13 @@
     }
 }
 
-fn get_type_struct_id(typ: &Type) -> Option<TypeId> {
-    match typ {
+fn get_type_type_id(typ: &Type) -> Option<TypeId> {
+    match typ.follow_bindings_shallow().deref() {
         Type::DataType(struct_type, _) => Some(struct_type.borrow().id),
         Type::Alias(type_alias, generics) => {
             let type_alias = type_alias.borrow();
             let typ = type_alias.get_type(generics);
-            get_type_struct_id(&typ)
+            get_type_type_id(&typ)
         }
         _ => None,
     }
@@ -1920,8 +1935,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
@@ -1958,7 +1973,7 @@
 fn module_def_id_from_reference_id(reference_id: ReferenceId) -> Option<ModuleDefId> {
     match reference_id {
         ReferenceId::Module(module_id) => Some(ModuleDefId::ModuleId(module_id)),
-        ReferenceId::Type(struct_id) => Some(ModuleDefId::TypeId(struct_id)),
+        ReferenceId::Type(type_id) => Some(ModuleDefId::TypeId(type_id)),
         ReferenceId::Trait(trait_id) => Some(ModuleDefId::TraitId(trait_id)),
         ReferenceId::Function(func_id) => Some(ModuleDefId::FunctionId(func_id)),
         ReferenceId::Alias(type_alias_id) => Some(ModuleDefId::TypeAliasId(type_alias_id)),

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

@@ -86,7 +86,14 @@ impl<'a> NodeFinder<'a> {
                 None,  // trait_id
                 false, // self_prefix
             ),
-            ModuleDefId::TypeId(struct_id) => vec![self.struct_completion_item(name, struct_id)],
+            ModuleDefId::TypeId(type_id) => {
+                let data_type = self.interner.get_type(type_id);
+                if data_type.borrow().is_struct() {
+                    vec![self.struct_completion_item(name, type_id)]
+                } else {
+                    vec![self.enum_completion_item(name, type_id)]
+                }
+            }
             ModuleDefId::TypeAliasId(id) => vec![self.type_alias_completion_item(name, id)],
             ModuleDefId::TraitId(trait_id) => vec![self.trait_completion_item(name, trait_id)],
             ModuleDefId::GlobalId(global_id) => vec![self.global_completion_item(name, global_id)],
@@ -106,14 +113,18 @@ impl<'a> NodeFinder<'a> {
         name: impl Into<String>,
         id: ModuleId,
     ) -> CompletionItem {
-        let completion_item = module_completion_item(name);
-        self.completion_item_with_doc_comments(ReferenceId::Module(id), completion_item)
+        let item = module_completion_item(name);
+        self.completion_item_with_doc_comments(ReferenceId::Module(id), item)
     }
 
-    fn struct_completion_item(&self, name: String, struct_id: TypeId) -> CompletionItem {
-        let completion_item =
-            simple_completion_item(name.clone(), CompletionItemKind::STRUCT, Some(name));
-        self.completion_item_with_doc_comments(ReferenceId::Type(struct_id), completion_item)
+    fn struct_completion_item(&self, name: String, type_id: TypeId) -> CompletionItem {
+        let items = simple_completion_item(name.clone(), CompletionItemKind::STRUCT, Some(name));
+        self.completion_item_with_doc_comments(ReferenceId::Type(type_id), items)
+    }
+
+    fn enum_completion_item(&self, name: String, type_id: TypeId) -> CompletionItem {
+        let item = simple_completion_item(name.clone(), CompletionItemKind::ENUM, Some(name));
+        self.completion_item_with_doc_comments(ReferenceId::Type(type_id), item)
     }
 
     pub(super) fn struct_field_completion_item(
@@ -124,33 +135,42 @@ impl<'a> NodeFinder<'a> {
         field_index: usize,
         self_type: bool,
     ) -> CompletionItem {
-        let completion_item = struct_field_completion_item(field, typ, self_type);
-        self.completion_item_with_doc_comments(
-            ReferenceId::StructMember(struct_id, field_index),
-            completion_item,
-        )
+        let item = struct_field_completion_item(field, typ, self_type);
+        let reference_id = ReferenceId::StructMember(struct_id, field_index);
+        self.completion_item_with_doc_comments(reference_id, item)
     }
 
     fn type_alias_completion_item(&self, name: String, id: TypeAliasId) -> CompletionItem {
-        let completion_item =
-            simple_completion_item(name.clone(), CompletionItemKind::STRUCT, Some(name));
-        self.completion_item_with_doc_comments(ReferenceId::Alias(id), completion_item)
+        let item = simple_completion_item(name.clone(), CompletionItemKind::STRUCT, Some(name));
+        self.completion_item_with_doc_comments(ReferenceId::Alias(id), item)
     }
 
     fn trait_completion_item(&self, name: String, trait_id: TraitId) -> CompletionItem {
-        let completion_item =
-            simple_completion_item(name.clone(), CompletionItemKind::INTERFACE, Some(name));
-        self.completion_item_with_doc_comments(ReferenceId::Trait(trait_id), completion_item)
+        let item = simple_completion_item(name.clone(), CompletionItemKind::INTERFACE, Some(name));
+        self.completion_item_with_doc_comments(ReferenceId::Trait(trait_id), item)
     }
 
     fn global_completion_item(&self, name: String, global_id: GlobalId) -> CompletionItem {
         let global = self.interner.get_global(global_id);
         let typ = self.interner.definition_type(global.definition_id);
         let description = typ.to_string();
+        let item = simple_completion_item(name, CompletionItemKind::CONSTANT, Some(description));
+        self.completion_item_with_doc_comments(ReferenceId::Global(global_id), item)
+    }
 
-        let completion_item =
-            simple_completion_item(name, CompletionItemKind::CONSTANT, Some(description));
-        self.completion_item_with_doc_comments(ReferenceId::Global(global_id), completion_item)
+    pub(super) fn enum_variant_completion_item(
+        &self,
+        name: String,
+        type_id: TypeId,
+        variant_index: usize,
+    ) -> CompletionItem {
+        let kind = CompletionItemKind::ENUM_MEMBER;
+        let item = simple_completion_item(name.clone(), kind, Some(name.clone()));
+        let item = completion_item_with_detail(item, name);
+        self.completion_item_with_doc_comments(
+            ReferenceId::EnumVariant(type_id, variant_index),
+            item,
+        )
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -354,6 +374,8 @@ impl<'a> NodeFinder<'a> {
         if let (Some(type_id), Some(variant_index)) =
             (func_meta.type_id, func_meta.enum_variant_index)
         {
+            completion_item.kind = Some(CompletionItemKind::ENUM_MEMBER);
+
             self.completion_item_with_doc_comments(
                 ReferenceId::EnumVariant(type_id, variant_index),
                 completion_item,

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

@@ -134,14 +134,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;
     }
@@ -304,7 +304,7 @@
             mod bar {
                 mod something {}
 
                 use super::foob>|<
             }
         "#;
 
@@ -1716,7 +1716,7 @@
     #[test]
     async fn test_auto_import_suggests_modules_too() {
         let src = r#"mod foo {
         pub mod barbaz {
             fn hello_world() {}
         }
     }
@@ -3094,6 +3094,7 @@
         assert_eq!(items.len(), 1);
 
         let item = &items[0];
+        assert_eq!(item.kind, Some(CompletionItemKind::ENUM_MEMBER));
         assert_eq!(item.label, "Variant(…)".to_string());
 
         let details = item.label_details.as_ref().unwrap();
@@ -3108,4 +3109,52 @@
         };
         assert!(markdown.value.contains("Some docs"));
     }
+
+    #[test]
+    async fn test_suggests_enum_variant_without_parameters() {
+        let src = r#"
+        enum Enum {
+            /// Some docs
+            Variant
+        }
+
+        fn foo() {
+            Enum::Var>|<
+        }
+        "#;
+        let items = get_completions(src).await;
+        assert_eq!(items.len(), 1);
+
+        let item = &items[0];
+        assert_eq!(item.kind, Some(CompletionItemKind::ENUM_MEMBER));
+        assert_eq!(item.label, "Variant".to_string());
+
+        let details = item.label_details.as_ref().unwrap();
+        assert_eq!(details.description, Some("Variant".to_string()));
+
+        assert_eq!(item.detail, Some("Variant".to_string()));
+        assert_eq!(item.insert_text, None);
+
+        let Documentation::MarkupContent(markdown) = item.documentation.as_ref().unwrap() else {
+            panic!("Expected markdown docs");
+        };
+        assert!(markdown.value.contains("Some docs"));
+    }
+
+    #[test]
+    async fn test_suggests_enum_type() {
+        let src = r#"
+        enum ThisIsAnEnum {
+        }
+
+        fn foo() {
+            ThisIsA>|<
+        }
+        "#;
+        let items = get_completions(src).await;
+        assert_eq!(items.len(), 1);
+
+        let item = &items[0];
+        assert_eq!(item.kind, Some(CompletionItemKind::ENUM));
+    }
 }