From 378cc297a2410aa4e900d41e1b56989ace9de7ff Mon Sep 17 00:00:00 2001
From: he1pa <18012015693@163.com>
Date: Tue, 10 Sep 2024 20:17:00 +0800
Subject: [PATCH] fix: fix complete after some special expr. No symbol will be
 added after these expressions.

Signed-off-by: he1pa <18012015693@163.com>
---
 kclvm/sema/src/advanced_resolver/node.rs      | 53 ++++++++++---------
 kclvm/tools/src/LSP/src/completion.rs         |  8 +++
 ..._tests__complete_after_compare_expr_1.snap |  5 ++
 .../completion_test/dot/compare/compare.k     |  3 ++
 .../dot/special_expr/compare.k                |  3 ++
 5 files changed, 48 insertions(+), 24 deletions(-)
 create mode 100644 kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__complete_after_compare_expr_1.snap
 create mode 100644 kclvm/tools/src/LSP/src/test_data/completion_test/dot/compare/compare.k
 create mode 100644 kclvm/tools/src/LSP/src/test_data/completion_test/dot/special_expr/compare.k

diff --git a/kclvm/sema/src/advanced_resolver/node.rs b/kclvm/sema/src/advanced_resolver/node.rs
index 04044f9c0..7ad4e7a91 100644
--- a/kclvm/sema/src/advanced_resolver/node.rs
+++ b/kclvm/sema/src/advanced_resolver/node.rs
@@ -1051,33 +1051,38 @@ impl<'ctx> AdvancedResolver<'ctx> {
                 .borrow()
                 .get(&self.ctx.get_node_key(&expr.id))
             {
-                Some(ty) => {
-                    if let ast::Expr::Missing(_) = expr.node {
-                        return Ok(None);
-                    }
-                    let (_, end) = expr.get_span_pos();
-                    let mut expr_symbol = ExpressionSymbol::new(
-                        format!("@{}", expr.node.get_expr_name()),
-                        end.clone(),
-                        end,
-                        None,
-                    );
-                    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())
+                Some(ty) => match expr.node {
+                    ast::Expr::Missing(_)
+                    | ast::Expr::Binary(_)
+                    | ast::Expr::CompClause(_)
+                    | ast::Expr::Keyword(_)
+                    | ast::Expr::Arguments(_)
+                    | ast::Expr::Compare(_) => return Ok(None),
+                    _ => {
+                        let (_, end) = expr.get_span_pos();
+                        let mut expr_symbol = ExpressionSymbol::new(
+                            format!("@{}", expr.node.get_expr_name()),
+                            end.clone(),
+                            end,
+                            None,
+                        );
+                        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())
-                        }
-                    } else {
-                        Some(ty.clone())
-                    };
+                        };
 
-                    Ok(self.gs.get_symbols_mut().alloc_expression_symbol(
-                        expr_symbol,
-                        self.ctx.get_node_key(&expr.id),
-                        self.ctx.current_pkgpath.clone().unwrap(),
-                    ))
-                }
+                        Ok(self.gs.get_symbols_mut().alloc_expression_symbol(
+                            expr_symbol,
+                            self.ctx.get_node_key(&expr.id),
+                            self.ctx.current_pkgpath.clone().unwrap(),
+                        ))
+                    }
+                },
                 None => Ok(None),
             },
             res => res,
diff --git a/kclvm/tools/src/LSP/src/completion.rs b/kclvm/tools/src/LSP/src/completion.rs
index bfa4d703e..8f22771d1 100644
--- a/kclvm/tools/src/LSP/src/completion.rs
+++ b/kclvm/tools/src/LSP/src/completion.rs
@@ -2151,4 +2151,12 @@ mod tests {
         15,
         Some('.')
     );
+
+    completion_label_test_snapshot!(
+        complete_after_compare_expr_1,
+        "src/test_data/completion_test/dot/special_expr/compare.k",
+        2,
+        23,
+        Some('.')
+    );
 }
diff --git a/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__complete_after_compare_expr_1.snap b/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__complete_after_compare_expr_1.snap
new file mode 100644
index 000000000..bbb46f4f8
--- /dev/null
+++ b/kclvm/tools/src/LSP/src/snapshots/kcl_language_server__completion__tests__complete_after_compare_expr_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/test_data/completion_test/dot/compare/compare.k b/kclvm/tools/src/LSP/src/test_data/completion_test/dot/compare/compare.k
new file mode 100644
index 000000000..326459b62
--- /dev/null
+++ b/kclvm/tools/src/LSP/src/test_data/completion_test/dot/compare/compare.k
@@ -0,0 +1,3 @@
+v = ""
+if v == "" and "" == v :
+    a = 1
diff --git a/kclvm/tools/src/LSP/src/test_data/completion_test/dot/special_expr/compare.k b/kclvm/tools/src/LSP/src/test_data/completion_test/dot/special_expr/compare.k
new file mode 100644
index 000000000..326459b62
--- /dev/null
+++ b/kclvm/tools/src/LSP/src/test_data/completion_test/dot/special_expr/compare.k
@@ -0,0 +1,3 @@
+v = ""
+if v == "" and "" == v :
+    a = 1