From 6c978a3768e56c05581c33bd1b273f0ba4816871 Mon Sep 17 00:00:00 2001
From: shruti2522 <shruti.apc01@gmail.com>
Date: Mon, 3 Jun 2024 17:38:51 +0530
Subject: [PATCH] update semantic tokens with function type

Signed-off-by: shruti2522 <shruti.apc01@gmail.com>
---
 kclvm/sema/src/namer/node.rs              | 30 ++++++-----------------
 kclvm/tools/src/LSP/src/semantic_token.rs | 13 +++++++++-
 2 files changed, 19 insertions(+), 24 deletions(-)

diff --git a/kclvm/sema/src/namer/node.rs b/kclvm/sema/src/namer/node.rs
index 4b4acab46..ab923cef7 100644
--- a/kclvm/sema/src/namer/node.rs
+++ b/kclvm/sema/src/namer/node.rs
@@ -1,7 +1,6 @@
 use crate::core::package::ImportInfo;
 use crate::core::symbol::{
-    AttributeSymbol, FunctionSymbol, RuleSymbol, SchemaSymbol, SymbolKind, SymbolRef,
-    TypeAliasSymbol, ValueSymbol,
+    AttributeSymbol, RuleSymbol, SchemaSymbol, SymbolKind, SymbolRef, TypeAliasSymbol, ValueSymbol,
 };
 
 use super::Namer;
@@ -115,18 +114,10 @@ impl<'ctx> MutSelfTypedResultWalker<'ctx> for Namer<'ctx> {
                     .value_fully_qualified_name_set
                     .contains(&value_fully_qualified_name)
                 {
-                    let value_ref = if let ast::Expr::Lambda(_lambda_expr) = &assign_stmt.value.node
-                    {
-                        self.gs.get_symbols_mut().alloc_function_symbol(
-                            FunctionSymbol::new(value_name, start_pos, end_pos, Some(owner), true),
-                            self.ctx.get_node_key(&target.id),
-                        )
-                    } else {
-                        self.gs.get_symbols_mut().alloc_value_symbol(
-                            ValueSymbol::new(value_name, start_pos, end_pos, Some(owner), true),
-                            self.ctx.get_node_key(&target.id),
-                        )
-                    };
+                    let value_ref = self.gs.get_symbols_mut().alloc_value_symbol(
+                        ValueSymbol::new(value_name, start_pos, end_pos, Some(owner), true),
+                        self.ctx.get_node_key(&target.id),
+                    );
                     self.ctx
                         .value_fully_qualified_name_set
                         .insert(value_fully_qualified_name);
@@ -264,15 +255,8 @@ impl<'ctx> MutSelfTypedResultWalker<'ctx> for Namer<'ctx> {
         None
     }
 
-    fn walk_call_expr(&mut self, call_expr: &'ctx ast::CallExpr) -> Self::Result {
-        let (start_pos, end_pos): Range = call_expr.func.get_span_pos();
-        let owner = self.ctx.owner_symbols.last().unwrap().clone();
-        let func_name = call_expr.func.node.get_expr_name();
-        let func_ref = self.gs.get_symbols_mut().alloc_function_symbol(
-            FunctionSymbol::new(func_name, start_pos, end_pos, Some(owner), true),
-            self.ctx.get_node_key(&call_expr.func.id),
-        );
-        Some(vec![func_ref])
+    fn walk_call_expr(&mut self, _call_expr: &'ctx ast::CallExpr) -> Self::Result {
+        None
     }
 
     fn walk_subscript(&mut self, _subscript: &'ctx ast::Subscript) -> Self::Result {
diff --git a/kclvm/tools/src/LSP/src/semantic_token.rs b/kclvm/tools/src/LSP/src/semantic_token.rs
index c8fa82614..44ad37536 100644
--- a/kclvm/tools/src/LSP/src/semantic_token.rs
+++ b/kclvm/tools/src/LSP/src/semantic_token.rs
@@ -5,6 +5,7 @@ use kclvm_sema::core::{
     global_state::GlobalState,
     symbol::{KCLSymbol, SymbolKind},
 };
+use kclvm_sema::ty::TypeKind;
 use lsp_types::{SemanticToken, SemanticTokenType, SemanticTokens, SemanticTokensResult};
 
 pub const LEGEND_TYPE: &[SemanticTokenType] = &[
@@ -63,7 +64,17 @@ pub(crate) fn get_kind(ty: SymbolKind, symbol: &KCLSymbol, gs: &GlobalState) ->
         SymbolKind::Attribute => Some(type_index(SemanticTokenType::PROPERTY)),
         SymbolKind::Package => Some(type_index(SemanticTokenType::NAMESPACE)),
         SymbolKind::TypeAlias => Some(type_index(SemanticTokenType::TYPE)),
-        SymbolKind::Value => Some(type_index(SemanticTokenType::VARIABLE)),
+        SymbolKind::Value => {
+            // Check if the value is a lambda or a function type
+            if let Some(ty) = &symbol.get_sema_info().ty {
+                match ty.kind {
+                    TypeKind::Function(_) => Some(type_index(SemanticTokenType::FUNCTION)),
+                    _ => Some(type_index(SemanticTokenType::VARIABLE)),
+                }
+            } else {
+                Some(type_index(SemanticTokenType::VARIABLE))
+            }
+        }
         SymbolKind::Function => Some(type_index(SemanticTokenType::FUNCTION)),
         SymbolKind::Rule => Some(type_index(SemanticTokenType::MACRO)),
         SymbolKind::Unresolved => match &symbol.get_definition() {