From b71a83b64a80c3bb0754b4952beb0aa73b88c6fc Mon Sep 17 00:00:00 2001
From: Gregory Anders <greg@gpanders.com>
Date: Sat, 7 May 2022 21:19:39 -0600
Subject: [PATCH] Limit references to current document for documentHighlight

---
 src/Server.zig     | 13 +++++++++-
 src/references.zig | 60 ++++++++++++++++++++++++----------------------
 src/rename.zig     |  2 +-
 3 files changed, 44 insertions(+), 31 deletions(-)

diff --git a/src/Server.zig b/src/Server.zig
index 1597f5e..a386f78 100644
--- a/src/Server.zig
+++ b/src/Server.zig
@@ -1151,6 +1151,7 @@ fn referencesDefinitionGlobal(
         &locs,
         std.ArrayList(types.Location).append,
         server.config.skip_std_references,
+        !highlight,
     );
 
     const result: types.ResponseParams = if (highlight) result: {
@@ -1185,7 +1186,17 @@ fn referencesDefinitionFieldAccess(
 
     const decl = (try server.getSymbolFieldAccess(handle, arena, position, range)) orelse return try respondGeneric(id, null_result_response);
     var locs = std.ArrayList(types.Location).init(arena.allocator());
-    try references.symbolReferences(arena, &server.document_store, decl, server.offset_encoding, include_decl, &locs, std.ArrayList(types.Location).append, server.config.skip_std_references);
+    try references.symbolReferences(
+        arena,
+        &server.document_store,
+        decl,
+        server.offset_encoding,
+        include_decl,
+        &locs,
+        std.ArrayList(types.Location).append,
+        server.config.skip_std_references,
+        !highlight,
+    );
     const result: types.ResponseParams = if (highlight) result: {
         var highlights = try std.ArrayList(types.DocumentHighlight).initCapacity(arena.allocator(), locs.items.len);
         for (locs.items) |loc| {
diff --git a/src/references.zig b/src/references.zig
index b5f58f7..fbd290f 100644
--- a/src/references.zig
+++ b/src/references.zig
@@ -492,7 +492,7 @@ fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentSto
     }
 }
 
-pub fn symbolReferences(arena: *std.heap.ArenaAllocator, store: *DocumentStore, decl_handle: analysis.DeclWithHandle, encoding: offsets.Encoding, include_decl: bool, context: anytype, comptime handler: anytype, skip_std_references: bool) !void {
+pub fn symbolReferences(arena: *std.heap.ArenaAllocator, store: *DocumentStore, decl_handle: analysis.DeclWithHandle, encoding: offsets.Encoding, include_decl: bool, context: anytype, comptime handler: anytype, skip_std_references: bool, workspace: bool) !void {
     std.debug.assert(decl_handle.decl.* != .label_decl);
     const curr_handle = decl_handle.handle;
     if (include_decl) {
@@ -503,41 +503,43 @@ pub fn symbolReferences(arena: *std.heap.ArenaAllocator, store: *DocumentStore,
         .ast_node => {
             try symbolReferencesInternal(arena, store, .{ .node = 0, .handle = curr_handle }, decl_handle, encoding, context, handler);
 
-            var imports = std.ArrayList(*DocumentStore.Handle).init(arena.allocator());
+            if (workspace) {
+                var imports = std.ArrayList(*DocumentStore.Handle).init(arena.allocator());
 
-            var handle_it = store.handles.iterator();
-            while (handle_it.next()) |entry| {
-                if (skip_std_references and std.mem.indexOf(u8, entry.key_ptr.*, "std") != null) {
-                    if (!include_decl or entry.value_ptr.* != curr_handle)
-                        continue;
-                }
+                var handle_it = store.handles.iterator();
+                while (handle_it.next()) |entry| {
+                    if (skip_std_references and std.mem.indexOf(u8, entry.key_ptr.*, "std") != null) {
+                        if (!include_decl or entry.value_ptr.* != curr_handle)
+                            continue;
+                    }
 
-                // Check entry's transitive imports
-                try imports.append(entry.value_ptr.*);
-                var i: usize = 0;
-                blk: while (i < imports.items.len) : (i += 1) {
-                    const import = imports.items[i];
-                    for (import.imports_used.items) |uri| {
-                        const h = store.getHandle(uri) orelse break;
+                    // Check entry's transitive imports
+                    try imports.append(entry.value_ptr.*);
+                    var i: usize = 0;
+                    blk: while (i < imports.items.len) : (i += 1) {
+                        const import = imports.items[i];
+                        for (import.imports_used.items) |uri| {
+                            const h = store.getHandle(uri) orelse break;
 
-                        if (h == curr_handle) {
-                            // entry does import curr_handle
-                            try symbolReferencesInternal(arena, store, .{ .node = 0, .handle = entry.value_ptr.* }, decl_handle, encoding, context, handler);
-                            break :blk;
-                        }
-
-                        select: {
-                            for (imports.items) |item| {
-                                if (item == h) {
-                                    // already checked this import
-                                    break :select;
-                                }
+                            if (h == curr_handle) {
+                                // entry does import curr_handle
+                                try symbolReferencesInternal(arena, store, .{ .node = 0, .handle = entry.value_ptr.* }, decl_handle, encoding, context, handler);
+                                break :blk;
+                            }
+
+                            select: {
+                                for (imports.items) |item| {
+                                    if (item == h) {
+                                        // already checked this import
+                                        break :select;
+                                    }
+                                }
+                                try imports.append(h);
                             }
-                            try imports.append(h);
                         }
                     }
+                    try imports.resize(0);
                 }
-                try imports.resize(0);
             }
         },
         .param_decl => |param| {
diff --git a/src/rename.zig b/src/rename.zig
index e5b4c18..8250de7 100644
--- a/src/rename.zig
+++ b/src/rename.zig
@@ -31,7 +31,7 @@ pub fn renameSymbol(arena: *std.heap.ArenaAllocator, store: *DocumentStore, decl
         .edits = edits,
         .allocator = arena.allocator(),
         .new_name = new_name,
-    }, refHandler, true);
+    }, refHandler, true, true);
 }
 
 pub fn renameLabel(arena: *std.heap.ArenaAllocator, decl_handle: analysis.DeclWithHandle, new_name: []const u8, edits: *std.StringHashMap([]types.TextEdit), encoding: offsets.Encoding) !void {