Merge pull request #2348 from xdBronch/references

speed up reference finding

Author: Techatrix

src/DocumentScope.zig

@@ -77,70 +77,13 @@ pub const DeclarationLookupContext = struct {
     }
 };
 
-/// Assumes that the `node` is not a container_field of a struct tuple field.
-/// Returns a `.identifier` or `.builtin` token.
-fn getDeclNameToken(tree: Ast, node: Ast.Node.Index) ?Ast.TokenIndex {
-    var buffer: [1]Ast.Node.Index = undefined;
-    const token_index = switch (tree.nodeTag(node)) {
-        .local_var_decl,
-        .global_var_decl,
-        .simple_var_decl,
-        .aligned_var_decl,
-        => tree.nodeMainToken(node) + 1,
-        .fn_proto,
-        .fn_proto_multi,
-        .fn_proto_one,
-        .fn_proto_simple,
-        .fn_decl,
-        => tree.fullFnProto(&buffer, node).?.name_token orelse return null,
-
-        .identifier => tree.nodeMainToken(node),
-        .test_decl => tree.nodeData(node).opt_token_and_node[0].unwrap() orelse return null,
-
-        .container_field,
-        .container_field_init,
-        .container_field_align,
-        => tree.nodeMainToken(node),
-
-        .root,
-        .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,
-        .block,
-        .block_semicolon,
-        .block_two,
-        .block_two_semicolon,
-        => return null,
-
-        else => unreachable,
-    };
-
-    return switch (tree.tokenTag(token_index)) {
-        .identifier, .builtin => token_index,
-        else => null,
-    };
-}
-
 pub const Declaration = union(enum) {
     /// Index of the ast node.
     /// Can have one of the following tags:
-    ///   - `.root`
-    ///   - `.container_decl`
-    ///   - `.tagged_union`
     ///   - `.container_field`
     ///   - `.fn_proto`
     ///   - `.fn_decl`
     ///   - `.var_decl`
-    ///   - `.block`
     ast_node: Ast.Node.Index,
     /// Function parameter
     function_parameter: Param,
@@ -250,7 +193,34 @@ pub const Declaration = union(enum) {
     /// Returns a `.identifier` or `.builtin` token.
     pub fn nameToken(decl: Declaration, tree: Ast) Ast.TokenIndex {
         return switch (decl) {
-            .ast_node => |n| getDeclNameToken(tree, n).?,
+            .ast_node => |node| {
+                var buffer: [1]Ast.Node.Index = undefined;
+                const token_index = switch (tree.nodeTag(node)) {
+                    .local_var_decl,
+                    .global_var_decl,
+                    .simple_var_decl,
+                    .aligned_var_decl,
+                    => tree.nodeMainToken(node) + 1,
+                    .fn_proto,
+                    .fn_proto_multi,
+                    .fn_proto_one,
+                    .fn_proto_simple,
+                    .fn_decl,
+                    => tree.fullFnProto(&buffer, node).?.name_token.?,
+
+                    .container_field,
+                    .container_field_init,
+                    .container_field_align,
+                    => tree.nodeMainToken(node),
+
+                    else => unreachable,
+                };
+
+                switch (tree.tokenTag(token_index)) {
+                    .identifier, .builtin => return token_index,
+                    else => unreachable,
+                }
+            },
             .function_parameter => |payload| payload.get(tree).?.name_token.?,
             .optional_payload => |payload| payload.identifier,
             .error_union_payload => |payload| payload.identifier,

src/analysis.zig

@@ -3867,7 +3867,7 @@ pub const Type = struct {
                             return;
                         }
                         if (token >= 1 and tree.tokenTag(token - 1) == .keyword_return) blk: {
-                            const function_scope = innermostFunctionScopeAtIndex(doc_scope, tree.tokenStart(token - 1)).unwrap() orelse break :blk;
+                            const function_scope = innermostScopeAtIndexWithTag(doc_scope, tree.tokenStart(token - 1), .initOne(.function)).unwrap() orelse break :blk;
                             const function_node = doc_scope.getScopeAstNode(function_scope).?;
                             var buf: [1]Ast.Node.Index = undefined;
                             const func = tree.fullFnProto(&buf, function_node).?;
@@ -5293,36 +5293,28 @@ pub fn iterateLabels(handle: *DocumentStore.Handle, source_index: usize, comptim
     }
 }
 
-pub fn innermostScopeAtIndex(document_scope: DocumentScope, source_index: usize) Scope.Index {
-    var scope_iterator = iterateEnclosingScopes(&document_scope, source_index);
-    var scope_index: Scope.Index = scope_iterator.next().unwrap().?; // the DocumentScope's root scope must exist
-    while (scope_iterator.next().unwrap()) |inner_scope| {
-        scope_index = inner_scope;
-    }
-    return scope_index;
+pub fn innermostScopeAtIndex(
+    document_scope: DocumentScope,
+    source_index: usize,
+) Scope.Index {
+    return innermostScopeAtIndexWithTag(document_scope, source_index, .initFull()).unwrap().?;
 }
 
-pub fn innermostFunctionScopeAtIndex(document_scope: DocumentScope, source_index: usize) Scope.OptionalIndex {
+pub fn innermostScopeAtIndexWithTag(
+    document_scope: DocumentScope,
+    source_index: usize,
+    tag_filter: std.EnumSet(Scope.Tag),
+) Scope.OptionalIndex {
     var scope_iterator = iterateEnclosingScopes(&document_scope, source_index);
     var scope_index: Scope.OptionalIndex = .none;
     while (scope_iterator.next().unwrap()) |inner_scope| {
-        if (document_scope.getScopeTag(inner_scope) != .function) continue;
+        const scope_tag = document_scope.getScopeTag(inner_scope);
+        if (!tag_filter.contains(scope_tag)) continue;
         scope_index = inner_scope.toOptional();
     }
     return scope_index;
 }
 
-pub fn innermostBlockScope(document_scope: DocumentScope, source_index: usize) Ast.Node.Index {
-    var scope_iterator = iterateEnclosingScopes(&document_scope, source_index);
-    var ast_node: ?Ast.Node.Index = null;
-    while (scope_iterator.next().unwrap()) |inner_scope| {
-        if (document_scope.getScopeAstNode(inner_scope)) |node| {
-            ast_node = node;
-        }
-    }
-    return ast_node.?; // the DocumentScope's root scope is guaranteed to have an Ast Node
-}
-
 pub fn innermostContainer(analyser: *Analyser, handle: *DocumentStore.Handle, source_index: usize) error{OutOfMemory}!Type {
     const tree = handle.tree;
     const document_scope = try handle.getDocumentScope();

src/features/references.zig

@@ -65,6 +65,8 @@ const Builder = struct {
     locations: std.ArrayListUnmanaged(types.Location) = .empty,
     /// this is the declaration we are searching for
     decl_handle: Analyser.DeclWithHandle,
+    /// the decl is local to a function, block, etc
+    local_only_decl: bool,
     /// Whether the `decl_handle` has been added
     did_add_decl_handle: bool = false,
     analyser: *Analyser,
@@ -104,6 +106,10 @@ const Builder = struct {
     fn referenceNode(self: *const Context, tree: Ast, node: Ast.Node.Index) error{OutOfMemory}!void {
         const builder = self.builder;
         const handle = self.handle;
+        const decl_name = offsets.identifierTokenToNameSlice(
+            builder.decl_handle.handle.tree,
+            builder.decl_handle.nameToken(),
+        );
 
         switch (tree.nodeTag(node)) {
             .identifier,
@@ -119,6 +125,7 @@ const Builder = struct {
                     else => unreachable,
                 };
                 const name = offsets.identifierTokenToNameSlice(tree, name_token);
+                if (!std.mem.eql(u8, name, decl_name)) return;
 
                 const child = try builder.analyser.lookupSymbolGlobal(
                     handle,
@@ -131,15 +138,18 @@ const Builder = struct {
                 }
             },
             .field_access => {
-                const lhs_node, const field_name = tree.nodeData(node).node_and_token;
+                if (builder.local_only_decl) return;
+                const lhs_node, const field_token = tree.nodeData(node).node_and_token;
+                const name = offsets.identifierTokenToNameSlice(tree, field_token);
+                if (!std.mem.eql(u8, name, decl_name)) return;
+
                 const lhs = try builder.analyser.resolveTypeOfNode(.of(lhs_node, handle)) orelse return;
                 const deref_lhs = try builder.analyser.resolveDerefType(lhs) orelse lhs;
 
-                const symbol = offsets.identifierTokenToNameSlice(tree, field_name);
-                const child = try deref_lhs.lookupSymbol(builder.analyser, symbol) orelse return;
+                const child = try deref_lhs.lookupSymbol(builder.analyser, name) orelse return;
 
                 if (builder.decl_handle.eql(child)) {
-                    try builder.add(handle, field_name);
+                    try builder.add(handle, field_token);
                 }
             },
             .struct_init_one,
@@ -151,23 +161,28 @@ const Builder = struct {
             .struct_init_dot_two,
             .struct_init_dot_two_comma,
             => {
+                if (builder.local_only_decl) return;
                 var buffer: [2]Ast.Node.Index = undefined;
                 const struct_init = tree.fullStructInit(&buffer, node).?;
                 for (struct_init.ast.fields) |value_node| { // the node of `value` in `.name = value`
                     const name_token = tree.firstToken(value_node) - 2; // math our way two token indexes back to get the `name`
                     const name_loc = offsets.tokenToLoc(tree, name_token);
                     const name = offsets.locToSlice(tree.source, name_loc);
+                    if (!std.mem.eql(u8, name, decl_name)) continue;
 
                     const lookup = try builder.analyser.lookupSymbolFieldInit(handle, name, node, &.{}) orelse continue;
 
                     if (builder.decl_handle.eql(lookup)) {
                         try builder.add(handle, name_token);
+                        return;
                     }
                 }
             },
             .enum_literal => {
+                if (builder.local_only_decl) return;
                 const name_token = tree.nodeMainToken(node);
                 const name = offsets.identifierTokenToNameSlice(handle.tree, name_token);
+                if (!std.mem.eql(u8, name, decl_name)) return;
                 const lookup = try builder.analyser.getSymbolEnumLiteral(handle, tree.tokenStart(name_token), name) orelse return;
 
                 if (builder.decl_handle.eql(lookup)) {
@@ -242,48 +257,61 @@ fn symbolReferences(
 
     std.debug.assert(decl_handle.decl != .label); // use `labelReferences` instead
 
-    var builder = Builder{
-        .allocator = allocator,
-        .analyser = analyser,
-        .decl_handle = decl_handle,
-        .encoding = encoding,
-    };
-    errdefer builder.deinit();
-
-    const curr_handle = decl_handle.handle;
-    if (include_decl) try builder.add(curr_handle, decl_handle.nameToken());
-
-    switch (decl_handle.decl) {
-        .ast_node => {
-            try builder.collectReferences(curr_handle, .root);
-
-            const source_index = decl_handle.handle.tree.tokenStart(decl_handle.nameToken());
-            // highlight requests only pertain to the current document, otherwise we can try to narrow things down
-            const workspace = if (request == .highlight) false else blk: {
-                const doc_scope = try curr_handle.getDocumentScope();
-                const scope_index = Analyser.innermostScopeAtIndex(doc_scope, source_index);
-                break :blk switch (doc_scope.getScopeTag(scope_index)) {
-                    .function, .block => false,
-                    .container, .container_usingnamespace => decl_handle.isPublic(),
-                    .other => true,
-                };
-            };
-            if (workspace) {
-                try gatherReferences(allocator, analyser, curr_handle, skip_std_references, include_decl, &builder, .get);
-            }
+    const doc_scope = try decl_handle.handle.getDocumentScope();
+    const source_index = decl_handle.handle.tree.tokenStart(decl_handle.nameToken());
+    const scope_index = Analyser.innermostScopeAtIndexWithTag(doc_scope, source_index, .init(.{
+        .block = true,
+        .container = true,
+        .container_usingnamespace = true,
+        .function = false,
+        .other = false,
+    })).unwrap().?;
+    const scope_node = doc_scope.getScopeAstNode(scope_index).?;
+
+    // If `local_node != null`, references to the declaration can only be
+    // found inside of the given ast node.
+    const local_node: ?Ast.Node.Index = switch (decl_handle.decl) {
+        .ast_node => switch (doc_scope.getScopeTag(scope_index)) {
+            .block => scope_node,
+            .container, .container_usingnamespace => null,
+            .function, .other => unreachable,
         },
         .optional_payload,
         .error_union_payload,
         .error_union_error,
         .for_loop_payload,
         .assign_destructure,
-        .switch_payload,
-        => {
-            try builder.collectReferences(curr_handle, .root);
-        },
-        .function_parameter => |payload| try builder.collectReferences(curr_handle, payload.func),
+        => scope_node,
+        .switch_payload => |payload| payload.node,
+        .function_parameter => |payload| payload.func,
         .label => unreachable, // handled separately by labelReferences
-        .error_token => {},
+        .error_token => return .empty,
+    };
+
+    var builder: Builder = .{
+        .allocator = allocator,
+        .analyser = analyser,
+        .decl_handle = decl_handle,
+        .local_only_decl = local_node != null,
+        .encoding = encoding,
+    };
+    errdefer builder.deinit();
+
+    if (include_decl) try builder.add(decl_handle.handle, decl_handle.nameToken());
+
+    try builder.collectReferences(decl_handle.handle, local_node orelse .root);
+
+    const workspace = local_node == null and request != .highlight and decl_handle.isPublic();
+    if (workspace) {
+        try gatherReferences(
+            allocator,
+            analyser,
+            decl_handle.handle,
+            skip_std_references,
+            include_decl,
+            &builder,
+            .get,
+        );
     }
 
     return builder.locations;

src/features/signature_help.zig

@@ -78,7 +78,14 @@ pub fn getSignatureInfo(
     markup_kind: types.MarkupKind,
 ) !?types.SignatureInformation {
     const document_scope = try handle.getDocumentScope();
-    const innermost_block = Analyser.innermostBlockScope(document_scope, absolute_index);
+    const innermost_block_scope = Analyser.innermostScopeAtIndexWithTag(document_scope, absolute_index, .init(.{
+        .block = true,
+        .container = true,
+        .container_usingnamespace = true,
+        .function = true,
+        .other = false,
+    })).unwrap().?;
+    const innermost_block = document_scope.getScopeAstNode(innermost_block_scope).?;
     const tree = handle.tree;
 
     // Use the innermost scope to determine the earliest token we would need