References improvements (#653)

src/Server.zig

@@ -9,7 +9,6 @@ const types = @import("types.zig");
 const analysis = @import("analysis.zig");
 const ast = @import("ast.zig");
 const references = @import("references.zig");
-const rename = @import("rename.zig");
 const offsets = @import("offsets.zig");
 const semantic_tokens = @import("semantic_tokens.zig");
 const inlay_hints = @import("inlay_hints.zig");
@@ -986,206 +985,6 @@ fn gotoDefinitionString(
     });
 }
 
-fn renameDefinitionGlobal(
-    server: *Server,
-    writer: anytype,
-    id: types.RequestId,
-    handle: *DocumentStore.Handle,
-    pos_index: usize,
-    new_name: []const u8,
-) !void {
-    const tracy_zone = tracy.trace(@src());
-    defer tracy_zone.end();
-
-    const decl = (try server.getSymbolGlobal(pos_index, handle)) orelse return try respondGeneric(writer, id, null_result_response);
-
-    var workspace_edit = types.WorkspaceEdit{ .changes = .{} };
-    try rename.renameSymbol(&server.arena, &server.document_store, decl, new_name, &workspace_edit.changes, server.offset_encoding);
-
-    try send(writer, server.arena.allocator(), types.Response{
-        .id = id,
-        .result = .{ .WorkspaceEdit = workspace_edit },
-    });
-}
-
-fn renameDefinitionFieldAccess(
-    server: *Server,
-    writer: anytype,
-    id: types.RequestId,
-    handle: *DocumentStore.Handle,
-    source_index: usize,
-    loc: offsets.Loc,
-    new_name: []const u8,
-) !void {
-    const tracy_zone = tracy.trace(@src());
-    defer tracy_zone.end();
-
-    const decl = (try server.getSymbolFieldAccess(handle, source_index, loc)) orelse return try respondGeneric(writer, id, null_result_response);
-
-    var workspace_edit = types.WorkspaceEdit{ .changes = .{} };
-    try rename.renameSymbol(&server.arena, &server.document_store, decl, new_name, &workspace_edit.changes, server.offset_encoding);
-
-    try send(writer, server.arena.allocator(), types.Response{
-        .id = id,
-        .result = .{ .WorkspaceEdit = workspace_edit },
-    });
-}
-
-fn renameDefinitionLabel(
-    server: *Server,
-    writer: anytype,
-    id: types.RequestId,
-    handle: *DocumentStore.Handle,
-    pos_index: usize,
-    new_name: []const u8,
-) !void {
-    const tracy_zone = tracy.trace(@src());
-    defer tracy_zone.end();
-
-    const decl = (try getLabelGlobal(pos_index, handle)) orelse return try respondGeneric(writer, id, null_result_response);
-
-    var workspace_edit = types.WorkspaceEdit{ .changes = .{} };
-    try rename.renameLabel(&server.arena, decl, new_name, &workspace_edit.changes, server.offset_encoding);
-
-    try send(writer, server.arena.allocator(), types.Response{
-        .id = id,
-        .result = .{ .WorkspaceEdit = workspace_edit },
-    });
-}
-
-fn referencesDefinitionGlobal(
-    server: *Server,
-    writer: anytype,
-    id: types.RequestId,
-    handle: *DocumentStore.Handle,
-    pos_index: usize,
-    include_decl: bool,
-    comptime highlight: bool,
-) !void {
-    const tracy_zone = tracy.trace(@src());
-    defer tracy_zone.end();
-
-    const decl = (try server.getSymbolGlobal(pos_index, handle)) orelse return try respondGeneric(writer, id, null_result_response);
-    var locs = std.ArrayList(types.Location).init(server.arena.allocator());
-    try references.symbolReferences(
-        &server.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 = std.ArrayListUnmanaged(types.DocumentHighlight){};
-        try highlights.ensureTotalCapacity(server.arena.allocator(), locs.items.len);
-        const uri = handle.uri();
-        for (locs.items) |loc| {
-            if (std.mem.eql(u8, loc.uri, uri)) {
-                highlights.appendAssumeCapacity(.{
-                    .range = loc.range,
-                    .kind = .Text,
-                });
-            }
-        }
-        break :result .{ .DocumentHighlight = highlights.items };
-    } else .{ .Locations = locs.items };
-
-    try send(writer, server.arena.allocator(), types.Response{
-        .id = id,
-        .result = result,
-    });
-}
-
-fn referencesDefinitionFieldAccess(
-    server: *Server,
-    writer: anytype,
-    id: types.RequestId,
-    handle: *DocumentStore.Handle,
-    source_index: usize,
-    loc: offsets.Loc,
-    include_decl: bool,
-    comptime highlight: bool,
-) !void {
-    const tracy_zone = tracy.trace(@src());
-    defer tracy_zone.end();
-
-    var allocator = server.arena.allocator();
-
-    const decl = (try server.getSymbolFieldAccess(handle, source_index, loc)) orelse return try respondGeneric(writer, id, null_result_response);
-    var locations = std.ArrayList(types.Location).init(allocator);
-    try references.symbolReferences(
-        &server.arena,
-        &server.document_store,
-        decl,
-        server.offset_encoding,
-        include_decl,
-        &locations,
-        std.ArrayList(types.Location).append,
-        server.config.skip_std_references,
-        !highlight,
-    );
-    const result: types.ResponseParams = if (highlight) result: {
-        var highlights = std.ArrayListUnmanaged(types.DocumentHighlight){};
-        try highlights.ensureTotalCapacity(allocator, locations.items.len);
-        const uri = handle.uri();
-        for (locations.items) |location| {
-            if (std.mem.eql(u8, location.uri, uri)) {
-                highlights.appendAssumeCapacity(.{
-                    .range = location.range,
-                    .kind = .Text,
-                });
-            }
-        }
-        break :result .{ .DocumentHighlight = highlights.items };
-    } else .{ .Locations = locations.items };
-    try send(writer, allocator, types.Response{
-        .id = id,
-        .result = result,
-    });
-}
-
-fn referencesDefinitionLabel(
-    server: *Server,
-    writer: anytype,
-    id: types.RequestId,
-    handle: *DocumentStore.Handle,
-    pos_index: usize,
-    include_decl: bool,
-    comptime highlight: bool,
-) !void {
-    const tracy_zone = tracy.trace(@src());
-    defer tracy_zone.end();
-
-    var allocator = server.arena.allocator();
-
-    const decl = (try getLabelGlobal(pos_index, handle)) orelse return try respondGeneric(writer, id, null_result_response);
-    var locs = std.ArrayList(types.Location).init(allocator);
-    try references.labelReferences(&server.arena, decl, server.offset_encoding, include_decl, &locs, std.ArrayList(types.Location).append);
-    const result: types.ResponseParams = if (highlight) result: {
-        var highlights = std.ArrayListUnmanaged(types.DocumentHighlight){};
-        try highlights.ensureTotalCapacity(allocator, locs.items.len);
-        const uri = handle.uri();
-        for (locs.items) |loc| {
-            if (std.mem.eql(u8, loc.uri, uri)) {
-                highlights.appendAssumeCapacity(.{
-                    .range = loc.range,
-                    .kind = .Text,
-                });
-            }
-        }
-        break :result .{ .DocumentHighlight = highlights.items };
-    } else .{ .Locations = locs.items };
-
-    try send(writer, allocator, types.Response{
-        .id = id,
-        .result = result,
-    });
-}
-
 fn hasComment(tree: Ast.Tree, start_token: Ast.TokenIndex, end_token: Ast.TokenIndex) bool {
     const tracy_zone = tracy.trace(@src());
     defer tracy_zone.end();
@@ -2223,29 +2022,6 @@ fn formattingHandler(server: *Server, writer: anytype, id: types.RequestId, req:
     return try respondGeneric(writer, id, null_result_response);
 }
 
-fn renameHandler(server: *Server, writer: anytype, id: types.RequestId, req: requests.Rename) !void {
-    const tracy_zone = tracy.trace(@src());
-    defer tracy_zone.end();
-
-    const handle = server.document_store.getHandle(req.params.textDocument.uri) orelse {
-        log.warn("Trying to rename in non existent document {s}", .{req.params.textDocument.uri});
-        return try respondGeneric(writer, id, null_result_response);
-    };
-
-    if (req.params.position.character >= 0) {
-        const source_index = offsets.positionToIndex(handle.document.text, req.params.position, server.offset_encoding);
-        const pos_context = try analysis.getPositionContext(server.arena.allocator(), handle.document, source_index);
-        switch (pos_context) {
-            .var_access => try server.renameDefinitionGlobal(writer, id, handle, source_index, req.params.newName),
-            .field_access => |loc| try server.renameDefinitionFieldAccess(writer, id, handle, source_index, loc, req.params.newName),
-            .label => try server.renameDefinitionLabel(writer, id, handle, source_index, req.params.newName),
-            else => try respondGeneric(writer, id, null_result_response),
-        }
-    } else {
-        try respondGeneric(writer, id, null_result_response);
-    }
-}
-
 fn didChangeConfigurationHandler(server: *Server, writer: anytype, id: types.RequestId, maybe_req: std.json.Value) !void {
     const tracy_zone = tracy.trace(@src());
     defer tracy_zone.end();
@@ -2265,53 +2041,129 @@ fn didChangeConfigurationHandler(server: *Server, writer: anytype, id: types.Req
         try server.requestConfiguration(writer);
 }
 
+fn renameHandler(server: *Server, writer: anytype, id: types.RequestId, req: requests.Rename) !void {
+    const tracy_zone = tracy.trace(@src());
+    defer tracy_zone.end();
+
+    try generalReferencesHandler(server, writer, id, .{ .rename = req });
+}
+
 fn referencesHandler(server: *Server, writer: anytype, id: types.RequestId, req: requests.References) !void {
     const tracy_zone = tracy.trace(@src());
     defer tracy_zone.end();
 
-    const handle = server.document_store.getHandle(req.params.textDocument.uri) orelse {
+    try generalReferencesHandler(server, writer, id, .{ .references = req });
-        log.warn("Trying to get references in non existent document {s}", .{req.params.textDocument.uri});
-        return try respondGeneric(writer, id, null_result_response);
-    };
-
-    if (req.params.position.character >= 0) {
-        const source_index = offsets.positionToIndex(handle.document.text, req.params.position, server.offset_encoding);
-        const pos_context = try analysis.getPositionContext(server.arena.allocator(), handle.document, source_index);
-
-        const include_decl = req.params.context.includeDeclaration;
-        switch (pos_context) {
-            .var_access => try server.referencesDefinitionGlobal(writer, id, handle, source_index, include_decl, false),
-            .field_access => |loc| try server.referencesDefinitionFieldAccess(writer, id, handle, source_index, loc, include_decl, false),
-            .label => try server.referencesDefinitionLabel(writer, id, handle, source_index, include_decl, false),
-            else => try respondGeneric(writer, id, null_result_response),
-        }
-    } else {
-        try respondGeneric(writer, id, null_result_response);
-    }
 }
 
 fn documentHighlightHandler(server: *Server, writer: anytype, id: types.RequestId, req: requests.DocumentHighlight) !void {
     const tracy_zone = tracy.trace(@src());
     defer tracy_zone.end();
 
-    const handle = server.document_store.getHandle(req.params.textDocument.uri) orelse {
+    try generalReferencesHandler(server, writer, id, .{ .highlight = req });
-        log.warn("Trying to highlight references in non existent document {s}", .{req.params.textDocument.uri});
+}
+
+const GeneralReferencesRequest = union(enum) {
+    rename: requests.Rename,
+    references: requests.References,
+    highlight: requests.DocumentHighlight,
+
+    pub fn uri(self: @This()) []const u8 {
+        return switch (self) {
+            .rename => |rename| rename.params.textDocument.uri,
+            .references => |ref| ref.params.textDocument.uri,
+            .highlight => |highlight| highlight.params.textDocument.uri,
+        };
+    }
+
+    pub fn position(self: @This()) types.Position {
+        return switch (self) {
+            .rename => |rename| rename.params.position,
+            .references => |ref| ref.params.position,
+            .highlight => |highlight| highlight.params.position,
+        };
+    }
+
+    pub fn name(self: @This()) []const u8 {
+        return switch (self) {
+            .rename => "rename",
+            .references => "references",
+            .highlight => "highlight references",
+        };
+    }
+};
+
+fn generalReferencesHandler(server: *Server, writer: anytype, id: types.RequestId, req: GeneralReferencesRequest) !void {
+    const tracy_zone = tracy.trace(@src());
+    defer tracy_zone.end();
+
+    const allocator = server.arena.allocator();
+
+    const handle = server.document_store.getHandle(req.uri()) orelse {
+        log.warn("Trying to get {s} in non existent document {s}", .{ req.name(), req.uri() });
         return try respondGeneric(writer, id, null_result_response);
     };
 
-    if (req.params.position.character >= 0) {
+    if (req.position().character <= 0) return try respondGeneric(writer, id, null_result_response);
-        const source_index = offsets.positionToIndex(handle.document.text, req.params.position, server.offset_encoding);
-        const pos_context = try analysis.getPositionContext(server.arena.allocator(), handle.document, source_index);
 
-        switch (pos_context) {
+    const source_index = offsets.positionToIndex(handle.document.text, req.position(), server.offset_encoding);
-            .var_access => try server.referencesDefinitionGlobal(writer, id, handle, source_index, true, true),
+    const pos_context = try analysis.getPositionContext(server.arena.allocator(), handle.document, source_index);
-            .field_access => |loc| try server.referencesDefinitionFieldAccess(writer, id, handle, source_index, loc, true, true),
+
-            .label => try server.referencesDefinitionLabel(writer, id, handle, source_index, true, true),
+    const decl = switch (pos_context) {
-            else => try respondGeneric(writer, id, null_result_response),
+        .var_access => try server.getSymbolGlobal(source_index, handle),
-        }
+        .field_access => |range| try server.getSymbolFieldAccess(handle, source_index, range),
-    } else {
+        .label => try getLabelGlobal(source_index, handle),
-        try respondGeneric(writer, id, null_result_response);
+        else => null,
-    }
+    } orelse return try respondGeneric(writer, id, null_result_response);
+
+    const include_decl = switch (req) {
+        .references => |ref| ref.params.context.includeDeclaration,
+        else => true,
+    };
+
+    const locations = if (pos_context == .label)
+        try references.labelReferences(allocator, decl, server.offset_encoding, include_decl)
+    else
+        try references.symbolReferences(
+            &server.arena,
+            &server.document_store,
+            decl,
+            server.offset_encoding,
+            include_decl,
+            server.config.skip_std_references,
+            req != .highlight, // scan the entire workspace except for highlight
+        );
+
+    const result: types.ResponseParams = switch (req) {
+        .rename => |rename| blk: {
+            var edits: types.WorkspaceEdit = .{ .changes = .{} };
+            for (locations.items) |loc| {
+                const gop = try edits.changes.getOrPutValue(allocator, loc.uri, .{});
+                try gop.value_ptr.append(allocator, .{
+                    .range = loc.range,
+                    .newText = rename.params.newName,
+                });
+            }
+            break :blk .{ .WorkspaceEdit = edits };
+        },
+        .references => .{ .Locations = locations.items },
+        .highlight => blk: {
+            var highlights = try std.ArrayListUnmanaged(types.DocumentHighlight).initCapacity(allocator, locations.items.len);
+            const uri = handle.uri();
+            for (locations.items) |loc| {
+                if (!std.mem.eql(u8, loc.uri, uri)) continue;
+                highlights.appendAssumeCapacity(.{
+                    .range = loc.range,
+                    .kind = .Text,
+                });
+            }
+            break :blk .{ .DocumentHighlight = highlights.items };
+        },
+    };
+
+    try send(writer, allocator, types.Response{
+        .id = id,
+        .result = result,
+    });
 }
 
 fn isPositionBefore(lhs: types.Position, rhs: types.Position) bool {

src/references.zig

@@ -7,16 +7,12 @@ const offsets = @import("offsets.zig");
 const log = std.log.scoped(.references);
 const ast = @import("ast.zig");
 
-fn tokenReference(handle: *DocumentStore.Handle, tok: Ast.TokenIndex, encoding: offsets.Encoding, context: anytype, comptime handler: anytype) !void {
+pub fn labelReferences(
-    try handler(context, types.Location{
+    allocator: std.mem.Allocator,
-        .uri = handle.uri(),
+    decl: analysis.DeclWithHandle,
-        .range = offsets.tokenToRange(handle.tree, tok, encoding),
+    encoding: offsets.Encoding,
-    });
+    include_decl: bool,
-}
+) error{OutOfMemory}!std.ArrayListUnmanaged(types.Location) {
-
-pub fn labelReferences(arena: *std.heap.ArenaAllocator, decl: analysis.DeclWithHandle, encoding: offsets.Encoding, include_decl: bool, context: anytype, comptime handler: anytype) !void {
-    _ = arena;
-
     std.debug.assert(decl.decl.* == .label_decl);
     const handle = decl.handle;
     const tree = handle.tree;
@@ -25,43 +21,88 @@ pub fn labelReferences(arena: *std.heap.ArenaAllocator, decl: analysis.DeclWithH
     // Find while / for / block from label -> iterate over children nodes, find break and continues, change their labels if they match.
     // This case can be implemented just by scanning tokens.
     const first_tok = tree.firstToken(decl.decl.label_decl);
-    const last_tok = tree.firstToken(decl.decl.label_decl);
+    const last_tok = ast.lastToken(tree, decl.decl.label_decl);
+
+    var locations = std.ArrayListUnmanaged(types.Location){};
+    errdefer locations.deinit(allocator);
 
     if (include_decl) {
         // The first token is always going to be the label
-        try tokenReference(handle, first_tok, encoding, context, handler);
+        try locations.append(allocator, .{
+            .uri = handle.uri(),
+            .range = offsets.tokenToRange(handle.tree, first_tok, encoding),
+        });
     }
 
     var curr_tok = first_tok + 1;
     while (curr_tok < last_tok - 2) : (curr_tok += 1) {
         const curr_id = token_tags[curr_tok];
-        if ((curr_id == .keyword_break or curr_id == .keyword_continue) and token_tags[curr_tok + 1] == .colon and
+
-            token_tags[curr_tok + 2] == .identifier)
+        if (curr_id != .keyword_break and curr_id != .keyword_continue) continue;
-        {
+        if (token_tags[curr_tok + 1] != .colon) continue;
-            if (std.mem.eql(u8, tree.tokenSlice(curr_tok + 2), tree.tokenSlice(first_tok))) {
+        if (token_tags[curr_tok + 2] != .identifier) continue;
-                try tokenReference(handle, first_tok, encoding, context, handler);
+
-            }
+        if (!std.mem.eql(u8, tree.tokenSlice(curr_tok + 2), tree.tokenSlice(first_tok))) continue;
-        }
+
+        try locations.append(allocator, .{
+            .uri = handle.uri(),
+            .range = offsets.tokenToRange(handle.tree, curr_tok + 2, encoding),
+        });
     }
+
+    return locations;
 }
 
-fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentStore, node_handle: analysis.NodeWithHandle, decl: analysis.DeclWithHandle, encoding: offsets.Encoding, context: anytype, comptime handler: anytype) error{OutOfMemory}!void {
+const Builder = struct {
-    const node = node_handle.node;
+    arena: *std.heap.ArenaAllocator,
-    const handle = node_handle.handle;
+    locations: std.ArrayListUnmanaged(types.Location),
+    store: *DocumentStore,
+    decl: analysis.DeclWithHandle,
+    encoding: offsets.Encoding,
+
+    pub fn init(arena: *std.heap.ArenaAllocator, store: *DocumentStore, decl: analysis.DeclWithHandle, encoding: offsets.Encoding) Builder {
+        return Builder{
+            .arena = arena,
+            .locations = .{},
+            .store = store,
+            .decl = decl,
+            .encoding = encoding,
+        };
+    }
+
+    pub fn add(self: *Builder, handle: *DocumentStore.Handle, token_index: Ast.TokenIndex) !void {
+        try self.locations.append(self.arena.allocator(), .{
+            .uri = handle.uri(),
+            .range = offsets.tokenToRange(handle.tree, token_index, self.encoding),
+        });
+    }
+};
+
+fn symbolReferencesInternal(
+    builder: *Builder,
+    node: Ast.Node.Index,
+    handle: *DocumentStore.Handle,
+) error{OutOfMemory}!void {
     const tree = handle.tree;
-    if (node > tree.nodes.len) return;
+
+    if (node == 0 or node > tree.nodes.len) return;
+
     const node_tags = tree.nodes.items(.tag);
     const datas = tree.nodes.items(.data);
     const main_tokens = tree.nodes.items(.main_token);
     const starts = tree.tokens.items(.start);
 
     switch (node_tags[node]) {
-        .block, .block_semicolon, .block_two, .block_two_semicolon => {
+        .block,
+        .block_semicolon,
+        .block_two,
+        .block_two_semicolon,
+        => {
             var buffer: [2]Ast.Node.Index = undefined;
             const statements = ast.blockStatements(tree, node, &buffer).?;
 
             for (statements) |stmt|
-                try symbolReferencesInternal(arena, store, .{ .node = stmt, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, stmt, handle);
         },
         .container_decl,
         .container_decl_trailing,
@@ -80,7 +121,7 @@ fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentSto
         => {
             var buf: [2]Ast.Node.Index = undefined;
             for (ast.declMembers(tree, node, &buf)) |member|
-                try symbolReferencesInternal(arena, store, .{ .node = member, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, member, handle);
         },
         .global_var_decl,
         .local_var_decl,
@@ -88,34 +129,20 @@ fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentSto
         .aligned_var_decl,
         => {
             const var_decl = ast.varDecl(tree, node).?;
-            if (var_decl.ast.type_node != 0) {
+            try symbolReferencesInternal(builder, var_decl.ast.type_node, handle);
-                try symbolReferencesInternal(arena, store, .{ .node = var_decl.ast.type_node, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, var_decl.ast.init_node, handle);
-            }
-            if (var_decl.ast.init_node != 0) {
-                try symbolReferencesInternal(arena, store, .{ .node = var_decl.ast.init_node, .handle = handle }, decl, encoding, context, handler);
-            }
-        },
-        .@"usingnamespace" => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
         },
         .container_field,
         .container_field_align,
         .container_field_init,
         => {
             const field = ast.containerField(tree, node).?;
-            if (field.ast.type_expr != 0) {
+            try symbolReferencesInternal(builder, field.ast.type_expr, handle);
-                try symbolReferencesInternal(arena, store, .{ .node = field.ast.type_expr, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, field.ast.value_expr, handle);
-            }
-            if (field.ast.value_expr != 0) {
-                try symbolReferencesInternal(arena, store, .{ .node = field.ast.value_expr, .handle = handle }, decl, encoding, context, handler);
-            }
         },
         .identifier => {
-            if (try analysis.lookupSymbolGlobal(store, arena, handle, offsets.nodeToSlice(tree, node), starts[main_tokens[node]])) |child| {
+            const child = (try analysis.lookupSymbolGlobal(builder.store, builder.arena, handle, tree.getNodeSource(node), starts[main_tokens[node]])) orelse return;
-                if (std.meta.eql(decl, child)) {
+            if (std.meta.eql(builder.decl, child)) try builder.add(handle, main_tokens[node]);
-                    try tokenReference(handle, main_tokens[node], encoding, context, handler);
-                }
-            }
         },
         .fn_proto,
         .fn_proto_multi,
@@ -127,62 +154,39 @@ fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentSto
             const fn_proto = ast.fnProto(tree, node, &buf).?;
             var it = fn_proto.iterate(&tree);
             while (ast.nextFnParam(&it)) |param| {
-                if (param.type_expr != 0)
+                try symbolReferencesInternal(builder, param.type_expr, handle);
-                    try symbolReferencesInternal(arena, store, .{ .node = param.type_expr, .handle = handle }, decl, encoding, context, handler);
             }
 
-            if (fn_proto.ast.return_type != 0) {
+            try symbolReferencesInternal(builder, fn_proto.ast.return_type, handle);
-                try symbolReferencesInternal(arena, store, .{ .node = fn_proto.ast.return_type, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, fn_proto.ast.align_expr, handle);
-            }
+            try symbolReferencesInternal(builder, fn_proto.ast.section_expr, handle);
-            if (fn_proto.ast.align_expr != 0) {
+            try symbolReferencesInternal(builder, fn_proto.ast.callconv_expr, handle);
-                try symbolReferencesInternal(arena, store, .{ .node = fn_proto.ast.align_expr, .handle = handle }, decl, encoding, context, handler);
-            }
-            if (fn_proto.ast.section_expr != 0) {
-                try symbolReferencesInternal(arena, store, .{ .node = fn_proto.ast.section_expr, .handle = handle }, decl, encoding, context, handler);
-            }
-            if (fn_proto.ast.callconv_expr != 0) {
-                try symbolReferencesInternal(arena, store, .{ .node = fn_proto.ast.callconv_expr, .handle = handle }, decl, encoding, context, handler);
-            }
             if (node_tags[node] == .fn_decl) {
-                try symbolReferencesInternal(arena, store, .{ .node = datas[node].rhs, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, datas[node].rhs, handle);
             }
         },
-        .anyframe_type => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].rhs, .handle = handle }, decl, encoding, context, handler);
-        },
-        .@"defer" => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].rhs, .handle = handle }, decl, encoding, context, handler);
-        },
-        .@"comptime" => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
-        },
-        .@"nosuspend" => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
-        },
         .@"switch",
         .switch_comma,
         => {
             // TODO When renaming a union(enum) field, also rename switch items that refer to it.
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, datas[node].lhs, handle);
             const extra = tree.extraData(datas[node].rhs, Ast.Node.SubRange);
             const cases = tree.extra_data[extra.start..extra.end];
             for (cases) |case| {
-                try symbolReferencesInternal(arena, store, .{ .node = case, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, case, handle);
             }
         },
         .switch_case_one => {
             const case_one = tree.switchCaseOne(node);
-            if (case_one.ast.target_expr != 0)
+            try symbolReferencesInternal(builder, case_one.ast.target_expr, handle);
-                try symbolReferencesInternal(arena, store, .{ .node = case_one.ast.target_expr, .handle = handle }, decl, encoding, context, handler);
             for (case_one.ast.values) |val|
-                try symbolReferencesInternal(arena, store, .{ .node = val, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, val, handle);
         },
         .switch_case => {
             const case = tree.switchCase(node);
-            if (case.ast.target_expr != 0)
+            try symbolReferencesInternal(builder, case.ast.target_expr, handle);
-                try symbolReferencesInternal(arena, store, .{ .node = case.ast.target_expr, .handle = handle }, decl, encoding, context, handler);
             for (case.ast.values) |val|
-                try symbolReferencesInternal(arena, store, .{ .node = val, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, val, handle);
         },
         .@"while",
         .while_simple,
@@ -191,31 +195,17 @@ fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentSto
         .@"for",
         => {
             const loop = ast.whileAst(tree, node).?;
-            try symbolReferencesInternal(arena, store, .{ .node = loop.ast.cond_expr, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, loop.ast.cond_expr, handle);
-            if (loop.ast.cont_expr != 0) {
+            try symbolReferencesInternal(builder, loop.ast.then_expr, handle);
-                try symbolReferencesInternal(arena, store, .{ .node = loop.ast.cont_expr, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, loop.ast.else_expr, handle);
-            }
-            try symbolReferencesInternal(arena, store, .{ .node = loop.ast.then_expr, .handle = handle }, decl, encoding, context, handler);
-            if (loop.ast.else_expr != 0) {
-                try symbolReferencesInternal(arena, store, .{ .node = loop.ast.else_expr, .handle = handle }, decl, encoding, context, handler);
-            }
         },
         .@"if",
         .if_simple,
         => {
             const if_node = ast.ifFull(tree, node);
-
+            try symbolReferencesInternal(builder, if_node.ast.cond_expr, handle);
-            try symbolReferencesInternal(arena, store, .{ .node = if_node.ast.cond_expr, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, if_node.ast.then_expr, handle);
-            try symbolReferencesInternal(arena, store, .{ .node = if_node.ast.then_expr, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, if_node.ast.else_expr, handle);
-            if (if_node.ast.else_expr != 0) {
-                try symbolReferencesInternal(arena, store, .{ .node = if_node.ast.else_expr, .handle = handle }, decl, encoding, context, handler);
-            }
-        },
-        .array_type,
-        .array_type_sentinel,
-        => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].rhs, .handle = handle }, decl, encoding, context, handler);
         },
         .ptr_type,
         .ptr_type_aligned,
@@ -225,26 +215,15 @@ fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentSto
             const ptr_type = ast.ptrType(tree, node).?;
 
             if (ptr_type.ast.align_node != 0) {
-                try symbolReferencesInternal(arena, store, .{ .node = ptr_type.ast.align_node, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, ptr_type.ast.align_node, handle);
                 if (node_tags[node] == .ptr_type_bit_range) {
-                    try symbolReferencesInternal(arena, store, .{
+                    try symbolReferencesInternal(builder, ptr_type.ast.bit_range_start, handle);
-                        .node = ptr_type.ast.bit_range_start,
+                    try symbolReferencesInternal(builder, ptr_type.ast.bit_range_end, handle);
-                        .handle = handle,
-                    }, decl, encoding, context, handler);
-                    try symbolReferencesInternal(arena, store, .{
-                        .node = ptr_type.ast.bit_range_end,
-                        .handle = handle,
-                    }, decl, encoding, context, handler);
                 }
             }
-            if (ptr_type.ast.sentinel != 0) {
-                try symbolReferencesInternal(arena, store, .{ .node = ptr_type.ast.sentinel, .handle = handle }, decl, encoding, context, handler);
-            }
 
-            try symbolReferencesInternal(arena, store, .{ .node = ptr_type.ast.child_type, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, ptr_type.ast.sentinel, handle);
-        },
+            try symbolReferencesInternal(builder, ptr_type.ast.child_type, handle);
-        .address_of, .@"await", .bit_not, .bool_not, .optional_type, .negation, .negation_wrap, .@"resume", .@"try" => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
         },
         .array_init,
         .array_init_comma,
@@ -254,19 +233,18 @@ fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentSto
         .array_init_one_comma,
         .array_init_dot_two,
         .array_init_dot_two_comma,
-        => |n| {
+        => |tag| {
             var buf: [2]Ast.Node.Index = undefined;
-            const array_init = switch (n) {
+            const array_init = switch (tag) {
                 .array_init, .array_init_comma => tree.arrayInit(node),
                 .array_init_dot, .array_init_dot_comma => tree.arrayInitDot(node),
                 .array_init_one, .array_init_one_comma => tree.arrayInitOne(buf[0..1], node),
                 .array_init_dot_two, .array_init_dot_two_comma => tree.arrayInitDotTwo(&buf, node),
                 else => unreachable,
             };
-            if (array_init.ast.type_expr != 0)
+            try symbolReferencesInternal(builder, array_init.ast.type_expr, handle);
-                try symbolReferencesInternal(arena, store, .{ .node = array_init.ast.type_expr, .handle = handle }, decl, encoding, context, handler);
             for (array_init.ast.elements) |e|
-                try symbolReferencesInternal(arena, store, .{ .node = e, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, e, handle);
         },
         .struct_init,
         .struct_init_comma,
@@ -276,19 +254,18 @@ fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentSto
         .struct_init_dot_two_comma,
         .struct_init_one,
         .struct_init_one_comma,
-        => |n| {
+        => |tag| {
             var buf: [2]Ast.Node.Index = undefined;
-            const struct_init: Ast.full.StructInit = switch (n) {
+            const struct_init: Ast.full.StructInit = switch (tag) {
                 .struct_init, .struct_init_comma => tree.structInit(node),
                 .struct_init_dot, .struct_init_dot_comma => tree.structInitDot(node),
                 .struct_init_one, .struct_init_one_comma => tree.structInitOne(buf[0..1], node),
                 .struct_init_dot_two, .struct_init_dot_two_comma => tree.structInitDotTwo(&buf, node),
                 else => unreachable,
             };
-            if (struct_init.ast.type_expr != 0)
+            try symbolReferencesInternal(builder, struct_init.ast.type_expr, handle);
-                try symbolReferencesInternal(arena, store, .{ .node = struct_init.ast.type_expr, .handle = handle }, decl, encoding, context, handler);
             for (struct_init.ast.fields) |field|
-                try symbolReferencesInternal(arena, store, .{ .node = field, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, field, handle);
         },
         .call,
         .call_comma,
@@ -302,55 +279,27 @@ fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentSto
             var buf: [1]Ast.Node.Index = undefined;
             const call = ast.callFull(tree, node, &buf).?;
 
-            if (call.ast.fn_expr != 0)
+            try symbolReferencesInternal(builder, call.ast.fn_expr, handle);
-                try symbolReferencesInternal(arena, store, .{ .node = call.ast.fn_expr, .handle = handle }, decl, encoding, context, handler);
 
             for (call.ast.params) |param| {
-                try symbolReferencesInternal(arena, store, .{ .node = param, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, param, handle);
             }
         },
         .slice,
         .slice_sentinel,
         .slice_open,
-        => |s| {
+        => |tag| {
-            const slice: Ast.full.Slice = switch (s) {
+            const slice: Ast.full.Slice = switch (tag) {
                 .slice => tree.slice(node),
                 .slice_open => tree.sliceOpen(node),
                 .slice_sentinel => tree.sliceSentinel(node),
                 else => unreachable,
             };
 
-            try symbolReferencesInternal(arena, store, .{ .node = slice.ast.sliced, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, slice.ast.sliced, handle);
-            try symbolReferencesInternal(arena, store, .{ .node = slice.ast.start, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, slice.ast.start, handle);
-            if (slice.ast.end != 0)
+            try symbolReferencesInternal(builder, slice.ast.end, handle);
-                try symbolReferencesInternal(arena, store, .{ .node = slice.ast.end, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, slice.ast.sentinel, handle);
-            if (slice.ast.sentinel != 0)
-                try symbolReferencesInternal(arena, store, .{ .node = slice.ast.sentinel, .handle = handle }, decl, encoding, context, handler);
-        },
-        .array_access => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].rhs, .handle = handle }, decl, encoding, context, handler);
-        },
-        .deref,
-        .unwrap_optional,
-        => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
-        },
-        .grouped_expression => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
-        },
-        .@"return",
-        .@"break",
-        .@"continue",
-        => {
-            if (datas[node].lhs != 0) {
-                try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
-            }
-        },
-        .@"suspend" => {
-            if (datas[node].lhs != 0) {
-                try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
-            }
         },
         .builtin_call,
         .builtin_call_comma,
@@ -361,33 +310,32 @@ fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentSto
             const params = ast.builtinCallParams(tree, node, &buffer).?;
 
             for (params) |param|
-                try symbolReferencesInternal(arena, store, .{ .node = param, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, param, handle);
         },
         .@"asm",
         .asm_simple,
-        => |a| {
+        => |tag| {
-            const _asm: Ast.full.Asm = if (a == .@"asm") tree.asmFull(node) else tree.asmSimple(node);
+            const full_asm: Ast.full.Asm = if (tag == .@"asm") tree.asmFull(node) else tree.asmSimple(node);
-            if (_asm.ast.items.len == 0)
+            if (full_asm.ast.items.len == 0)
-                try symbolReferencesInternal(arena, store, .{ .node = _asm.ast.template, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, full_asm.ast.template, handle);
 
-            for (_asm.inputs) |input|
+            for (full_asm.inputs) |input|
-                try symbolReferencesInternal(arena, store, .{ .node = input, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, input, handle);
 
-            for (_asm.outputs) |output|
+            for (full_asm.outputs) |output|
-                try symbolReferencesInternal(arena, store, .{ .node = output, .handle = handle }, decl, encoding, context, handler);
+                try symbolReferencesInternal(builder, output, handle);
-        },
-        .test_decl => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].rhs, .handle = handle }, decl, encoding, context, handler);
         },
+        .asm_output => unreachable,
+        .asm_input => unreachable,
         .field_access => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, datas[node].lhs, handle);
 
             const rhs_str = ast.tokenSlice(tree, datas[node].rhs) catch return;
             var bound_type_params = analysis.BoundTypeParams{};
             const left_type = try analysis.resolveFieldAccessLhsType(
-                store,
+                builder.store,
-                arena,
+                builder.arena,
-                (try analysis.resolveTypeOfNodeInternal(store, arena, .{
+                (try analysis.resolveTypeOfNodeInternal(builder.store, builder.arena, .{
                     .node = datas[node].lhs,
                     .handle = handle,
                 }, &bound_type_params)) orelse return,
@@ -399,152 +347,198 @@ fn symbolReferencesInternal(arena: *std.heap.ArenaAllocator, store: *DocumentSto
                 else => return,
             };
 
-            if (try analysis.lookupSymbolContainer(
+            const child = (try analysis.lookupSymbolContainer(
-                store,
+                builder.store,
-                arena,
+                builder.arena,
                 .{ .node = left_type_node, .handle = left_type.handle },
                 rhs_str,
                 !left_type.type.is_type_val,
-            )) |child| {
+            )) orelse return;
-                if (std.meta.eql(child, decl)) {
+
-                    try tokenReference(handle, datas[node].rhs, encoding, context, handler);
+            if (std.meta.eql(child, builder.decl)) try builder.add(handle, datas[node].rhs);
-                }
-            }
         },
-        .add,
+        .@"usingnamespace",
-        .add_wrap,
+        .unwrap_optional,
-        .array_cat,
+        .bool_not,
-        .array_mult,
+        .negation,
-        .assign,
+        .bit_not,
-        .assign_bit_and,
+        .negation_wrap,
-        .assign_bit_or,
+        .address_of,
-        .assign_shl,
+        .@"try",
-        .assign_shr,
+        .@"await",
-        .assign_bit_xor,
+        .optional_type,
-        .assign_div,
+        .deref,
-        .assign_sub,
+        .@"suspend",
-        .assign_sub_wrap,
+        .@"resume",
-        .assign_mod,
+        .@"continue",
-        .assign_add,
+        .@"break",
-        .assign_add_wrap,
+        .@"return",
-        .assign_mul,
+        .grouped_expression,
-        .assign_mul_wrap,
+        .@"comptime",
-        .bang_equal,
+        .@"nosuspend",
-        .bit_and,
+        => try symbolReferencesInternal(builder, datas[node].lhs, handle),
-        .bit_or,
+        .test_decl,
-        .shl,
+        .@"errdefer",
-        .shr,
+        .@"defer",
-        .bit_xor,
+        .anyframe_type,
-        .bool_or,
+        => try symbolReferencesInternal(builder, datas[node].rhs, handle),
-        .div,
         .equal_equal,
-        .error_union,
+        .bang_equal,
-        .greater_or_equal,
+        .less_than,
         .greater_than,
         .less_or_equal,
-        .less_than,
+        .greater_or_equal,
+        .assign_mul,
+        .assign_div,
+        .assign_mod,
+        .assign_add,
+        .assign_sub,
+        .assign_shl,
+        .assign_shl_sat,
+        .assign_shr,
+        .assign_bit_and,
+        .assign_bit_xor,
+        .assign_bit_or,
+        .assign_mul_wrap,
+        .assign_add_wrap,
+        .assign_sub_wrap,
+        .assign_mul_sat,
+        .assign_add_sat,
+        .assign_sub_sat,
+        .assign,
         .merge_error_sets,
-        .mod,
         .mul,
+        .div,
+        .mod,
+        .array_mult,
         .mul_wrap,
-        .switch_range,
+        .mul_sat,
+        .add,
         .sub,
+        .array_cat,
+        .add_wrap,
         .sub_wrap,
+        .add_sat,
+        .sub_sat,
+        .shl,
+        .shl_sat,
+        .shr,
+        .bit_and,
+        .bit_xor,
+        .bit_or,
         .@"orelse",
+        .bool_and,
+        .bool_or,
+        .array_type,
+        .array_type_sentinel,
+        .array_access,
+        .@"catch",
+        .switch_range,
+        .error_union,
         => {
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].lhs, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, datas[node].lhs, handle);
-            try symbolReferencesInternal(arena, store, .{ .node = datas[node].rhs, .handle = handle }, decl, encoding, context, handler);
+            try symbolReferencesInternal(builder, datas[node].rhs, handle);
         },
-        else => {},
+        .anyframe_literal,
+        .char_literal,
+        .number_literal,
+        .unreachable_literal,
+        .enum_literal,
+        .string_literal,
+        .multiline_string_literal,
+        .error_value,
+        => {},
     }
 }
 
-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 {
+pub fn symbolReferences(
+    arena: *std.heap.ArenaAllocator,
+    store: *DocumentStore,
+    decl_handle: analysis.DeclWithHandle,
+    encoding: offsets.Encoding,
+    include_decl: bool,
+    skip_std_references: bool,
+    workspace: bool,
+) error{OutOfMemory}!std.ArrayListUnmanaged(types.Location) {
     std.debug.assert(decl_handle.decl.* != .label_decl);
+
+    var builder = Builder.init(arena, store, decl_handle, encoding);
+
     const curr_handle = decl_handle.handle;
-    if (include_decl) {
+    if (include_decl) try builder.add(curr_handle, decl_handle.nameToken());
-        try tokenReference(curr_handle, decl_handle.nameToken(), encoding, context, handler);
-    }
 
     switch (decl_handle.decl.*) {
-        .ast_node => {
+        .pointer_payload,
-            try symbolReferencesInternal(arena, store, .{ .node = 0, .handle = curr_handle }, decl_handle, encoding, context, handler);
+        .switch_payload,
+        .array_payload,
+        .array_index,
+        .ast_node,
+        => {
+            try symbolReferencesInternal(&builder, 0, curr_handle);
 
-            if (workspace) {
+            if (decl_handle.decl.* != .ast_node) return builder.locations;
-                var imports = std.ArrayListUnmanaged(*DocumentStore.Handle){};
+            if (!workspace) return builder.locations;
 
-                var handle_it = store.handles.iterator();
+            var imports = std.ArrayListUnmanaged(*DocumentStore.Handle){};
-                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
+            var handle_it = store.handles.iterator();
-                    try imports.append(arena.allocator(), entry.value_ptr.*);
+            while (handle_it.next()) |entry| {
-                    var i: usize = 0;
+                if (skip_std_references and std.mem.indexOf(u8, entry.key_ptr.*, "std") != null) {
-                    blk: while (i < imports.items.len) : (i += 1) {
+                    if (!include_decl or entry.value_ptr.* != curr_handle)
-                        const import = imports.items[i];
+                        continue;
-                        for (import.imports_used.items) |uri| {
+                }
-                            const h = store.getHandle(uri) orelse break;
 
-                            if (h == curr_handle) {
+                // Check entry's transitive imports
-                                // entry does import curr_handle
+                try imports.append(arena.allocator(), entry.value_ptr.*);
-                                try symbolReferencesInternal(arena, store, .{ .node = 0, .handle = entry.value_ptr.* }, decl_handle, encoding, context, handler);
+                var i: usize = 0;
-                                break :blk;
+                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;
 
-                            select: {
+                        if (h == curr_handle) {
-                                for (imports.items) |item| {
+                            // entry does import curr_handle
-                                    if (item == h) {
+                            try symbolReferencesInternal(&builder, 0, entry.value_ptr.*);
-                                        // already checked this import
+                            break :blk;
-                                        break :select;
+                        }
-                                    }
+
+                        select: {
+                            for (imports.items) |item| {
+                                if (item == h) {
+                                    // already checked this import
+                                    break :select;
                                 }
-                                try imports.append(arena.allocator(), h);
                             }
+                            try imports.append(arena.allocator(), h);
                         }
                     }
-                    try imports.resize(arena.allocator(), 0);
                 }
+                try imports.resize(arena.allocator(), 0);
             }
         },
         .param_decl => |param| {
             // Rename the param tok.
-            const fn_node: Ast.full.FnProto = loop: for (curr_handle.document_scope.scopes.items) |scope| {
+            for (curr_handle.document_scope.scopes.items) |scope| {
-                switch (scope.data) {
+                if (scope.data != .function) continue;
-                    .function => |proto| {
+
-                        var buf: [1]Ast.Node.Index = undefined;
+                const proto = scope.data.function;
-                        const fn_proto = ast.fnProto(curr_handle.tree, proto, &buf).?;
+
-                        var it = fn_proto.iterate(&curr_handle.tree);
+                var buf: [1]Ast.Node.Index = undefined;
-                        while (ast.nextFnParam(&it)) |candidate| {
+                const fn_proto = ast.fnProto(curr_handle.tree, proto, &buf).?;
-                            if (std.meta.eql(candidate, param)) {
+
-                                if (curr_handle.tree.nodes.items(.tag)[proto] == .fn_decl) {
+                var it = fn_proto.iterate(&curr_handle.tree);
-                                    try symbolReferencesInternal(
+                while (ast.nextFnParam(&it)) |candidate| {
-                                        arena,
+                    if (!std.meta.eql(candidate, param)) continue;
-                                        store,
+
-                                        .{ .node = curr_handle.tree.nodes.items(.data)[proto].rhs, .handle = curr_handle },
+                    if (curr_handle.tree.nodes.items(.tag)[proto] != .fn_decl) break;
-                                        decl_handle,
+                    try symbolReferencesInternal(&builder, curr_handle.tree.nodes.items(.data)[proto].rhs, curr_handle);
-                                        encoding,
+                    break;
-                                        context,
-                                        handler,
-                                    );
-                                }
-                                break :loop fn_proto;
-                            }
-                        }
-                    },
-                    else => {},
                 }
-            } else {
+            }
-                log.warn("Could not find param decl's function", .{});
+            log.warn("Could not find param decl's function", .{});
-                return;
-            };
-            _ = fn_node;
         },
-        .pointer_payload, .switch_payload, .array_payload, .array_index => {
+        .label_decl => unreachable, // handled separately by labelReferences
-            try symbolReferencesInternal(arena, store, .{ .node = 0, .handle = curr_handle }, decl_handle, encoding, context, handler);
-        },
-        .label_decl => unreachable,
     }
+
+    return builder.locations;
 }

src/rename.zig

@@ -1,52 +0,0 @@
-const std = @import("std");
-const DocumentStore = @import("DocumentStore.zig");
-const analysis = @import("analysis.zig");
-const references = @import("references.zig");
-const types = @import("types.zig");
-const offsets = @import("offsets.zig");
-
-// TODO Use a map to array lists and collect at the end instead?
-const RefHandlerContext = struct {
-    allocator: std.mem.Allocator,
-    edits: *std.StringHashMapUnmanaged(std.ArrayListUnmanaged(types.TextEdit)),
-    new_name: []const u8,
-};
-
-fn refHandler(context: RefHandlerContext, loc: types.Location) !void {
-    const gop = try context.edits.getOrPutValue(context.allocator, loc.uri, .{});
-    try gop.value_ptr.append(context.allocator, .{
-        .range = loc.range,
-        .newText = context.new_name,
-    });
-}
-
-pub fn renameSymbol(
-    arena: *std.heap.ArenaAllocator,
-    store: *DocumentStore,
-    decl_handle: analysis.DeclWithHandle,
-    new_name: []const u8,
-    edits: *std.StringHashMapUnmanaged(std.ArrayListUnmanaged(types.TextEdit)),
-    encoding: offsets.Encoding,
-) !void {
-    std.debug.assert(decl_handle.decl.* != .label_decl);
-    try references.symbolReferences(arena, store, decl_handle, encoding, true, RefHandlerContext{
-        .edits = edits,
-        .allocator = arena.allocator(),
-        .new_name = new_name,
-    }, refHandler, true, true);
-}
-
-pub fn renameLabel(
-    arena: *std.heap.ArenaAllocator,
-    decl_handle: analysis.DeclWithHandle,
-    new_name: []const u8,
-    edits: *std.StringHashMapUnmanaged(std.ArrayListUnmanaged(types.TextEdit)),
-    encoding: offsets.Encoding,
-) !void {
-    std.debug.assert(decl_handle.decl.* == .label_decl);
-    try references.labelReferences(arena, decl_handle, encoding, true, RefHandlerContext{
-        .edits = edits,
-        .allocator = arena.allocator(),
-        .new_name = new_name,
-    }, refHandler);
-}