LS: Use LookupItemEx universally and handle trait functions there (#…

…5805)
starkware-libs · Jun 14, 2024 · 57ffe18 · 57ffe18
1 parent fed0031
commit 57ffe18
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 35 deletions.
diff --git a/crates/cairo-lang-language-server/src/ide/completion/completions.rs b/crates/cairo-lang-language-server/src/ide/completion/completions.rs
@@ -1,6 +1,6 @@
 use cairo_lang_defs::ids::{
-    FunctionWithBodyId, ImplItemId, LanguageElementId, LookupItemId, ModuleFileId, ModuleId,
-    ModuleItemId, NamedLanguageElementId, TopLevelLanguageElementId, TraitFunctionId,
+    LanguageElementId, LookupItemId, ModuleFileId, ModuleId, NamedLanguageElementId,
+    TopLevelLanguageElementId, TraitFunctionId,
 };
 use cairo_lang_filesystem::ids::FileId;
 use cairo_lang_filesystem::span::TextOffset;
@@ -56,16 +56,8 @@ pub fn generic_completions(
     let Some(lookup_item_id) = lookup_items.into_iter().next() else {
         return completions;
     };
-    let function_id = match lookup_item_id {
-        LookupItemId::ModuleItem(ModuleItemId::FreeFunction(free_function_id)) => {
-            FunctionWithBodyId::Free(free_function_id)
-        }
-        LookupItemId::ImplItem(ImplItemId::Function(impl_function_id)) => {
-            FunctionWithBodyId::Impl(impl_function_id)
-        }
-        _ => {
-            return completions;
-        }
+    let Some(function_id) = lookup_item_id.function_with_body() else {
+        return completions;
     };
     let Ok(signature) = db.function_with_body_signature(function_id) else {
         return completions;

diff --git a/crates/cairo-lang-language-server/src/ide/hover/render/legacy.rs b/crates/cairo-lang-language-server/src/ide/hover/render/legacy.rs
@@ -1,8 +1,9 @@
 use cairo_lang_compiler::db::RootDatabase;
-use cairo_lang_defs::ids::{FunctionWithBodyId, ImplItemId, LookupItemId, ModuleItemId};
+use cairo_lang_defs::ids::{FunctionWithBodyId, LookupItemId};
 use cairo_lang_diagnostics::ToOption;
 use cairo_lang_semantic::db::SemanticGroup;
 use cairo_lang_semantic::items::function_with_body::SemanticExprLookup;
+use cairo_lang_semantic::lookup_item::LookupItemEx;
 use cairo_lang_semantic::Mutability;
 use cairo_lang_syntax::node::ast::{Expr, Pattern, TerminalIdentifier};
 use cairo_lang_syntax::node::kind::SyntaxKind;
@@ -20,17 +21,7 @@ use crate::markdown::Markdown;
 pub fn legacy(db: &RootDatabase, identifier: &TerminalIdentifier) -> Option<Hover> {
     let node = identifier.as_syntax_node();
     let lookup_item_id = db.find_lookup_item(&node)?;
-    let function_id = match lookup_item_id {
-        LookupItemId::ModuleItem(ModuleItemId::FreeFunction(free_function_id)) => {
-            FunctionWithBodyId::Free(free_function_id)
-        }
-        LookupItemId::ImplItem(ImplItemId::Function(impl_function_id)) => {
-            FunctionWithBodyId::Impl(impl_function_id)
-        }
-        _ => {
-            return None;
-        }
-    };
+    let function_id = lookup_item_id.function_with_body()?;
 
     // Build texts.
     let mut hints = Vec::new();

diff --git a/crates/cairo-lang-language-server/src/ide/semantic_highlighting/token_kind.rs b/crates/cairo-lang-language-server/src/ide/semantic_highlighting/token_kind.rs
@@ -1,6 +1,6 @@
-use cairo_lang_defs::ids::{FunctionWithBodyId, ImplItemId, LookupItemId, ModuleItemId};
 use cairo_lang_semantic::db::SemanticGroup;
 use cairo_lang_semantic::items::function_with_body::SemanticExprLookup;
+use cairo_lang_semantic::lookup_item::LookupItemEx;
 use cairo_lang_semantic::resolve::{ResolvedConcreteItem, ResolvedGenericItem};
 use cairo_lang_syntax::node::kind::SyntaxKind;
 use cairo_lang_syntax::node::utils::grandparent_kind;
@@ -171,16 +171,8 @@ impl SemanticTokenKind {
                 }
 
                 // Exprs and patterns..
-                let function_id = match lookup_item_id {
-                    LookupItemId::ModuleItem(ModuleItemId::FreeFunction(free_function_id)) => {
-                        FunctionWithBodyId::Free(free_function_id)
-                    }
-                    LookupItemId::ImplItem(ImplItemId::Function(impl_function_id)) => {
-                        FunctionWithBodyId::Impl(impl_function_id)
-                    }
-                    _ => {
-                        continue;
-                    }
+                let Some(function_id) = lookup_item_id.function_with_body() else {
+                    continue;
                 };
                 if let Some(expr_path_ptr) = expr_path_ptr {
                     if db.lookup_pattern_by_ptr(function_id, expr_path_ptr.into()).is_ok() {

diff --git a/crates/cairo-lang-semantic/src/lookup_item.rs b/crates/cairo-lang-semantic/src/lookup_item.rs
@@ -30,6 +30,9 @@ impl LookupItemEx for LookupItemId {
             LookupItemId::ModuleItem(ModuleItemId::FreeFunction(free_function_id)) => {
                 Some(FunctionWithBodyId::Free(*free_function_id))
             }
+            LookupItemId::TraitItem(TraitItemId::Function(trait_function_id)) => {
+                Some(FunctionWithBodyId::Trait(*trait_function_id))
+            }
             LookupItemId::ImplItem(ImplItemId::Function(impl_function_id)) => {
                 Some(FunctionWithBodyId::Impl(*impl_function_id))
             }