From 4b8502beebec9e7faeca5128a8f579aed28de3f0 Mon Sep 17 00:00:00 2001
From: "xiarui.xr" <xiarui1994@gmail.com>
Date: Fri, 13 Oct 2023 16:56:35 +0800
Subject: [PATCH] fix find refs from schema attr def

Signed-off-by: xiarui.xr <xiarui1994@gmail.com>
---
 kclvm/tools/src/LSP/src/find_refs.rs | 14 +++++++-------
 kclvm/tools/src/LSP/src/goto_def.rs  | 12 ++++--------
 kclvm/tools/src/LSP/src/util.rs      | 18 ++++++++++++++++++
 3 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/kclvm/tools/src/LSP/src/find_refs.rs b/kclvm/tools/src/LSP/src/find_refs.rs
index aaa162b6d..2697df3b4 100644
--- a/kclvm/tools/src/LSP/src/find_refs.rs
+++ b/kclvm/tools/src/LSP/src/find_refs.rs
@@ -226,13 +226,13 @@ mod tests {
                     },
                 };
                 let expect = vec![
-                    // Location {
-                    //     uri: url.clone(),
-                    //     range: Range {
-                    //         start: Position::new(5, 4),
-                    //         end: Position::new(5, 8),
-                    //     },
-                    // },
+                    Location {
+                        uri: url.clone(),
+                        range: Range {
+                            start: Position::new(5, 4),
+                            end: Position::new(5, 8),
+                        },
+                    },
                     Location {
                         uri: url.clone(),
                         range: Range {
diff --git a/kclvm/tools/src/LSP/src/goto_def.rs b/kclvm/tools/src/LSP/src/goto_def.rs
index 5e957020d..810b19a8b 100644
--- a/kclvm/tools/src/LSP/src/goto_def.rs
+++ b/kclvm/tools/src/LSP/src/goto_def.rs
@@ -147,7 +147,7 @@ pub(crate) fn find_def(
                     ),
                 );
                 let id = pre_process_identifier(id_node, kcl_pos);
-                                match parent {
+                match parent {
                     Some(schema_expr) => {
                         if let Expr::Schema(schema_expr) = schema_expr.node {
                             let schema_def =
@@ -171,7 +171,7 @@ pub(crate) fn find_def(
                     }
                     None => {
                         if let Some(inner_most_scope) = prog_scope.inner_most_scope(kcl_pos) {
-                                                        return resolve_var(
+                            return resolve_var(
                                 &id.names,
                                 &inner_most_scope,
                                 &prog_scope.scope_map,
@@ -220,7 +220,7 @@ pub(crate) fn find_def(
             },
             _ => {}
         }
-        }
+    }
     None
 }
 
@@ -1018,10 +1018,6 @@ mod tests {
         compare_goto_res(res, (&file, 43, 4, 43, 9));
     }
 
-// this test case if to ensure the go-to-def result of schema attribute def
-// schema Person {
-//     target: str
-// }
     #[test]
     #[bench_test]
     fn schema_attribute_def_goto_def() {
@@ -1037,6 +1033,6 @@ mod tests {
         };
 
         let res = goto_definition(&program, &pos, &prog_scope);
-        compare_goto_res(res, (&file, 19, 4, 19, 8));
+        compare_goto_res(res, (&file, 18, 4, 18, 8));
     }
 }
diff --git a/kclvm/tools/src/LSP/src/util.rs b/kclvm/tools/src/LSP/src/util.rs
index 3a8535450..e93f184e5 100644
--- a/kclvm/tools/src/LSP/src/util.rs
+++ b/kclvm/tools/src/LSP/src/util.rs
@@ -333,6 +333,24 @@ pub(crate) fn inner_most_expr_in_stmt(
             (None, schema_def)
         }
         Stmt::SchemaAttr(schema_attr_expr) => {
+            walk_if_contains!(
+                Node::node_with_pos(
+                    Expr::Identifier(Identifier {
+                        names: vec![*schema_attr_expr.name.clone()],
+                        pkgpath: "".to_string(),
+                        ctx: kclvm_ast::ast::ExprContext::Load,
+                    }),
+                    (
+                        schema_attr_expr.name.filename.clone(),
+                        schema_attr_expr.name.line,
+                        schema_attr_expr.name.column,
+                        schema_attr_expr.name.end_line,
+                        schema_attr_expr.name.end_column,
+                    ),
+                ),
+                pos,
+                schema_def
+            );
             if schema_attr_expr.ty.contains_pos(pos) {
                 return (
                     build_identifier_from_ty_string(&schema_attr_expr.ty, pos),