From a959f161e9197dbddb589a1937336eea6d07f642 Mon Sep 17 00:00:00 2001
From: Techatrix <19954306+Techatrix@users.noreply.github.com>
Date: Sun, 12 Mar 2023 06:24:54 +0000
Subject: [PATCH] find references on test decl with identifier name (#1051)

---
 src/analysis.zig                  | 24 ++++++++++++++----------
 src/references.zig                |  9 ++++++---
 tests/lsp_features/references.zig |  7 +++++++
 3 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/src/analysis.zig b/src/analysis.zig
index 411cd51..d53fbd6 100644
--- a/src/analysis.zig
+++ b/src/analysis.zig
@@ -276,7 +276,9 @@ pub fn isSnakeCase(name: []const u8) bool {
 
 pub fn getDeclNameToken(tree: Ast, node: Ast.Node.Index) ?Ast.TokenIndex {
     const tags = tree.nodes.items(.tag);
+    const datas = tree.nodes.items(.data);
     const main_token = tree.nodes.items(.main_token)[node];
+
     return switch (tags[node]) {
         // regular declaration names. + 1 to mut token because name comes after 'const'/'var'
         .local_var_decl,
@@ -307,21 +309,23 @@ pub fn getDeclNameToken(tree: Ast, node: Ast.Node.Index) ?Ast.TokenIndex {
         .identifier => main_token,
         .error_value => main_token + 2, // 'error'.<main_token +2>
 
-        // lhs of main token is name token, so use `node` - 1
-        .test_decl => if (tree.tokens.items(.tag)[main_token + 1] == .string_literal)
-            return main_token + 1
-        else
-            null,
+        .test_decl => datas[node].lhs,
+
         else => null,
     };
 }
 
 pub fn getDeclName(tree: Ast, node: Ast.Node.Index) ?[]const u8 {
-    const name = tree.tokenSlice(getDeclNameToken(tree, node) orelse return null);
-    return switch (tree.nodes.items(.tag)[node]) {
-        .test_decl => name[1 .. name.len - 1],
-        else => name,
-    };
+    const name_token = getDeclNameToken(tree, node) orelse return null;
+    const name = offsets.tokenToSlice(tree, name_token);
+
+    if (tree.nodes.items(.tag)[node] == .test_decl and
+        tree.tokens.items(.tag)[name_token] == .string_literal)
+    {
+        return name[1 .. name.len - 1];
+    }
+
+    return name;
 }
 
 fn resolveVarDeclAliasInternal(arena: std.mem.Allocator, store: *DocumentStore, node_handle: NodeWithHandle, root: bool) error{OutOfMemory}!?DeclWithHandle {
diff --git a/src/references.zig b/src/references.zig
index e32568f..08305da 100644
--- a/src/references.zig
+++ b/src/references.zig
@@ -91,12 +91,15 @@ const Builder = struct {
 
         const node_tags = tree.nodes.items(.tag);
         const datas = tree.nodes.items(.data);
-        const main_tokens = tree.nodes.items(.main_token);
+        const token_tags = tree.tokens.items(.tag);
         const starts = tree.tokens.items(.start);
 
         switch (node_tags[node]) {
-            .identifier => {
-                const identifier_token = main_tokens[node];
+            .identifier,
+            .test_decl,
+            => {
+                const identifier_token = analysis.getDeclNameToken(tree, node).?;
+                if (token_tags[identifier_token] != .identifier) return;
 
                 const child = (try analysis.lookupSymbolGlobal(
                     builder.allocator,
diff --git a/tests/lsp_features/references.zig b/tests/lsp_features/references.zig
index a9a04de..dcfb466 100644
--- a/tests/lsp_features/references.zig
+++ b/tests/lsp_features/references.zig
@@ -105,6 +105,13 @@ test "references - while continue expression" {
     );
 }
 
+test "references - test with identifier" {
+    try testReferences(
+        \\pub fn <0>() bool {}
+        \\test <0> {}
+    );
+}
+
 test "references - label" {
     if (true) return error.SkipZigTest; // TODO
     try testReferences(