simplify symbol references handler (#978)

src/ast.zig

@@ -1395,14 +1395,18 @@ pub fn iterateChildren(
             var buffer: [1]Node.Index = undefined;
             const fn_proto = tree.fullFnProto(&buffer, node).?;
 
-            for (fn_proto.ast.params) |child| {
+            var it = fn_proto.iterate(&tree);
-                try callback(context, child);
+            while (nextFnParam(&it)) |param| {
+                try callback(context, param.type_expr);
             }
             try callback(context, fn_proto.ast.align_expr);
             try callback(context, fn_proto.ast.addrspace_expr);
             try callback(context, fn_proto.ast.section_expr);
             try callback(context, fn_proto.ast.callconv_expr);
             try callback(context, fn_proto.ast.return_type);
+            if (node_tags[node] == .fn_decl) {
+                try callback(context, node_data[node].rhs);
+            }
         },
 
         .container_decl_arg,

src/references.zig

@@ -13,7 +13,7 @@ pub fn labelReferences(
     encoding: offsets.Encoding,
     include_decl: bool,
 ) error{OutOfMemory}!std.ArrayListUnmanaged(types.Location) {
-    std.debug.assert(decl.decl.* == .label_decl);
+    std.debug.assert(decl.decl.* == .label_decl); // use `symbolReferences` instead
     const handle = decl.handle;
     const tree = handle.tree;
     const token_tags = tree.tokens.items(.tag);
@@ -55,276 +55,68 @@ pub fn labelReferences(
 
 const Builder = struct {
     arena: *std.heap.ArenaAllocator,
-    locations: std.ArrayListUnmanaged(types.Location),
+    locations: std.ArrayListUnmanaged(types.Location) = .{},
+    /// this is the declaration we are searching for
+    decl_handle: analysis.DeclWithHandle,
     store: *DocumentStore,
-    decl: analysis.DeclWithHandle,
     encoding: offsets.Encoding,
 
-    pub fn init(arena: *std.heap.ArenaAllocator, store: *DocumentStore, decl: analysis.DeclWithHandle, encoding: offsets.Encoding) Builder {
+    const Context = struct {
-        return Builder{
+        builder: *Builder,
-            .arena = arena,
+        handle: *const DocumentStore.Handle,
-            .locations = .{},
-            .store = store,
-            .decl = decl,
-            .encoding = encoding,
     };
-    }
 
-    pub fn add(self: *Builder, handle: *const DocumentStore.Handle, token_index: Ast.TokenIndex) !void {
+    pub fn add(self: *Builder, handle: *const DocumentStore.Handle, token_index: Ast.TokenIndex) error{OutOfMemory}!void {
         try self.locations.append(self.arena.allocator(), .{
             .uri = handle.uri,
             .range = offsets.tokenToRange(handle.tree, token_index, self.encoding),
         });
     }
-};
 
-fn symbolReferencesInternal(
+    fn collectReferences(self: *Builder, handle: *const DocumentStore.Handle, node: Ast.Node.Index) error{OutOfMemory}!void {
-    builder: *Builder,
+        const context = Context{
-    node: Ast.Node.Index,
+            .builder = self,
-    handle: *const DocumentStore.Handle,
+            .handle = handle,
-    is_root: bool,
+        };
-) error{OutOfMemory}!void {
+        try ast.iterateChildrenRecursive(handle.tree, node, &context, error{OutOfMemory}, referenceNode);
-    const tree = handle.tree;
+    }
 
-    if (!is_root and node == 0 or node > tree.nodes.len) return;
+    fn referenceNode(self: *const Context, node: Ast.Node.Index) error{OutOfMemory}!void {
+        const builder = self.builder;
+        const handle = self.handle;
 
-    const node_tags = tree.nodes.items(.tag);
+        const node_tags = handle.tree.nodes.items(.tag);
-    const datas = tree.nodes.items(.data);
+        const datas = handle.tree.nodes.items(.data);
-    const main_tokens = tree.nodes.items(.main_token);
+        const main_tokens = handle.tree.nodes.items(.main_token);
-    const starts = tree.tokens.items(.start);
+        const starts = handle.tree.tokens.items(.start);
 
         switch (node_tags[node]) {
-        .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(builder, stmt, handle, false);
-        },
-        .container_decl,
-        .container_decl_trailing,
-        .container_decl_arg,
-        .container_decl_arg_trailing,
-        .container_decl_two,
-        .container_decl_two_trailing,
-        .tagged_union,
-        .tagged_union_trailing,
-        .tagged_union_two,
-        .tagged_union_two_trailing,
-        .tagged_union_enum_tag,
-        .tagged_union_enum_tag_trailing,
-        .root,
-        => {
-            var buf: [2]Ast.Node.Index = undefined;
-            const container_decl = tree.fullContainerDecl(&buf, node).?;
-            for (container_decl.ast.members) |member|
-                try symbolReferencesInternal(builder, member, handle, false);
-        },
-        .error_set_decl => {},
-        .global_var_decl,
-        .local_var_decl,
-        .simple_var_decl,
-        .aligned_var_decl,
-        => {
-            const var_decl = tree.fullVarDecl(node).?;
-            try symbolReferencesInternal(builder, var_decl.ast.type_node, handle, false);
-            try symbolReferencesInternal(builder, var_decl.ast.init_node, handle, false);
-        },
-        .container_field,
-        .container_field_align,
-        .container_field_init,
-        => {
-            const field = tree.fullContainerField(node).?;
-            try symbolReferencesInternal(builder, field.ast.type_expr, handle, false);
-            try symbolReferencesInternal(builder, field.ast.value_expr, handle, false);
-        },
             .identifier => {
-            const child = (try analysis.lookupSymbolGlobal(builder.store, builder.arena, handle, tree.getNodeSource(node), starts[main_tokens[node]])) orelse return;
+                const identifier_token = main_tokens[node];
-            if (std.meta.eql(builder.decl, child)) try builder.add(handle, main_tokens[node]);
-        },
-        .fn_proto,
-        .fn_proto_multi,
-        .fn_proto_one,
-        .fn_proto_simple,
-        .fn_decl,
-        => {
-            var buf: [1]Ast.Node.Index = undefined;
-            const fn_proto = tree.fullFnProto(&buf, node).?;
-            var it = fn_proto.iterate(&tree);
-            while (ast.nextFnParam(&it)) |param| {
-                try symbolReferencesInternal(builder, param.type_expr, handle, false);
-            }
 
-            try symbolReferencesInternal(builder, fn_proto.ast.return_type, handle, false);
+                const child = (try analysis.lookupSymbolGlobal(
-            try symbolReferencesInternal(builder, fn_proto.ast.align_expr, handle, false);
+                    builder.store,
-            try symbolReferencesInternal(builder, fn_proto.ast.section_expr, handle, false);
+                    builder.arena,
-            try symbolReferencesInternal(builder, fn_proto.ast.callconv_expr, handle, false);
+                    handle,
-            if (node_tags[node] == .fn_decl) {
+                    offsets.tokenToSlice(handle.tree, identifier_token),
-                try symbolReferencesInternal(builder, datas[node].rhs, handle, false);
+                    starts[identifier_token],
+                )) orelse return;
+
+                if (std.meta.eql(builder.decl_handle, child)) {
+                    try builder.add(handle, identifier_token);
                 }
             },
-        .@"switch",
-        .switch_comma,
-        => {
-            // TODO When renaming a union(enum) field, also rename switch items that refer to it.
-            try symbolReferencesInternal(builder, datas[node].lhs, handle, false);
-            const extra = tree.extraData(datas[node].rhs, Ast.Node.SubRange);
-            const cases = tree.extra_data[extra.start..extra.end];
-            for (cases) |case| {
-                try symbolReferencesInternal(builder, case, handle, false);
-            }
-        },
-        .switch_case_one,
-        .switch_case_inline_one,
-        .switch_case,
-        .switch_case_inline,
-        => {
-            const case = tree.fullSwitchCase(node).?;
-            try symbolReferencesInternal(builder, case.ast.target_expr, handle, false);
-            for (case.ast.values) |val|
-                try symbolReferencesInternal(builder, val, handle, false);
-        },
-        .@"while",
-        .while_simple,
-        .while_cont,
-        .for_simple,
-        .@"for",
-        => {
-            const loop = ast.fullWhile(tree, node).?;
-            try symbolReferencesInternal(builder, loop.ast.cond_expr, handle, false);
-            try symbolReferencesInternal(builder, loop.ast.then_expr, handle, false);
-            try symbolReferencesInternal(builder, loop.ast.cont_expr, handle, false);
-            try symbolReferencesInternal(builder, loop.ast.else_expr, handle, false);
-        },
-        .@"if",
-        .if_simple,
-        => {
-            const if_node = ast.fullIf(tree, node).?;
-            try symbolReferencesInternal(builder, if_node.ast.cond_expr, handle, false);
-            try symbolReferencesInternal(builder, if_node.ast.then_expr, handle, false);
-            try symbolReferencesInternal(builder, if_node.ast.else_expr, handle, false);
-        },
-        .ptr_type,
-        .ptr_type_aligned,
-        .ptr_type_bit_range,
-        .ptr_type_sentinel,
-        => {
-            const ptr_type = ast.fullPtrType(tree, node).?;
-
-            if (ptr_type.ast.align_node != 0) {
-                try symbolReferencesInternal(builder, ptr_type.ast.align_node, handle, false);
-                if (node_tags[node] == .ptr_type_bit_range) {
-                    try symbolReferencesInternal(builder, ptr_type.ast.bit_range_start, handle, false);
-                    try symbolReferencesInternal(builder, ptr_type.ast.bit_range_end, handle, false);
-                }
-            }
-
-            try symbolReferencesInternal(builder, ptr_type.ast.sentinel, handle, false);
-            try symbolReferencesInternal(builder, ptr_type.ast.child_type, handle, false);
-        },
-        .array_init,
-        .array_init_comma,
-        .array_init_dot,
-        .array_init_dot_comma,
-        .array_init_one,
-        .array_init_one_comma,
-        .array_init_dot_two,
-        .array_init_dot_two_comma,
-        => {
-            var buf: [2]Ast.Node.Index = undefined;
-            const array_init = tree.fullArrayInit(&buf, node).?;
-
-            try symbolReferencesInternal(builder, array_init.ast.type_expr, handle, false);
-            for (array_init.ast.elements) |e|
-                try symbolReferencesInternal(builder, e, handle, false);
-        },
-        .struct_init,
-        .struct_init_comma,
-        .struct_init_dot,
-        .struct_init_dot_comma,
-        .struct_init_dot_two,
-        .struct_init_dot_two_comma,
-        .struct_init_one,
-        .struct_init_one_comma,
-        => {
-            var buf: [2]Ast.Node.Index = undefined;
-            const struct_init: Ast.full.StructInit = tree.fullStructInit(&buf, node).?;
-
-            try symbolReferencesInternal(builder, struct_init.ast.type_expr, handle, false);
-            for (struct_init.ast.fields) |field|
-                try symbolReferencesInternal(builder, field, handle, false);
-        },
-        .call,
-        .call_comma,
-        .call_one,
-        .call_one_comma,
-        .async_call,
-        .async_call_comma,
-        .async_call_one,
-        .async_call_one_comma,
-        => {
-            var buf: [1]Ast.Node.Index = undefined;
-            const call = tree.fullCall(&buf, node).?;
-
-            try symbolReferencesInternal(builder, call.ast.fn_expr, handle, false);
-
-            for (call.ast.params) |param| {
-                try symbolReferencesInternal(builder, param, handle, false);
-            }
-        },
-        .slice,
-        .slice_sentinel,
-        .slice_open,
-        => {
-            const slice: Ast.full.Slice = tree.fullSlice(node).?;
-
-            try symbolReferencesInternal(builder, slice.ast.sliced, handle, false);
-            try symbolReferencesInternal(builder, slice.ast.start, handle, false);
-            try symbolReferencesInternal(builder, slice.ast.end, handle, false);
-            try symbolReferencesInternal(builder, slice.ast.sentinel, handle, false);
-        },
-        .builtin_call,
-        .builtin_call_comma,
-        .builtin_call_two,
-        .builtin_call_two_comma,
-        => {
-            var buffer: [2]Ast.Node.Index = undefined;
-            const params = ast.builtinCallParams(tree, node, &buffer).?;
-
-            for (params) |param|
-                try symbolReferencesInternal(builder, param, handle, false);
-        },
-        .@"asm",
-        .asm_simple,
-        => |tag| {
-            const full_asm: Ast.full.Asm = if (tag == .@"asm") tree.asmFull(node) else tree.asmSimple(node);
-            if (full_asm.ast.items.len == 0)
-                try symbolReferencesInternal(builder, full_asm.ast.template, handle, false);
-
-            for (full_asm.inputs) |input|
-                try symbolReferencesInternal(builder, input, handle, false);
-
-            for (full_asm.outputs) |output|
-                try symbolReferencesInternal(builder, output, handle, false);
-        },
-        // TODO implement references for asm
-        .asm_output => {},
-        .asm_input => {},
             .field_access => {
-            try symbolReferencesInternal(builder, datas[node].lhs, handle, false);
-
                 var bound_type_params = analysis.BoundTypeParams{};
                 const left_type = try analysis.resolveFieldAccessLhsType(
                     builder.store,
                     builder.arena,
-                (try analysis.resolveTypeOfNodeInternal(builder.store, builder.arena, .{
+                    (try analysis.resolveTypeOfNodeInternal(
-                    .node = datas[node].lhs,
+                        builder.store,
-                    .handle = handle,
+                        builder.arena,
-                }, &bound_type_params)) orelse return,
+                        .{ .node = datas[node].lhs, .handle = handle },
+                        &bound_type_params,
+                    )) orelse return,
                     &bound_type_params,
                 );
 
@@ -337,127 +129,53 @@ fn symbolReferencesInternal(
                     builder.store,
                     builder.arena,
                     .{ .node = left_type_node, .handle = left_type.handle },
-                tree.tokenSlice(datas[node].rhs),
+                    offsets.tokenToSlice(handle.tree, datas[node].rhs),
                     !left_type.type.is_type_val,
                 )) orelse return;
 
-            if (std.meta.eql(child, builder.decl)) try builder.add(handle, datas[node].rhs);
+                if (std.meta.eql(builder.decl_handle, child)) {
-        },
+                    try builder.add(handle, datas[node].rhs);
-        .@"usingnamespace",
-        .unwrap_optional,
-        .bool_not,
-        .negation,
-        .bit_not,
-        .negation_wrap,
-        .address_of,
-        .@"try",
-        .@"await",
-        .optional_type,
-        .deref,
-        .@"suspend",
-        .@"resume",
-        .@"continue",
-        .@"break",
-        .@"return",
-        .grouped_expression,
-        .@"comptime",
-        .@"nosuspend",
-        => try symbolReferencesInternal(builder, datas[node].lhs, handle, false),
-        .test_decl,
-        .@"errdefer",
-        .@"defer",
-        .anyframe_type,
-        => try symbolReferencesInternal(builder, datas[node].rhs, handle, false),
-        .equal_equal,
-        .bang_equal,
-        .less_than,
-        .greater_than,
-        .less_or_equal,
-        .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,
-        .mul,
-        .div,
-        .mod,
-        .array_mult,
-        .mul_wrap,
-        .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(builder, datas[node].lhs, handle, false);
-            try symbolReferencesInternal(builder, datas[node].rhs, handle, false);
-        },
-        .anyframe_literal,
-        .char_literal,
-        .number_literal,
-        .unreachable_literal,
-        .enum_literal,
-        .string_literal,
-        .multiline_string_literal,
-        .error_value,
-        => {},
                 }
-}
+            },
+            else => {},
+        }
+    }
+};
 
 pub fn symbolReferences(
     arena: *std.heap.ArenaAllocator,
     store: *DocumentStore,
     decl_handle: analysis.DeclWithHandle,
     encoding: offsets.Encoding,
+    /// add `decl_handle` as a references
     include_decl: bool,
+    /// exclude references from the std library
     skip_std_references: bool,
+    /// search other files for references
     workspace: bool,
 ) error{OutOfMemory}!std.ArrayListUnmanaged(types.Location) {
-    std.debug.assert(decl_handle.decl.* != .label_decl);
+    std.debug.assert(decl_handle.decl.* != .label_decl); // use `labelReferences` instead
 
-    var builder = Builder.init(arena, store, decl_handle, encoding);
+    var builder = Builder{
+        .arena = arena,
+        .store = store,
+        .decl_handle = decl_handle,
+        .encoding = encoding,
+    };
 
     const curr_handle = decl_handle.handle;
     if (include_decl) try builder.add(curr_handle, decl_handle.nameToken());
 
     switch (decl_handle.decl.*) {
-        .ast_node => {
+        .ast_node,
-            try symbolReferencesInternal(&builder, 0, curr_handle, true);
+        .pointer_payload,
+        .switch_payload,
+        .array_payload,
+        .array_index,
+        => {
+            try builder.collectReferences(curr_handle, 0);
 
-            if (!workspace) return builder.locations;
+            if (decl_handle.decl.* != .ast_node or !workspace) return builder.locations;
 
             var dependencies = std.StringArrayHashMapUnmanaged(void){};
 
@@ -476,24 +194,14 @@ pub fn symbolReferences(
             }
 
             for (dependencies.keys()) |uri| {
+                if (std.mem.eql(u8, uri, curr_handle.uri)) continue;
                 const handle = store.getHandle(uri) orelse continue;
-                if (std.mem.eql(u8, handle.uri, curr_handle.uri)) continue;
 
-                try symbolReferencesInternal(&builder, 0, handle, true);
+                try builder.collectReferences(handle, 0);
             }
         },
-        .pointer_payload,
+        .param_payload => |payload| blk: {
-        .switch_payload,
-        .array_payload,
-        .array_index,
-        => {
-            try symbolReferencesInternal(&builder, 0, curr_handle, true);
-
-            return builder.locations;
-        },
-        .param_payload => |pay| blk: {
             // Rename the param tok.
-            const param = pay.param;
             for (curr_handle.document_scope.scopes.items(.data)) |scope_data| {
                 if (scope_data != .function) continue;
 
@@ -504,10 +212,10 @@ pub fn symbolReferences(
 
                 var it = fn_proto.iterate(&curr_handle.tree);
                 while (ast.nextFnParam(&it)) |candidate| {
-                    if (!std.meta.eql(candidate, param)) continue;
+                    if (!std.meta.eql(candidate, payload.param)) continue;
 
                     if (curr_handle.tree.nodes.items(.tag)[proto] != .fn_decl) break :blk;
-                    try symbolReferencesInternal(&builder, curr_handle.tree.nodes.items(.data)[proto].rhs, curr_handle, false);
+                    try builder.collectReferences(curr_handle, curr_handle.tree.nodes.items(.data)[proto].rhs);
                     break :blk;
                 }
             }

tests/lsp_features/references.zig

@@ -11,9 +11,6 @@ const offsets = zls.offsets;
 
 const allocator: std.mem.Allocator = std.testing.allocator;
 
-// TODO fix references so that we can stop skipping these tests
-const skip_references_tests = true;
-
 test "references" {
     try testReferences(
         \\const <0> = 0;
@@ -59,7 +56,7 @@ test "references - local scope" {
         \\    return <0> + bar;
         \\}
     );
-    if (skip_references_tests) return error.SkipZigTest;
+    if (true) return error.SkipZigTest; // TODO
     try testReferences(
         \\const foo = blk: {
         \\    _ = blk: {
@@ -73,6 +70,32 @@ test "references - local scope" {
     );
 }
 
+test "references - struct field access" {
+    if (true) return error.SkipZigTest; // TODO
+    try testReferences(
+        \\const S = struct {placeholder: u32 = 3};
+        \\pub fn foo() bool {
+        \\    const s: S = .{};
+        \\    return s.<0> == s.<0>;
+        \\}
+    );
+}
+
+test "references - struct decl access" {
+    try testReferences(
+        \\const S = struct {
+        \\    fn <0>() void {}
+        \\};
+        \\pub fn foo() bool {
+        \\    const s: S = .{};
+        \\    s.<0>();
+        \\    s.<0>();
+        \\    <1>();
+        \\}
+        \\fn <1>() void {}
+    );
+}
+
 test "references - while continue expression" {
     try testReferences(
         \\ pub fn foo() void {
@@ -83,7 +106,7 @@ test "references - while continue expression" {
 }
 
 test "references - label" {
-    if (skip_references_tests) return error.SkipZigTest;
+    if (true) return error.SkipZigTest; // TODO
     try testReferences(
         \\const foo = <0>: {
         \\    break :<0> 0;
@@ -106,6 +129,8 @@ fn testReferences(source: []const u8) !void {
 
     try ctx.requestDidOpen(file_uri, phr.new_source);
 
+    try std.testing.expect(phr.locations.len != 0);
+
     var i: usize = 0;
     while (i < phr.locations.len) : (i += 1) {
         const var_loc = phr.locations.items(.old)[i];
@@ -120,6 +145,14 @@ fn testReferences(source: []const u8) !void {
 
         const response = try ctx.requestGetResponse(?[]types.Location, "textDocument/references", params);
 
+        var error_builder = ErrorBuilder.init(allocator, phr.new_source);
+        defer error_builder.deinit();
+        errdefer {
+            const note_loc = phr.locations.items(.new)[i];
+            error_builder.msgAtLoc("asked for references here", note_loc, .info, .{}) catch {};
+            error_builder.writeDebug();
+        }
+
         const locations: []types.Location = response.result orelse {
             std.debug.print("Server returned `null` as the result\n", .{});
             return error.InvalidResponse;
@@ -150,14 +183,6 @@ fn testReferences(source: []const u8) !void {
         };
         defer allocator.free(expected_locs);
 
-        var error_builder = ErrorBuilder.init(allocator, phr.new_source);
-        defer error_builder.deinit();
-        errdefer {
-            const note_loc = phr.locations.items(.new)[i];
-            error_builder.msgAtLoc("asked for references here", note_loc, .info, .{}) catch {};
-            error_builder.writeDebug();
-        }
-
         // keeps track of expected locations that have been given by the server
         // used to detect double references and missing references
         var visited = try std.DynamicBitSetUnmanaged.initEmpty(allocator, expected_locs.len);