fix inefficient semanticTokens/range at the end of files

Author: xdBronch

src/ast.zig

@@ -1641,6 +1641,28 @@ pub fn nodesOverlappingIndexIncludingParseErrors(allocator: std.mem.Allocator, t
     return nodes;
 }
 
+/// Returns a shallow list of nodes that overlap with a loc, i.e. may contain a `fn_decl` node but not its prototype or body.
+/// Caller owns returned memory.
+pub fn nodesIntersectingLoc(allocator: std.mem.Allocator, tree: *const Ast, loc: offsets.Loc) error{OutOfMemory}![]Ast.Node.Index {
+    std.debug.assert(loc.start <= loc.end and loc.end <= tree.source.len);
+
+    var nodes: std.ArrayList(Ast.Node.Index) = .empty;
+    defer nodes.deinit(allocator);
+
+    var walker: Walker = try .init(allocator, tree, .root);
+    defer walker.deinit(allocator);
+
+    while (try walker.nextIgnoreClose(allocator, tree)) |node| {
+        const node_loc = offsets.nodeToLoc(tree, node);
+        if (offsets.locIntersect(loc, node_loc)) {
+            try nodes.append(allocator, node);
+        }
+        walker.skip();
+    }
+
+    return try nodes.toOwnedSlice(allocator);
+}
+
 /// returns a list of nodes that together encloses the given source code range
 /// caller owns the returned memory
 pub fn nodesAtLoc(allocator: std.mem.Allocator, tree: *const Ast, loc: offsets.Loc) error{OutOfMemory}![]Ast.Node.Index {

src/features/inlay_hints.zig

@@ -565,7 +565,7 @@ pub fn writeRangeInlayHint(
         .hover_kind = hover_kind,
     };
 
-    const nodes = try ast.nodesAtLoc(arena, &handle.tree, loc);
+    const nodes = try ast.nodesIntersectingLoc(arena, &handle.tree, loc);
 
     for (nodes) |child| {
         try writeNodeInlayHint(&builder, &handle.tree, child);

src/features/semantic_tokens.zig

@@ -1173,7 +1173,7 @@ pub fn writeSemanticTokens(
         .overlappingTokenSupport = overlappingTokenSupport,
     };
 
-    var nodes = if (loc) |l| try ast.nodesAtLoc(arena, &handle.tree, l) else handle.tree.rootDecls();
+    var nodes = if (loc) |l| try ast.nodesIntersectingLoc(arena, &handle.tree, l) else handle.tree.rootDecls();
     if (nodes.len == 1 and nodes[0] == .root) {
         nodes = handle.tree.rootDecls();
     }

tests/utility/ast.zig

@@ -115,3 +115,70 @@ fn testNodesAtLoc(source: []const u8) !void {
         return error.LocEndMismatch;
     }
 }
+
+test "nodesIntersectingLoc" {
+    try testNodesIntersectingLoc(
+        \\<loc>const a = 0;
+        \\<loc>const b = 0;
+    , &.{.simple_var_decl});
+    try testNodesIntersectingLoc(
+        \\<loc> const a = 0;
+        \\const <loc> b = 0;
+    , &.{ .simple_var_decl, .simple_var_decl });
+    try testNodesIntersectingLoc(
+        \\const a = 0;
+        \\fn foo(_:<loc> u32) void {
+        \\    const b = 0;
+        \\<loc>}
+        \\const c = 0;
+    , &.{.fn_decl});
+    try testNodesIntersectingLoc(
+        \\const a = 0;
+        \\fn foo(_: u32) void {<loc>
+        \\    const b = 0;
+        \\}
+        \\const <loc> c = 0;
+    , &.{ .fn_decl, .simple_var_decl });
+}
+
+fn testNodesIntersectingLoc(
+    source: [:0]const u8,
+    expected_tags: []const std.zig.Ast.Node.Tag,
+) !void {
+    var ccp = try helper.collectClearPlaceholders(allocator, source);
+    defer ccp.deinit(allocator);
+
+    const locs = ccp.locations.items(.new);
+    std.debug.assert(ccp.locations.len == 2);
+
+    const loc: offsets.Loc = .{
+        .start = locs[0].start,
+        .end = locs[1].start,
+    };
+
+    const new_source = try allocator.dupeSentinel(u8, ccp.new_source, 0);
+    defer allocator.free(new_source);
+
+    var tree: std.zig.Ast = try .parse(allocator, new_source, .zig);
+    defer tree.deinit(allocator);
+
+    const nodes = try ast.nodesIntersectingLoc(allocator, &tree, loc);
+    defer allocator.free(nodes);
+
+    const uri = "file.zig";
+    var error_builder: ErrorBuilder = .init(allocator);
+    defer error_builder.deinit();
+    errdefer error_builder.writeDebug();
+
+    try error_builder.addFile(uri, new_source);
+
+    for (nodes, expected_tags) |node, tag| {
+        if (tree.nodeTag(node) != tag) {
+            try error_builder.msgAtIndex("expected '{t}' found '{t}'", uri, loc.start, .err, .{
+                tag,
+                tree.nodeTag(node),
+            });
+            return error.NodeMismatch;
+        }
+    }
+}