From 21231497e3335c7b3efdce66fb13ebc2d7e68706 Mon Sep 17 00:00:00 2001
From: He1pa <56333845+He1pa@users.noreply.github.com>
Date: Thu, 4 Jul 2024 12:32:20 +0800
Subject: [PATCH] fix: lsp completion depend on func return type for func call
 expr (#1472)

Signed-off-by: he1pa <18012015693@163.com>
---
 kclvm/sema/src/advanced_resolver/node.rs      | 11 ++++++++-
 kclvm/tools/src/LSP/src/completion.rs         | 24 +++++++++++++++++++
 ...__completion__tests__func_return_ty_1.snap |  5 ++++
 ...__completion__tests__func_return_ty_2.snap |  5 ++++
 ...__completion__tests__func_return_ty_3.snap |  5 ++++
 .../func_return/func_return_1/func_return_1.k |  4 ++++
 .../func_return/func_return_2/func_return_2.k |  8 +++++++
 .../func_return/func_return_3/func_return_3.k |  3 +++
 8 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_1.snap
 create mode 100644 kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_2.snap
 create mode 100644 kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_3.snap
 create mode 100644 kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_1/func_return_1.k
 create mode 100644 kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_2/func_return_2.k
 create mode 100644 kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_3/func_return_3.k

diff --git a/kclvm/sema/src/advanced_resolver/node.rs b/kclvm/sema/src/advanced_resolver/node.rs
index c7bce6cf7..359bf2289 100644
--- a/kclvm/sema/src/advanced_resolver/node.rs
+++ b/kclvm/sema/src/advanced_resolver/node.rs
@@ -910,7 +910,16 @@ impl<'ctx> AdvancedResolver<'ctx> {
                         end,
                         None,
                     );
-                    expr_symbol.sema_info.ty = Some(ty.clone());
+                    expr_symbol.sema_info.ty = if matches!(&expr.node, | ast::Expr::Call(_)) {
+                        if let TypeKind::Function(func_ty) = &ty.kind {
+                            Some(func_ty.return_ty.clone())
+                        } else {
+                            Some(ty.clone())
+                        }
+                    } else {
+                        Some(ty.clone())
+                    };
+
                     Ok(self.gs.get_symbols_mut().alloc_expression_symbol(
                         expr_symbol,
                         self.ctx.get_node_key(&expr.id),
diff --git a/kclvm/tools/src/LSP/src/completion.rs b/kclvm/tools/src/LSP/src/completion.rs
index 5df3efe94..288d15e18 100644
--- a/kclvm/tools/src/LSP/src/completion.rs
+++ b/kclvm/tools/src/LSP/src/completion.rs
@@ -2024,4 +2024,28 @@ mod tests {
         8,
         None
     );
+
+    completion_label_without_builtin_func_test_snapshot!(
+        func_return_ty_1,
+        "src/test_data/completion_test/dot/func_return/func_return_1/func_return_1.k",
+        4,
+        8,
+        Some('.')
+    );
+
+    completion_label_without_builtin_func_test_snapshot!(
+        func_return_ty_2,
+        "src/test_data/completion_test/dot/func_return/func_return_2/func_return_2.k",
+        8,
+        12,
+        Some('.')
+    );
+
+    completion_label_without_builtin_func_test_snapshot!(
+        func_return_ty_3,
+        "src/test_data/completion_test/dot/func_return/func_return_3/func_return_3.k",
+        3,
+        2,
+        Some('.')
+    );
 }
diff --git a/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_1.snap b/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_1.snap
new file mode 100644
index 000000000..bbb46f4f8
--- /dev/null
+++ b/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_1.snap
@@ -0,0 +1,5 @@
+---
+source: tools/src/LSP/src/completion.rs
+expression: "format!(\"{:?}\", got_labels)"
+---
+["capitalize(…)", "count(…)", "endswith(…)", "find(…)", "format(…)", "index(…)", "isalnum(…)", "isalpha(…)", "isdigit(…)", "islower(…)", "isspace(…)", "istitle(…)", "isupper(…)", "join(…)", "lower(…)", "lstrip(…)", "removeprefix(…)", "removesuffix(…)", "replace(…)", "rfind(…)", "rindex(…)", "rsplit(…)", "rstrip(…)", "split(…)", "splitlines(…)", "startswith(…)", "strip(…)", "title(…)", "upper(…)"]
diff --git a/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_2.snap b/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_2.snap
new file mode 100644
index 000000000..8b28f7593
--- /dev/null
+++ b/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_2.snap
@@ -0,0 +1,5 @@
+---
+source: tools/src/LSP/src/completion.rs
+expression: "format!(\"{:?}\", got_labels)"
+---
+["name"]
diff --git a/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_3.snap b/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_3.snap
new file mode 100644
index 000000000..f9d531ae0
--- /dev/null
+++ b/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__func_return_ty_3.snap
@@ -0,0 +1,5 @@
+---
+source: tools/src/LSP/src/completion.rs
+expression: "format!(\"{:?}\", got_labels)"
+---
+[]
diff --git a/kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_1/func_return_1.k b/kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_1/func_return_1.k
new file mode 100644
index 000000000..84932f666
--- /dev/null
+++ b/kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_1/func_return_1.k
@@ -0,0 +1,4 @@
+f: () -> str = lambda {
+    ""
+}
+a = f()
diff --git a/kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_2/func_return_2.k b/kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_2/func_return_2.k
new file mode 100644
index 000000000..bad592256
--- /dev/null
+++ b/kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_2/func_return_2.k
@@ -0,0 +1,8 @@
+schema Name:
+    name: str
+
+f1: (str) -> Name = lambda s: str -> Name {
+    Name {name: s}
+}
+
+b = f1("a")
diff --git a/kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_3/func_return_3.k b/kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_3/func_return_3.k
new file mode 100644
index 000000000..2159d44b7
--- /dev/null
+++ b/kclvm/tools/src/LSP/src/test_data/completion_test/dot/func_return/func_return_3/func_return_3.k
@@ -0,0 +1,3 @@
+f: () -> str = lambda {
+    ""
+}
\ No newline at end of file