support enum completion

Author: CppCXY

crates/code_analysis/src/db_index/type/mod.rs

@@ -10,8 +10,7 @@ use flagset::FlagSet;
 use rowan::TextRange;
 use std::collections::HashMap;
 pub use type_assert::TypeAssertion;
-use type_decl::LuaDeclLocation;
-pub use type_decl::{LuaDeclTypeKind, LuaTypeAttribute, LuaTypeDecl, LuaTypeDeclId};
+pub use type_decl::{LuaDeclTypeKind, LuaTypeAttribute, LuaTypeDecl, LuaTypeDeclId, LuaDeclLocation};
 pub use types::*;
 
 #[derive(Debug)]

crates/emmylua_ls/src/handlers/completion/providers/type_special_provider.rs

@@ -1,14 +1,19 @@
 use code_analysis::{
-    InferGuard, LuaFunctionType, LuaMemberId, LuaPropertyOwnerId, LuaType, LuaTypeDeclId,
-    LuaUnionType,
+    InferGuard, LuaDeclLocation, LuaFunctionType, LuaMemberId, LuaMemberKey, LuaMemberOwner,
+    LuaPropertyOwnerId, LuaType, LuaTypeDeclId, LuaUnionType,
 };
 use emmylua_parser::{
-    LuaAstNode, LuaCallArgList, LuaCallExpr, LuaExpr, LuaSyntaxKind, LuaSyntaxToken, LuaTokenKind,
+    LuaAst, LuaAstNode, LuaAstToken, LuaCallArgList, LuaCallExpr, LuaComment, LuaExpr,
+    LuaNameToken, LuaSyntaxId, LuaSyntaxKind, LuaSyntaxToken, LuaTokenKind, LuaVarExpr,
 };
 use lsp_types::CompletionItem;
 
-use crate::handlers::{
-    completion::completion_builder::CompletionBuilder, signature_helper::get_current_param_index,
+use crate::{
+    handlers::{
+        completion::completion_builder::CompletionBuilder,
+        signature_helper::get_current_param_index,
+    },
+    util::humanize_type,
 };
 
 pub fn add_completion(builder: &mut CompletionBuilder) -> Option<()> {
@@ -68,10 +73,43 @@ fn add_type_ref_completion(
         }
         builder.stop_here();
     } else if type_decl.is_enum() {
-        // let member_ids = type_decl.get_enum_members()?.to_vec();
-        // for member_id in member_ids {
-        //     add_alias_member_completion(builder, &member_id);
-        // }
+        let owner_id = LuaMemberOwner::Type(type_ref_id.clone());
+        let member_map = builder
+            .semantic_model
+            .get_db()
+            .get_member_index()
+            .get_member_map(owner_id)?;
+
+        if type_decl.is_enum_key() {
+            let mut completion_items = Vec::new();
+            for member_key in member_map.keys() {
+                let label = match member_key {
+                    LuaMemberKey::Name(str) => to_enum_label(builder, str.as_str()),
+                    LuaMemberKey::Integer(i) => i.to_string(),
+                    LuaMemberKey::None => continue,
+                };
+
+                let completion_item = CompletionItem {
+                    label,
+                    kind: Some(lsp_types::CompletionItemKind::ENUM_MEMBER),
+                    ..Default::default()
+                };
+
+                completion_items.push(completion_item);
+            }
+            for completion_item in completion_items {
+                builder.add_completion_item(completion_item);
+            }
+        } else {
+            let locations = type_decl
+                .get_locations()
+                .iter()
+                .map(|it| it.clone())
+                .collect::<Vec<_>>();
+            let member_ids = member_map.values().map(|it| it.clone()).collect::<Vec<_>>();
+            add_enum_members_completion(builder, member_ids, &type_ref_id, locations);
+        }
+
         builder.stop_here();
     }
 
@@ -220,3 +258,130 @@ fn add_lambda_completion(builder: &mut CompletionBuilder, func: &LuaFunctionType
     builder.add_completion_item(completion_item);
     Some(())
 }
+
+fn add_enum_members_completion(
+    builder: &mut CompletionBuilder,
+    member_ids: Vec<LuaMemberId>,
+    type_id: &LuaTypeDeclId,
+    locations: Vec<LuaDeclLocation>,
+) -> Option<()> {
+    let file_id = builder.semantic_model.get_file_id();
+    let is_same_file = locations.iter().all(|it| it.file_id == file_id);
+    if let Some(variable_name) = get_enum_decl_variable_name(builder, locations, is_same_file) {
+        for member_id in member_ids {
+            let member = builder
+                .semantic_model
+                .get_db()
+                .get_member_index()
+                .get_member(&member_id)?;
+            let key = member.get_key();
+            let label = match key {
+                LuaMemberKey::Name(str) => format!("{}.{}", variable_name, str.to_string()),
+                LuaMemberKey::Integer(i) => format!("{}[{}]", variable_name, i),
+                LuaMemberKey::None => continue,
+            };
+
+            let description = format!("{}", type_id.get_name());
+            let completion_item = CompletionItem {
+                label,
+                kind: Some(lsp_types::CompletionItemKind::ENUM_MEMBER),
+                label_details: Some(lsp_types::CompletionItemLabelDetails {
+                    detail: None,
+                    description: Some(description.clone()),
+                }),
+                ..Default::default()
+            };
+
+            builder.add_completion_item(completion_item);
+        }
+    } else {
+        for member_id in member_ids {
+            let member = builder
+                .semantic_model
+                .get_db()
+                .get_member_index()
+                .get_member(&member_id)?;
+            let label = humanize_type(builder.semantic_model.get_db(), member.get_decl_type());
+            let description = format!("{}", type_id.get_name());
+            let completion_item = CompletionItem {
+                label,
+                kind: Some(lsp_types::CompletionItemKind::ENUM_MEMBER),
+                label_details: Some(lsp_types::CompletionItemLabelDetails {
+                    detail: None,
+                    description: Some(description.clone()),
+                }),
+                ..Default::default()
+            };
+
+            builder.add_completion_item(completion_item);
+        }
+    }
+    Some(())
+}
+
+fn get_enum_decl_variable_name(
+    builder: &CompletionBuilder,
+    locations: Vec<LuaDeclLocation>,
+    is_same_file: bool,
+) -> Option<String> {
+    let completion_file_id = builder.semantic_model.get_file_id();
+    if is_same_file {
+        let same_location = locations
+            .iter()
+            .find(|it| it.file_id == completion_file_id)?;
+        let root = builder
+            .semantic_model
+            .get_root_by_file_id(same_location.file_id)?;
+        let syntax_id = LuaSyntaxId::new(LuaTokenKind::TkName.into(), same_location.range);
+        let token = LuaNameToken::cast(syntax_id.to_token_from_root(root.syntax())?)?;
+        let comment = token.ancestors::<LuaComment>().next()?;
+        let comment_owner = comment.get_owner()?;
+        match comment_owner {
+            LuaAst::LuaLocalStat(local_stat) => {
+                return Some(
+                    local_stat
+                        .get_local_name_list()
+                        .next()?
+                        .get_name_token()?
+                        .get_name_text()
+                        .to_string(),
+                )
+            }
+            LuaAst::LuaAssignStat(assign_stat) => {
+                return Some(
+                    assign_stat
+                        .child::<LuaVarExpr>()?
+                        .syntax()
+                        .text()
+                        .to_string(),
+                )
+            }
+            _ => {}
+        }
+    } else {
+        for location in locations {
+            let root = builder
+                .semantic_model
+                .get_root_by_file_id(location.file_id)?;
+            let syntax_id = LuaSyntaxId::new(LuaTokenKind::TkName.into(), location.range);
+            let token = LuaNameToken::cast(syntax_id.to_token_from_root(root.syntax())?)?;
+            let comment = token.ancestors::<LuaComment>().next()?;
+            let comment_owner = comment.get_owner()?;
+            match comment_owner {
+                LuaAst::LuaLocalStat(_) => return None,
+                LuaAst::LuaAssignStat(assign_stat) => {
+                    return Some(
+                        assign_stat
+                            .child::<LuaVarExpr>()?
+                            .syntax()
+                            .text()
+                            .to_string(),
+                    );
+                }
+                _ => {}
+            }
+        }
+    }
+
+    None
+}