bugfix: goto def (and hover) for complex select expr (#755)

Signed-off-by: He1pa <[email protected]>
kcl-lang · Oct 11, 2023 · c7c851b · c7c851b
1 parent e5f0482
commit c7c851b
Show file tree

Hide file tree

Showing 4 changed files with 127 additions and 32 deletions.
diff --git a/kclvm/tools/src/LSP/src/goto_def.rs b/kclvm/tools/src/LSP/src/goto_def.rs
@@ -128,46 +128,91 @@ pub(crate) fn find_def(
 
     let (inner_expr, parent) = inner_most_expr_in_stmt(&node.node, kcl_pos, None);
     if let Some(expr) = inner_expr {
-        if let Expr::Identifier(id) = expr.node {
-            let id_node = Node::node_with_pos(
-                id.clone(),
-                (
-                    expr.filename,
-                    expr.line,
-                    expr.column,
-                    expr.end_line,
-                    expr.end_column,
-                ),
-            );
-            let id = pre_process_identifier(id_node, kcl_pos);
-            match parent {
-                Some(schema_expr) => {
-                    if let Expr::Schema(schema_expr) = schema_expr.node {
-                        let schema_def =
-                            find_def(node, &schema_expr.name.get_end_pos(), prog_scope);
-                        if let Some(schema) = schema_def {
-                            match schema {
-                                Definition::Object(obj) => {
-                                    let schema_type = obj.ty.into_schema_type();
+        match expr.node {
+            Expr::Identifier(id) => {
+                let id_node = Node::node_with_pos(
+                    id.clone(),
+                    (
+                        expr.filename,
+                        expr.line,
+                        expr.column,
+                        expr.end_line,
+                        expr.end_column,
+                    ),
+                );
+                let id = pre_process_identifier(id_node, kcl_pos);
+                match parent {
+                    Some(schema_expr) => {
+                        if let Expr::Schema(schema_expr) = schema_expr.node {
+                            let schema_def =
+                                find_def(node, &schema_expr.name.get_end_pos(), prog_scope);
+                            if let Some(schema) = schema_def {
+                                match schema {
+                                    Definition::Object(obj) => match &obj.ty.kind {
+                                        kclvm_sema::ty::TypeKind::Schema(schema_type) => {
+                                            return find_attr_in_schema(
+                                                &schema_type,
+                                                &id.names,
+                                                &prog_scope.scope_map,
+                                            )
+                                        }
+                                        _ => {}
+                                    },
+                                    Definition::Scope(_) => {}
+                                }
+                            }
+                        }
+                    }
+                    None => {
+                        if let Some(inner_most_scope) = prog_scope.inner_most_scope(kcl_pos) {
+                            return resolve_var(
+                                &id.names,
+                                &inner_most_scope,
+                                &prog_scope.scope_map,
+                            );
+                        }
+                    }
+                }
+            }
+            Expr::Selector(select_expr) => {
+                if select_expr.attr.contains_pos(kcl_pos) {
+                    let value_def = find_def(node, &select_expr.value.get_end_pos(), prog_scope);
+                    let id = select_expr.attr;
+                    match value_def {
+                        Some(def) => match def {
+                            Definition::Object(obj) => match &obj.ty.kind {
+                                kclvm_sema::ty::TypeKind::Schema(schema_type) => {
                                     return find_attr_in_schema(
                                         &schema_type,
-                                        &id.names,
+                                        &id.node.names,
                                         &prog_scope.scope_map,
-                                    );
-                                }
-                                Definition::Scope(_) => {
-                                    //todo
+                                    )
                                 }
+                                _ => {}
+                            },
+                            Definition::Scope(_) => {}
+                        },
+                        None => {
+                            if let Some(inner_most_scope) = prog_scope.inner_most_scope(kcl_pos) {
+                                return resolve_var(
+                                    &id.node.names,
+                                    &inner_most_scope,
+                                    &prog_scope.scope_map,
+                                );
                             }
                         }
                     }
                 }
-                None => {
-                    if let Some(inner_most_scope) = prog_scope.inner_most_scope(kcl_pos) {
-                        return resolve_var(&id.names, &inner_most_scope, &prog_scope.scope_map);
+            }
+            Expr::Config(_) | Expr::ConfigIfEntry(_) => match parent {
+                Some(schema_expr) => {
+                    if let Expr::Schema(schema_expr) = schema_expr.node {
+                        return find_def(node, &schema_expr.name.get_end_pos(), prog_scope);
                     }
                 }
-            }
+                None => {}
+            },
+            _ => {}
         }
     }
     None
@@ -187,7 +232,6 @@ pub(crate) fn resolve_var(
         0 => None,
         1 => {
             let name = names[0].clone();
-
             match current_scope.lookup(&name) {
                 Some(obj) => match &obj.borrow().kind {
                     kclvm_sema::resolver::scope::ScopeObjectKind::Module(_) => {
@@ -920,4 +964,22 @@ mod tests {
         let res = goto_definition(&program, &pos, &prog_scope);
         compare_goto_res(res, (&file, 43, 4, 43, 9));
     }
+
+    #[test]
+    #[bench_test]
+    fn complex_select_goto_def() {
+        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+
+        let (file, program, prog_scope, _) =
+            compile_test_file("src/test_data/goto_def_test/goto_def.k");
+
+        let pos = KCLPos {
+            filename: file.clone(),
+            line: 52,
+            column: Some(22),
+        };
+
+        let res = goto_definition(&program, &pos, &prog_scope);
+        compare_goto_res(res, (&file, 43, 4, 43, 9));
+    }
 }
diff --git a/kclvm/tools/src/LSP/src/hover.rs b/kclvm/tools/src/LSP/src/hover.rs
@@ -357,4 +357,24 @@ mod tests {
             _ => unreachable!("test error"),
         }
     }
+
+    #[test]
+    #[bench_test]
+    fn complex_select_hover() {
+        let (file, program, prog_scope, _) = compile_test_file("src/test_data/hover_test/fib.k");
+        let pos = KCLPos {
+            filename: file.clone(),
+            line: 14,
+            column: Some(22),
+        };
+        let got = hover(&program, &pos, &prog_scope).unwrap();
+        match got.contents {
+            lsp_types::HoverContents::Scalar(marked_string) => {
+                if let MarkedString::String(s) = marked_string {
+                    assert_eq!(s, "value: int");
+                }
+            }
+            _ => unreachable!("test error"),
+        }
+    }
 }
diff --git a/kclvm/tools/src/LSP/src/test_data/hover_test/fib.k b/kclvm/tools/src/LSP/src/test_data/hover_test/fib.k
@@ -0,0 +1,14 @@
+schema Fib:
+    n1 = n - 1
+    n2 = n1 - 1
+    n: int
+    value: int
+
+    if n <= 1:
+        value = 1
+    elif n == 2:
+        value = 1
+    else:
+        value = Fib {n = n1}.value + Fib {n = n2}.value
+
+fib8 = Fib {n = 8}.value
diff --git a/kclvm/tools/src/LSP/src/util.rs b/kclvm/tools/src/LSP/src/util.rs
@@ -373,7 +373,6 @@ pub(crate) fn inner_most_expr(
     match &expr.node {
         Expr::Identifier(_) => (Some(expr.clone()), schema_def),
         Expr::Selector(select_expr) => {
-            walk_if_contains_with_new_expr!(select_expr.attr, pos, schema_def, Expr::Identifier);
             walk_if_contains!(select_expr.value, pos, schema_def);
             (Some(expr.clone()), schema_def)
         }