From bacd47adc064d7b4866c67933e3b6e273e1f2f3b Mon Sep 17 00:00:00 2001
From: shruti2522 <shruti.apc01@gmail.com>
Date: Tue, 2 Jul 2024 21:09:14 +0530
Subject: [PATCH] inlay hints for str, int, float, bool, any

Signed-off-by: shruti2522 <shruti.apc01@gmail.com>
---
 kclvm/tools/src/LSP/src/inlay_hints.rs | 60 +++++++++++---------------
 kclvm/tools/src/LSP/src/request.rs     |  2 +-
 2 files changed, 26 insertions(+), 36 deletions(-)

diff --git a/kclvm/tools/src/LSP/src/inlay_hints.rs b/kclvm/tools/src/LSP/src/inlay_hints.rs
index cd2b0dd34..c2363756d 100644
--- a/kclvm/tools/src/LSP/src/inlay_hints.rs
+++ b/kclvm/tools/src/LSP/src/inlay_hints.rs
@@ -1,10 +1,12 @@
+use kclvm_ast::ast::{self, Program};
+use kclvm_ast::pos::GetPos;
 use kclvm_error::Position as KCLPos;
 use kclvm_sema::core::{global_state::GlobalState, symbol::KCLSymbol};
 use kclvm_sema::ty::TypeKind;
 use lsp_types::{InlayHint, InlayHintLabelPart, Position as LspPosition, Range};
 use std::convert::TryInto;
 
-pub fn inlay_hints(file: &str, gs: &GlobalState) -> Option<Vec<InlayHint>> {
+pub fn inlay_hints(file: &str, gs: &GlobalState, program: &Program) -> Option<Vec<InlayHint>> {
     let mut inlay_hints: Vec<InlayHint> = vec![];
     let sema_db = gs.get_sema_db();
     if let Some(file_sema) = sema_db.get_file_sema(file) {
@@ -12,8 +14,10 @@ pub fn inlay_hints(file: &str, gs: &GlobalState) -> Option<Vec<InlayHint>> {
         for symbol_ref in symbols {
             if let Some(symbol) = gs.get_symbols().get_symbol(*symbol_ref) {
                 let (start, end) = symbol.get_range();
-                if let Some(hint) = generate_inlay_hint(symbol, gs, &start, &end) {
-                    inlay_hints.push(hint);
+                if has_type_assignment(program, &start) {
+                    if let Some(hint) = generate_inlay_hint(symbol, gs, &start, &end) {
+                        inlay_hints.push(hint);
+                    }
                 }
             }
         }
@@ -21,6 +25,22 @@ pub fn inlay_hints(file: &str, gs: &GlobalState) -> Option<Vec<InlayHint>> {
     Some(inlay_hints)
 }
 
+fn has_type_assignment(program: &Program, start: &KCLPos) -> bool {
+    if let Some(stmt_node) = program.pos_to_stmt(start) {
+        if let ast::Stmt::Assign(assign_stmt) = stmt_node.node {
+            if assign_stmt
+                .targets
+                .iter()
+                .any(|target| target.get_pos() == *start)
+                && assign_stmt.ty.is_none()
+            {
+                return true;
+            }
+        }
+    }
+    false
+}
+
 fn generate_inlay_hint(
     symbol: &KCLSymbol,
     gs: &GlobalState,
@@ -59,39 +79,9 @@ fn get_hint_label(symbol: &KCLSymbol, _gs: &GlobalState) -> Option<Vec<InlayHint
         let mut label_parts = Vec::new();
 
         match &ty.kind {
-            TypeKind::Str
-            | TypeKind::Bool
-            | TypeKind::Int
-            | TypeKind::Float
-            | TypeKind::Any
-            | TypeKind::None
-            | TypeKind::Named(_)
-            | TypeKind::NumberMultiplier(_)
-            | TypeKind::Union(_)
-            | TypeKind::Dict(_)
-            | TypeKind::List(_) => {
-                label_parts.push(InlayHintLabelPart {
-                    value: format!(": {}", ty.ty_str()),
-                    ..Default::default()
-                });
-            }
-            TypeKind::Module(module_ty) => {
-                label_parts.push(InlayHintLabelPart {
-                    value: format!(": {}", module_ty.pkgpath),
-                    ..Default::default()
-                });
-            }
-            TypeKind::Function(_) => {
-                let symbol_name = symbol.get_name().to_string();
-                label_parts.push(InlayHintLabelPart {
-                    value: format!("fn {}", symbol_name),
-                    ..Default::default()
-                });
-            }
-            TypeKind::Schema(schema_ty) => {
-                let fully_qualified_ty_name = format!("schema {}", schema_ty.name);
+            TypeKind::Str | TypeKind::Bool | TypeKind::Int | TypeKind::Float | TypeKind::Any => {
                 label_parts.push(InlayHintLabelPart {
-                    value: format!(": {}", fully_qualified_ty_name),
+                    value: format!("[: {}]", ty.ty_str()),
                     ..Default::default()
                 });
             }
diff --git a/kclvm/tools/src/LSP/src/request.rs b/kclvm/tools/src/LSP/src/request.rs
index 596d2810c..b4ad85a40 100644
--- a/kclvm/tools/src/LSP/src/request.rs
+++ b/kclvm/tools/src/LSP/src/request.rs
@@ -457,7 +457,7 @@ pub(crate) fn handle_inlay_hint(
         },
         Err(_) => return Ok(None),
     };
-    let res = inlay_hints(&file, &db.gs);
+    let res = inlay_hints(&file, &db.gs, &db.prog);
 
     if !snapshot.verify_request_version(db.version, &path)? {
         return Err(anyhow!(LSPError::Retry));