Optimize inlay hints (#948)

* optimize inlay hints

* update iterateChildren

* add tests for nodesAtLoc

src/Server.zig

@@ -2493,6 +2493,7 @@ fn inlayHintHandler(server: *Server, request: types.InlayHintParams) Error!?[]ty
     const handle = server.document_store.getHandle(request.textDocument.uri) orelse return null;
 
     const hover_kind: types.MarkupKind = if (server.client_capabilities.hover_supports_md) .markdown else .plaintext;
+    const loc = offsets.rangeToLoc(handle.text, request.range, server.offset_encoding);
 
     // TODO cache hints per document
     // because the function could be stored in a different document
@@ -2503,27 +2504,44 @@ fn inlayHintHandler(server: *Server, request: types.InlayHintParams) Error!?[]ty
         server.config.*,
         &server.document_store,
         handle,
-        request.range,
+        loc,
         hover_kind,
-        server.offset_encoding,
     );
 
-    // and only convert and return all hints in range for every request
-    var visible_hints = hints;
+    const helper = struct {
+        fn lessThan(_: void, lhs: inlay_hints.InlayHint, rhs: inlay_hints.InlayHint) bool {
+            return lhs.token_index < rhs.token_index;
+        }
+    };
 
-    // small_hints should roughly be sorted by position
+    std.sort.sort(inlay_hints.InlayHint, hints, {}, helper.lessThan);
+
+    var last_index: usize = 0;
+    var last_position: types.Position = .{ .line = 0, .character = 0 };
+
+    var converted_hints = try server.arena.allocator().alloc(types.InlayHint, hints.len);
     for (hints) |hint, i| {
-        if (isPositionBefore(hint.position, request.range.start)) continue;
-        visible_hints = hints[i..];
-        break;
-    }
-    for (visible_hints) |hint, i| {
-        if (isPositionBefore(hint.position, request.range.end)) continue;
-        visible_hints = visible_hints[0..i];
-        break;
+        const index = offsets.tokenToIndex(handle.tree, hint.token_index);
+        const position = offsets.advancePosition(
+            handle.tree.source,
+            last_position,
+            last_index,
+            index,
+            server.offset_encoding,
+        );
+        defer last_index = index;
+        defer last_position = position;
+        converted_hints[i] = types.InlayHint{
+            .position = position,
+            .label = .{ .string = hint.label },
+            .kind = hint.kind,
+            .tooltip = .{ .MarkupContent = hint.tooltip },
+            .paddingLeft = false,
+            .paddingRight = true,
+        };
     }
 
-    return visible_hints;
+    return converted_hints;
 }
 
 fn codeActionHandler(server: *Server, request: types.CodeActionParams) Error!?[]types.CodeAction {

src/ast.zig

@@ -3,6 +3,7 @@
 //! when there are parser errors.
 
 const std = @import("std");
+const offsets = @import("offsets.zig");
 const Ast = std.zig.Ast;
 const Node = Ast.Node;
 const full = Ast.full;
@@ -1140,3 +1141,452 @@ pub fn nextFnParam(it: *Ast.full.FnProto.Iterator) ?Ast.full.FnProto.Param {
         it.tok_flag = false;
     }
 }
+
+/// returns an Iterator that yields every child of the given node.
+/// see `nodeChildrenAlloc` for a non-iterator allocating variant.
+/// the order in which children are given corresponds to the order in which they are found in the source text
+pub fn iterateChildren(
+    tree: Ast,
+    node: Ast.Node.Index,
+    context: anytype,
+    comptime Error: type,
+    comptime callback: fn (@TypeOf(context), Ast.Node.Index) Error!void,
+) Error!void {
+    const node_tags = tree.nodes.items(.tag);
+    const node_data = tree.nodes.items(.data);
+
+    if (node > tree.nodes.len) return;
+
+    const tag = node_tags[node];
+    switch (tag) {
+        .@"usingnamespace",
+        .field_access,
+        .unwrap_optional,
+        .bool_not,
+        .negation,
+        .bit_not,
+        .negation_wrap,
+        .address_of,
+        .@"try",
+        .@"await",
+        .optional_type,
+        .deref,
+        .@"suspend",
+        .@"resume",
+        .@"return",
+        .grouped_expression,
+        .@"comptime",
+        .@"nosuspend",
+        .asm_simple,
+        => {
+            try callback(context, node_data[node].lhs);
+        },
+
+        .test_decl,
+        .@"errdefer",
+        .@"defer",
+        .@"break",
+        .anyframe_type,
+        => {
+            try callback(context, node_data[node].rhs);
+        },
+
+        .@"catch",
+        .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_access,
+        .array_init_one,
+        .array_init_one_comma,
+        .array_init_dot_two,
+        .array_init_dot_two_comma,
+        .struct_init_one,
+        .struct_init_one_comma,
+        .struct_init_dot_two,
+        .struct_init_dot_two_comma,
+        .call_one,
+        .call_one_comma,
+        .async_call_one,
+        .async_call_one_comma,
+        .switch_range,
+        .builtin_call_two,
+        .builtin_call_two_comma,
+        .container_decl_two,
+        .container_decl_two_trailing,
+        .tagged_union_two,
+        .tagged_union_two_trailing,
+        .container_field_init,
+        .container_field_align,
+        .block_two,
+        .block_two_semicolon,
+        .error_union,
+        => {
+            try callback(context, node_data[node].lhs);
+            try callback(context, node_data[node].rhs);
+        },
+
+        .root,
+        .array_init_dot,
+        .array_init_dot_comma,
+        .struct_init_dot,
+        .struct_init_dot_comma,
+        .builtin_call,
+        .builtin_call_comma,
+        .container_decl,
+        .container_decl_trailing,
+        .tagged_union,
+        .tagged_union_trailing,
+        .block,
+        .block_semicolon,
+        => {
+            for (tree.extra_data[node_data[node].lhs..node_data[node].rhs]) |child| {
+                try callback(context, child);
+            }
+        },
+
+        .global_var_decl,
+        .local_var_decl,
+        .simple_var_decl,
+        .aligned_var_decl,
+        => {
+            const var_decl = tree.fullVarDecl(node).?.ast;
+            try callback(context, var_decl.type_node);
+            try callback(context, var_decl.align_node);
+            try callback(context, var_decl.addrspace_node);
+            try callback(context, var_decl.section_node);
+            try callback(context, var_decl.init_node);
+        },
+
+        .array_type_sentinel => {
+            const array_type = tree.arrayTypeSentinel(node).ast;
+            try callback(context, array_type.elem_count);
+            try callback(context, array_type.sentinel);
+            try callback(context, array_type.elem_type);
+        },
+
+        .ptr_type_aligned,
+        .ptr_type_sentinel,
+        .ptr_type,
+        .ptr_type_bit_range,
+        => {
+            const ptr_type = fullPtrType(tree, node).?.ast;
+            try callback(context, ptr_type.sentinel);
+            try callback(context, ptr_type.align_node);
+            try callback(context, ptr_type.bit_range_start);
+            try callback(context, ptr_type.bit_range_end);
+            try callback(context, ptr_type.addrspace_node);
+            try callback(context, ptr_type.child_type);
+        },
+
+        .slice_open,
+        .slice,
+        .slice_sentinel,
+        => {
+            const slice = tree.fullSlice(node).?;
+            try callback(context, slice.ast.sliced);
+            try callback(context, slice.ast.start);
+            try callback(context, slice.ast.end);
+            try callback(context, slice.ast.sentinel);
+        },
+
+        .array_init,
+        .array_init_comma,
+        => {
+            const array_init = tree.arrayInit(node).ast;
+            try callback(context, array_init.type_expr);
+            for (array_init.elements) |child| {
+                try callback(context, child);
+            }
+        },
+
+        .struct_init,
+        .struct_init_comma,
+        => {
+            const struct_init = tree.structInit(node).ast;
+            try callback(context, struct_init.type_expr);
+            for (struct_init.fields) |child| {
+                try callback(context, child);
+            }
+        },
+
+        .call,
+        .call_comma,
+        .async_call,
+        .async_call_comma,
+        => {
+            const call = tree.callFull(node).ast;
+            try callback(context, call.fn_expr);
+            for (call.params) |child| {
+                try callback(context, child);
+            }
+        },
+
+        .@"switch",
+        .switch_comma,
+        => {
+            const cond = node_data[node].lhs;
+            const extra = tree.extraData(node_data[node].rhs, Ast.Node.SubRange);
+            const cases = tree.extra_data[extra.start..extra.end];
+            try callback(context, cond);
+            for (cases) |child| {
+                try callback(context, child);
+            }
+        },
+
+        .switch_case_one,
+        .switch_case_inline_one,
+        .switch_case,
+        .switch_case_inline,
+        => {
+            const switch_case = tree.fullSwitchCase(node).?.ast;
+            for (switch_case.values) |child| {
+                try callback(context, child);
+            }
+            try callback(context, switch_case.target_expr);
+        },
+
+        .while_simple,
+        .while_cont,
+        .@"while",
+        .for_simple,
+        .@"for",
+        => {
+            const while_ast = fullWhile(tree, node).?.ast;
+            try callback(context, while_ast.cond_expr);
+            try callback(context, while_ast.cont_expr);
+            try callback(context, while_ast.then_expr);
+            try callback(context, while_ast.else_expr);
+        },
+
+        .@"if",
+        .if_simple,
+        => {
+            const if_ast = ifFull(tree, node).ast;
+            try callback(context, if_ast.cond_expr);
+            try callback(context, if_ast.then_expr);
+            try callback(context, if_ast.else_expr);
+        },
+
+        .fn_proto_simple,
+        .fn_proto_multi,
+        .fn_proto_one,
+        .fn_proto,
+        .fn_decl,
+        => {
+            var buffer: [1]Node.Index = undefined;
+            const fn_proto = tree.fullFnProto(&buffer, node).?;
+
+            for (fn_proto.ast.params) |child| {
+                try callback(context, child);
+            }
+            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);
+        },
+
+        .container_decl_arg,
+        .container_decl_arg_trailing,
+        => {
+            const decl = tree.containerDeclArg(node).ast;
+            try callback(context, decl.arg);
+            for (decl.members) |child| {
+                try callback(context, child);
+            }
+        },
+
+        .tagged_union_enum_tag,
+        .tagged_union_enum_tag_trailing,
+        => {
+            const decl = tree.taggedUnionEnumTag(node).ast;
+            try callback(context, decl.arg);
+            for (decl.members) |child| {
+                try callback(context, child);
+            }
+        },
+
+        .container_field => {
+            const field = tree.containerField(node).ast;
+            try callback(context, field.type_expr);
+            try callback(context, field.align_expr);
+            try callback(context, field.value_expr);
+        },
+
+        .@"asm" => {
+            const asm_ast = tree.asmFull(node).ast;
+            try callback(context, asm_ast.template);
+            for (asm_ast.items) |child| {
+                try callback(context, child);
+            }
+        },
+
+        .asm_output,
+        .asm_input,
+        => {}, // TODO
+
+        .@"continue",
+        .anyframe_literal,
+        .char_literal,
+        .number_literal,
+        .unreachable_literal,
+        .identifier,
+        .enum_literal,
+        .string_literal,
+        .multiline_string_literal,
+        .error_set_decl,
+        .error_value,
+        => {},
+    }
+}
+
+/// returns an Iterator that recursively yields every child of the given node.
+/// see `nodeChildrenRecursiveAlloc` for a non-iterator allocating variant.
+pub fn iterateChildrenRecursive(
+    tree: Ast,
+    node: Ast.Node.Index,
+    context: anytype,
+    comptime Error: type,
+    comptime callback: fn (@TypeOf(context), Ast.Node.Index) Error!void,
+) Error!void {
+    const RecursiveContext = struct {
+        tree: Ast,
+        context: @TypeOf(context),
+
+        fn recursive_callback(self: @This(), child_node: Ast.Node.Index) Error!void {
+            if (child_node == 0) return;
+            try callback(self.context, child_node);
+            try iterateChildrenRecursive(self.tree, child_node, self.context, Error, callback);
+        }
+    };
+
+    try iterateChildren(tree, node, RecursiveContext{
+        .tree = tree,
+        .context = context,
+    }, Error, RecursiveContext.recursive_callback);
+}
+
+/// returns the children of the given node.
+/// see `iterateChildren` for a callback variant
+/// caller owns the returned memory
+pub fn nodeChildrenAlloc(allocator: std.mem.Allocator, tree: Ast, node: Ast.Node.Index) error{OutOfMemory}![]Ast.Node.Index {
+    const Context = struct {
+        children: *std.ArrayList(Ast.Node.Index),
+        fn callback(self: @This(), child_node: Ast.Node.Index) error{OutOfMemory}!void {
+            if (child_node == 0) return;
+            try self.children.append(child_node);
+        }
+    };
+
+    var children = std.ArrayList(Ast.Node.Index).init(allocator);
+    errdefer children.deinit();
+    try iterateChildren(tree, node, Context{ .children = &children }, error{OutOfMemory}, Context.callback);
+    return children.toOwnedSlice();
+}
+
+/// returns the children of the given node.
+/// see `iterateChildrenRecursive` for a callback variant
+/// caller owns the returned memory
+pub fn nodeChildrenRecursiveAlloc(allocator: std.mem.Allocator, tree: Ast, node: Ast.Node.Index) error{OutOfMemory}![]Ast.Node.Index {
+    const Context = struct {
+        children: *std.ArrayList(Ast.Node.Index),
+        fn callback(self: @This(), child_node: Ast.Node.Index) error{OutOfMemory}!void {
+            if (child_node == 0) return;
+            try self.children.append(child_node);
+        }
+    };
+
+    var children = std.ArrayList(Ast.Node.Index).init(allocator);
+    errdefer children.deinit();
+    try iterateChildrenRecursive(tree, node, .{ .children = &children }, Context.callback);
+    return children.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: 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.ArrayListUnmanaged(Ast.Node.Index){};
+    errdefer nodes.deinit(allocator);
+    var parent: Ast.Node.Index = 0; // root node
+
+    try nodes.ensureTotalCapacity(allocator, 32);
+
+    while (true) {
+        const children = try nodeChildrenAlloc(allocator, tree, parent);
+        defer allocator.free(children);
+
+        var children_loc: ?offsets.Loc = null;
+        for (children) |child_node| {
+            const child_loc = offsets.nodeToLoc(tree, child_node);
+
+            const merge_child = offsets.locIntersect(loc, child_loc) or offsets.locInside(child_loc, loc);
+
+            if (merge_child) {
+                children_loc = if (children_loc) |l| offsets.locMerge(l, child_loc) else child_loc;
+                try nodes.append(allocator, child_node);
+            } else {
+                if (nodes.items.len != 0) break;
+            }
+        }
+
+        if (children_loc == null or !offsets.locInside(loc, children_loc.?)) {
+            nodes.clearRetainingCapacity();
+            nodes.appendAssumeCapacity(parent); // capacity is never 0
+            return try nodes.toOwnedSlice(allocator);
+        }
+
+        if (nodes.items.len == 1) {
+            parent = nodes.items[0];
+            nodes.clearRetainingCapacity();
+        } else {
+            return try nodes.toOwnedSlice(allocator);
+        }
+    }
+}

src/inlay_hints.zig

@@ -16,79 +16,72 @@ const Config = @import("Config.zig");
 /// non-configurable at runtime
 pub const inlay_hints_exclude_builtins: []const u8 = &.{};
 
-/// max number of children in a declaration/array-init/struct-init or similar
-/// that will not get a visibility check
-pub const inlay_hints_max_inline_children = 12;
-
-/// checks whether node is inside the range
-fn isNodeInRange(tree: Ast, node: Ast.Node.Index, range: types.Range) bool {
-    const endLocation = tree.tokenLocation(0, ast.lastToken(tree, node));
-    if (endLocation.line < range.start.line) return false;
-
-    const beginLocation = tree.tokenLocation(0, tree.firstToken(node));
-    if (beginLocation.line > range.end.line) return false;
-
-    return true;
-}
+pub const InlayHint = struct {
+    token_index: Ast.TokenIndex,
+    label: []const u8,
+    kind: types.InlayHintKind,
+    tooltip: types.MarkupContent,
+};
 
 const Builder = struct {
-    arena: std.mem.Allocator,
+    arena: *std.heap.ArenaAllocator,
+    store: *DocumentStore,
     config: *const Config,
     handle: *const DocumentStore.Handle,
-    hints: std.ArrayListUnmanaged(types.InlayHint),
+    hints: std.ArrayListUnmanaged(InlayHint),
     hover_kind: types.MarkupKind,
-    encoding: offsets.Encoding,
 
-    fn appendParameterHint(self: *Builder, position: types.Position, label: []const u8, tooltip: []const u8, tooltip_noalias: bool, tooltip_comptime: bool) !void {
-        // TODO allocation could be avoided by extending InlayHint.jsonStringify
+    fn appendParameterHint(self: *Builder, token_index: Ast.TokenIndex, label: []const u8, tooltip: []const u8, tooltip_noalias: bool, tooltip_comptime: bool) !void {
         // adding tooltip_noalias & tooltip_comptime to InlayHint should be enough
         const tooltip_text = blk: {
             if (tooltip.len == 0) break :blk "";
             const prefix = if (tooltip_noalias) if (tooltip_comptime) "noalias comptime " else "noalias " else if (tooltip_comptime) "comptime " else "";
 
             if (self.hover_kind == .markdown) {
-                break :blk try std.fmt.allocPrint(self.arena, "```zig\n{s}{s}\n```", .{ prefix, tooltip });
+                break :blk try std.fmt.allocPrint(self.arena.allocator(), "```zig\n{s}{s}\n```", .{ prefix, tooltip });
             }
 
-            break :blk try std.fmt.allocPrint(self.arena, "{s}{s}", .{ prefix, tooltip });
+            break :blk try std.fmt.allocPrint(self.arena.allocator(), "{s}{s}", .{ prefix, tooltip });
         };
 
-        try self.hints.append(self.arena, .{
-            .position = position,
-            .label = .{ .string = try std.fmt.allocPrint(self.arena, "{s}:", .{label}) },
-            .kind = types.InlayHintKind.Parameter,
-            .tooltip = .{ .MarkupContent = .{
+        try self.hints.append(self.arena.allocator(), .{
+            .token_index = token_index,
+            .label = try std.fmt.allocPrint(self.arena.allocator(), "{s}:", .{label}),
+            .kind = .Parameter,
+            .tooltip = .{
                 .kind = self.hover_kind,
                 .value = tooltip_text,
-            } },
-            .paddingLeft = false,
-            .paddingRight = true,
+            },
         });
     }
 
-    fn toOwnedSlice(self: *Builder) error{OutOfMemory}![]types.InlayHint {
-        return self.hints.toOwnedSlice(self.arena);
+    fn toOwnedSlice(self: *Builder) error{OutOfMemory}![]InlayHint {
+        return self.hints.toOwnedSlice(self.arena.allocator());
     }
 };
 
 /// `call` is the function call
 /// `decl_handle` should be a function protototype
 /// writes parameter hints into `builder.hints`
-fn writeCallHint(builder: *Builder, arena: *std.heap.ArenaAllocator, store: *DocumentStore, call: Ast.full.Call, decl_handle: analysis.DeclWithHandle) !void {
+fn writeCallHint(builder: *Builder, call: Ast.full.Call, decl_handle: analysis.DeclWithHandle) !void {
     const handle = builder.handle;
     const tree = handle.tree;
 
     const decl = decl_handle.decl;
     const decl_tree = decl_handle.handle.tree;
 
-    switch (decl.*) {
-        .ast_node => |fn_node| {
+    const fn_node = switch (decl.*) {
+        .ast_node => |fn_node| fn_node,
+        else => return,
+    };
+
     var buffer: [1]Ast.Node.Index = undefined;
-            if (decl_tree.fullFnProto(&buffer, fn_node)) |fn_proto| {
+    const fn_proto = decl_tree.fullFnProto(&buffer, fn_node) orelse return;
+
     var i: usize = 0;
     var it = fn_proto.iterate(&decl_tree);
 
-                if (try analysis.hasSelfParam(arena, store, decl_handle.handle, fn_proto)) {
+    if (try analysis.hasSelfParam(builder.arena, builder.store, decl_handle.handle, fn_proto)) {
         _ = ast.nextFnParam(&it);
     }
 
@@ -123,17 +116,13 @@ fn writeCallHint(builder: *Builder, arena: *std.heap.ArenaAllocator, store: *Doc
             offsets.nodeToSlice(decl_tree, param.type_expr);
 
         try builder.appendParameterHint(
-                        offsets.tokenToPosition(tree, tree.firstToken(call.ast.params[i]), builder.encoding),
+            tree.firstToken(call.ast.params[i]),
             name,
             tooltip,
             no_alias,
             comp_time,
         );
     }
-            }
-        },
-        else => {},
-    }
 }
 
 /// takes parameter nodes from the ast and function parameter names from `Builtin.arguments` and writes parameter hints into `builder.hints`
@@ -164,7 +153,7 @@ fn writeBuiltinHint(builder: *Builder, parameters: []const Ast.Node.Index, argum
         }
 
         try builder.appendParameterHint(
-            offsets.tokenToPosition(tree, tree.firstToken(parameters[i]), builder.encoding),
+            tree.firstToken(parameters[i]),
             label orelse "",
             std.mem.trim(u8, type_expr, " \t\n"),
             no_alias,
@@ -174,7 +163,7 @@ fn writeBuiltinHint(builder: *Builder, parameters: []const Ast.Node.Index, argum
 }
 
 /// takes a Ast.full.Call (a function call), analysis its function expression, finds its declaration and writes parameter hints into `builder.hints`
-fn writeCallNodeHint(builder: *Builder, arena: *std.heap.ArenaAllocator, store: *DocumentStore, call: Ast.full.Call) !void {
+fn writeCallNodeHint(builder: *Builder, call: Ast.full.Call) !void {
     if (call.ast.params.len == 0) return;
     if (builder.config.inlay_hints_exclude_single_argument and call.ast.params.len == 1) return;
 
@@ -187,14 +176,11 @@ fn writeCallNodeHint(builder: *Builder, arena: *std.heap.ArenaAllocator, store:
 
     switch (node_tags[call.ast.fn_expr]) {
         .identifier => {
-            const location = tree.tokenLocation(0, main_tokens[call.ast.fn_expr]);
+            const source_index = offsets.tokenToIndex(tree, main_tokens[call.ast.fn_expr]);
+            const name = offsets.tokenToSlice(tree, main_tokens[call.ast.fn_expr]);
 
-            const absolute_index = location.line_start + location.column;
-
-            const name = tree.tokenSlice(main_tokens[call.ast.fn_expr]);
-
-            if (try analysis.lookupSymbolGlobal(store, arena, handle, name, absolute_index)) |decl_handle| {
-                try writeCallHint(builder, arena, store, call, decl_handle);
+            if (try analysis.lookupSymbolGlobal(builder.store, builder.arena, handle, name, source_index)) |decl_handle| {
+                try writeCallHint(builder, call, decl_handle);
             }
         },
         .field_access => {
@@ -205,23 +191,23 @@ fn writeCallNodeHint(builder: *Builder, arena: *std.heap.ArenaAllocator, store:
             const start = offsets.tokenToIndex(tree, lhsToken);
             const rhs_loc = offsets.tokenToLoc(tree, rhsToken);
 
-            var held_range = try arena.allocator().dupeZ(u8, handle.text[start..rhs_loc.end]);
+            var held_range = try builder.arena.allocator().dupeZ(u8, handle.text[start..rhs_loc.end]);
             var tokenizer = std.zig.Tokenizer.init(held_range);
 
             // note: we have the ast node, traversing it would probably yield better results
             // than trying to re-tokenize and re-parse it
-            if (try analysis.getFieldAccessType(store, arena, handle, rhs_loc.end, &tokenizer)) |result| {
+            if (try analysis.getFieldAccessType(builder.store, builder.arena, handle, rhs_loc.end, &tokenizer)) |result| {
                 const container_handle = result.unwrapped orelse result.original;
                 switch (container_handle.type.data) {
                     .other => |container_handle_node| {
                         if (try analysis.lookupSymbolContainer(
-                            store,
-                            arena,
+                            builder.store,
+                            builder.arena,
                             .{ .node = container_handle_node, .handle = container_handle.handle },
                             tree.tokenSlice(rhsToken),
                             true,
                         )) |decl_handle| {
-                            try writeCallHint(builder, arena, store, call, decl_handle);
+                            try writeCallHint(builder, call, decl_handle);
                         }
                     },
                     else => {},
@@ -234,44 +220,18 @@ fn writeCallNodeHint(builder: *Builder, arena: *std.heap.ArenaAllocator, store:
     }
 }
 
-/// HACK self-hosted has not implemented async yet
-fn callWriteNodeInlayHint(allocator: std.mem.Allocator, args: anytype) error{OutOfMemory}!void {
-    if (zig_builtin.zig_backend == .other or zig_builtin.zig_backend == .stage1) {
-        const FrameSize = @sizeOf(@Frame(writeNodeInlayHint));
-        var child_frame = try allocator.alignedAlloc(u8, std.Target.stack_align, FrameSize);
-        defer allocator.free(child_frame);
-        return await @asyncCall(child_frame, {}, writeNodeInlayHint, args);
-    } else {
-        // TODO find a non recursive solution
-        return @call(.auto, writeNodeInlayHint, args);
-    }
-}
-
-/// iterates over the ast and writes parameter hints into `builder.hints` for every function call and builtin call
-/// nodes outside the given range are excluded
-fn writeNodeInlayHint(builder: *Builder, arena: *std.heap.ArenaAllocator, store: *DocumentStore, maybe_node: ?Ast.Node.Index, range: types.Range) error{OutOfMemory}!void {
-    const node = maybe_node orelse return;
-
+fn writeNodeInlayHint(
+    builder: *Builder,
+    node: Ast.Node.Index,
+) error{OutOfMemory}!void {
     const handle = builder.handle;
     const tree = handle.tree;
     const node_tags = tree.nodes.items(.tag);
-    const node_data = tree.nodes.items(.data);
     const main_tokens = tree.nodes.items(.main_token);
 
-    // std.log.info("max: {d} | curr: {d}", .{ node_data.len, node });
-    // if (node == 0 or node >= node_data.len) return;
-    if (node == 0) return;
-    // std.log.info("tag: {any}", .{node_tags[node]});
-    // std.log.info("src: {s}", .{tree.getNodeSource(node)});
-
-    var allocator = arena.allocator();
-
     const tag = node_tags[node];
 
-    // NOTE traversing the ast instead of iterating over all nodes allows using visibility
-    // checks based on the given range which reduce runtimes by orders of magnitude for large files
     switch (tag) {
-        .root => unreachable,
         .call_one,
         .call_one_comma,
         .async_call_one,
@@ -283,26 +243,19 @@ fn writeNodeInlayHint(builder: *Builder, arena: *std.heap.ArenaAllocator, store:
         => {
             var params: [1]Ast.Node.Index = undefined;
             const call = tree.fullCall(&params, node).?;
-            try writeCallNodeHint(builder, arena, store, call);
-
-            for (call.ast.params) |param| {
-                if (call.ast.params.len > inlay_hints_max_inline_children) {
-                    if (!isNodeInRange(tree, param, range)) continue;
-                }
-
-                try callWriteNodeInlayHint(allocator, .{ builder, arena, store, param, range });
-            }
+            try writeCallNodeHint(builder, call);
         },
 
         .builtin_call_two,
         .builtin_call_two_comma,
         .builtin_call,
         .builtin_call_comma,
-        => {
+        => blk: {
             var buffer: [2]Ast.Node.Index = undefined;
             const params = ast.builtinCallParams(tree, node, &buffer).?;
 
-            if (builder.config.inlay_hints_show_builtin and params.len > 1) {
+            if (!builder.config.inlay_hints_show_builtin or params.len <= 1) break :blk;
+
             const name = tree.tokenSlice(main_tokens[node]);
 
             outer: for (data.builtins) |builtin| {
@@ -314,375 +267,37 @@ fn writeNodeInlayHint(builder: *Builder, arena: *std.heap.ArenaAllocator, store:
 
                 try writeBuiltinHint(builder, params, builtin.arguments);
             }
-            }
-
-            for (params) |param| {
-                if (params.len > inlay_hints_max_inline_children) {
-                    if (!isNodeInRange(tree, param, range)) continue;
-                }
-
-                try callWriteNodeInlayHint(allocator, .{ builder, arena, store, param, range });
-            }
         },
-
-        .optional_type,
-        .array_type,
-        .@"continue",
-        .anyframe_type,
-        .anyframe_literal,
-        .char_literal,
-        .number_literal,
-        .unreachable_literal,
-        .identifier,
-        .enum_literal,
-        .string_literal,
-        .multiline_string_literal,
-        .error_set_decl,
-        => {},
-
-        .array_type_sentinel => {
-            const array_type = tree.arrayTypeSentinel(node);
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, array_type.ast.sentinel, range });
-        },
-
-        .ptr_type_aligned,
-        .ptr_type_sentinel,
-        .ptr_type,
-        .ptr_type_bit_range,
-        => {
-            const ptr_type: Ast.full.PtrType = ast.fullPtrType(tree, node).?;
-
-            if (ptr_type.ast.sentinel != 0) {
-                return try callWriteNodeInlayHint(allocator, .{ builder, arena, store, ptr_type.ast.sentinel, range });
-            }
-
-            if (ptr_type.ast.align_node != 0) {
-                try callWriteNodeInlayHint(allocator, .{ builder, arena, store, ptr_type.ast.align_node, range });
-
-                if (ptr_type.ast.bit_range_start != 0) {
-                    try callWriteNodeInlayHint(allocator, .{ builder, arena, store, ptr_type.ast.bit_range_start, range });
-                    try callWriteNodeInlayHint(allocator, .{ builder, arena, store, ptr_type.ast.bit_range_end, range });
-                }
-            }
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, ptr_type.ast.child_type, range });
-        },
-
-        .@"usingnamespace",
-        .field_access,
-        .unwrap_optional,
-        .bool_not,
-        .negation,
-        .bit_not,
-        .negation_wrap,
-        .address_of,
-        .@"try",
-        .@"await",
-        .deref,
-        .@"suspend",
-        .@"resume",
-        .@"return",
-        .grouped_expression,
-        .@"comptime",
-        .@"nosuspend",
-        => try callWriteNodeInlayHint(allocator, .{ builder, arena, store, node_data[node].lhs, range }),
-
-        .test_decl,
-        .global_var_decl,
-        .local_var_decl,
-        .simple_var_decl,
-        .aligned_var_decl,
-        .@"errdefer",
-        .@"defer",
-        .@"break",
-        => try callWriteNodeInlayHint(allocator, .{ builder, arena, store, node_data[node].rhs, range }),
-
-        .@"catch",
-        .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_access,
-        .switch_range,
-        .error_union,
-        => {
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, node_data[node].lhs, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, node_data[node].rhs, range });
-        },
-
-        .slice_open,
-        .slice,
-        .slice_sentinel,
-        => {
-            const slice: Ast.full.Slice = tree.fullSlice(node).?;
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, slice.ast.sliced, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, slice.ast.start, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, slice.ast.end, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, slice.ast.sentinel, range });
-        },
-
-        .array_init_one,
-        .array_init_one_comma,
-        .array_init_dot_two,
-        .array_init_dot_two_comma,
-        .array_init_dot,
-        .array_init_dot_comma,
-        .array_init,
-        .array_init_comma,
-        => {
-            var buffer: [2]Ast.Node.Index = undefined;
-            const array_init: Ast.full.ArrayInit = tree.fullArrayInit(&buffer, node).?;
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, array_init.ast.type_expr, range });
-            for (array_init.ast.elements) |elem| {
-                try callWriteNodeInlayHint(allocator, .{ builder, arena, store, elem, range });
-            }
-        },
-
-        .struct_init_one,
-        .struct_init_one_comma,
-        .struct_init_dot_two,
-        .struct_init_dot_two_comma,
-        .struct_init_dot,
-        .struct_init_dot_comma,
-        .struct_init,
-        .struct_init_comma,
-        => {
-            var buffer: [2]Ast.Node.Index = undefined;
-            const struct_init: Ast.full.StructInit = tree.fullStructInit(&buffer, node).?;
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, struct_init.ast.type_expr, range });
-
-            for (struct_init.ast.fields) |field_init| {
-                if (struct_init.ast.fields.len > inlay_hints_max_inline_children) {
-                    if (!isNodeInRange(tree, field_init, range)) continue;
-                }
-
-                try callWriteNodeInlayHint(allocator, .{ builder, arena, store, field_init, range });
-            }
-        },
-
-        .@"switch",
-        .switch_comma,
-        => {
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, node_data[node].lhs, range });
-
-            const extra = tree.extraData(node_data[node].rhs, Ast.Node.SubRange);
-            const cases = tree.extra_data[extra.start..extra.end];
-
-            for (cases) |case_node| {
-                if (cases.len > inlay_hints_max_inline_children) {
-                    if (!isNodeInRange(tree, case_node, range)) continue;
-                }
-
-                try callWriteNodeInlayHint(allocator, .{ builder, arena, store, case_node, range });
-            }
-        },
-
-        .switch_case_one,
-        .switch_case,
-        .switch_case_inline_one,
-        .switch_case_inline,
-        => {
-            const switch_case = tree.fullSwitchCase(node).?;
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, switch_case.ast.target_expr, range });
-        },
-
-        .while_simple,
-        .while_cont,
-        .@"while",
-        .for_simple,
-        .@"for",
-        => {
-            const while_node = ast.fullWhile(tree, node).?;
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, while_node.ast.cond_expr, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, while_node.ast.cont_expr, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, while_node.ast.then_expr, range });
-
-            if (while_node.ast.else_expr != 0) {
-                try callWriteNodeInlayHint(allocator, .{ builder, arena, store, while_node.ast.else_expr, range });
-            }
-        },
-
-        .if_simple,
-        .@"if",
-        => {
-            const if_node = ast.fullIf(tree, node).?;
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, if_node.ast.cond_expr, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, if_node.ast.then_expr, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, if_node.ast.else_expr, range });
-        },
-
-        .fn_proto_simple,
-        .fn_proto_multi,
-        .fn_proto_one,
-        .fn_proto,
-        .fn_decl,
-        => {
-            var buffer: [1]Ast.Node.Index = undefined;
-            const fn_proto: Ast.full.FnProto = tree.fullFnProto(&buffer, node).?;
-
-            var it = fn_proto.iterate(&tree);
-            while (ast.nextFnParam(&it)) |param_decl| {
-                try callWriteNodeInlayHint(allocator, .{ builder, arena, store, param_decl.type_expr, range });
-            }
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, fn_proto.ast.align_expr, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, fn_proto.ast.addrspace_expr, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, fn_proto.ast.section_expr, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, fn_proto.ast.callconv_expr, range });
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, fn_proto.ast.return_type, range });
-
-            if (tag == .fn_decl) {
-                try callWriteNodeInlayHint(allocator, .{ builder, arena, store, node_data[node].rhs, range });
-            }
-        },
-
-        .container_decl,
-        .container_decl_trailing,
-        .container_decl_two,
-        .container_decl_two_trailing,
-        .container_decl_arg,
-        .container_decl_arg_trailing,
-        .tagged_union,
-        .tagged_union_trailing,
-        .tagged_union_two,
-        .tagged_union_two_trailing,
-        .tagged_union_enum_tag,
-        .tagged_union_enum_tag_trailing,
-        => {
-            var buffer: [2]Ast.Node.Index = undefined;
-            const decl: Ast.full.ContainerDecl = tree.fullContainerDecl(&buffer, node).?;
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, decl.ast.arg, range });
-
-            for (decl.ast.members) |child| {
-                if (decl.ast.members.len > inlay_hints_max_inline_children) {
-                    if (!isNodeInRange(tree, child, range)) continue;
-                }
-
-                try callWriteNodeInlayHint(allocator, .{ builder, arena, store, child, range });
-            }
-        },
-
-        .container_field_init,
-        .container_field_align,
-        .container_field,
-        => {
-            const container_field = tree.fullContainerField(node).?;
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, container_field.ast.value_expr, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, container_field.ast.align_expr, range });
-        },
-
-        .block_two,
-        .block_two_semicolon,
-        => {
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, node_data[node].lhs, range });
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, node_data[node].rhs, range });
-        },
-
-        .block,
-        .block_semicolon,
-        => {
-            const subrange = tree.extra_data[node_data[node].lhs..node_data[node].rhs];
-
-            for (subrange) |child| {
-                if (subrange.len > inlay_hints_max_inline_children) {
-                    if (!isNodeInRange(tree, child, range)) continue;
-                }
-
-                try callWriteNodeInlayHint(allocator, .{ builder, arena, store, child, range });
-            }
-        },
-
-        .asm_simple,
-        .@"asm",
-        .asm_output,
-        .asm_input,
-        => {
-            const asm_node: Ast.full.Asm = tree.fullAsm(node) orelse return;
-
-            try callWriteNodeInlayHint(allocator, .{ builder, arena, store, asm_node.ast.template, range });
-        },
-
-        .error_value => {},
+        else => {},
     }
 }
 
 /// creates a list of `InlayHint`'s from the given document
 /// only parameter hints are created
-/// only hints in the given range are created
+/// only hints in the given loc are created
 pub fn writeRangeInlayHint(
     arena: *std.heap.ArenaAllocator,
     config: Config,
     store: *DocumentStore,
     handle: *const DocumentStore.Handle,
-    range: types.Range,
+    loc: offsets.Loc,
     hover_kind: types.MarkupKind,
-    encoding: offsets.Encoding,
-) error{OutOfMemory}![]types.InlayHint {
+) error{OutOfMemory}![]InlayHint {
     var builder: Builder = .{
-        .arena = arena.allocator(),
+        .arena = arena,
+        .store = store,
         .config = &config,
         .handle = handle,
         .hints = .{},
         .hover_kind = hover_kind,
-        .encoding = encoding,
     };
 
-    for (handle.tree.rootDecls()) |child| {
-        if (!isNodeInRange(handle.tree, child, range)) continue;
-        try writeNodeInlayHint(&builder, arena, store, child, range);
+    const nodes = try ast.nodesAtLoc(arena.allocator(), handle.tree, loc);
+
+    for (nodes) |child| {
+        try writeNodeInlayHint(&builder, child);
+        try ast.iterateChildrenRecursive(handle.tree, child, &builder, error{OutOfMemory}, writeNodeInlayHint);
     }
 
-    return builder.toOwnedSlice();
+    return try builder.toOwnedSlice();
 }

src/offsets.zig

@@ -240,6 +240,29 @@ pub fn convertRangeEncoding(text: []const u8, range: types.Range, from_encoding:
     };
 }
 
+// returns true if a and b intersect
+pub fn locIntersect(a: Loc, b: Loc) bool {
+    std.debug.assert(a.start <= a.end and b.start <= b.end);
+    const a_start_in_b = b.start <= a.start and a.start <= b.end;
+    const a_end_in_b = b.start <= a.end and a.end <= b.end;
+    return a_start_in_b or a_end_in_b;
+}
+
+// returns true if a is inside b
+pub fn locInside(inner: Loc, outer: Loc) bool {
+    std.debug.assert(inner.start <= inner.end and outer.start <= outer.end);
+    return outer.start <= inner.start and inner.end <= outer.end;
+}
+
+// returns the union of a and b
+pub fn locMerge(a: Loc, b: Loc) Loc {
+    std.debug.assert(a.start <= a.end and b.start <= b.end);
+    return .{
+        .start = @min(a.start, b.start),
+        .end = @max(a.end, b.end),
+    };
+}
+
 // Helper functions
 
 /// advance `position` which starts at `from_index` to `to_index` accounting for line breaks

src/zls.zig

@@ -1,6 +1,7 @@
 // Used by tests as a package, can be used by tools such as
 // zigbot9001 to take advantage of zls' tools
 
+pub const ast = @import("ast.zig");
 pub const analysis = @import("analysis.zig");
 pub const Header = @import("Header.zig");
 pub const debug = @import("debug.zig");

tests/tests.zig

@@ -1,6 +1,7 @@
 comptime {
     _ = @import("helper.zig");
 
+    _ = @import("utility/ast.zig");
     _ = @import("utility/offsets.zig");
     _ = @import("utility/position_context.zig");
     _ = @import("utility/uri.zig");

tests/utility/ast.zig

@@ -0,0 +1,76 @@
+const std = @import("std");
+const zls = @import("zls");
+
+const helper = @import("../helper.zig");
+const Context = @import("../context.zig").Context;
+const ErrorBuilder = @import("../ErrorBuilder.zig");
+
+const types = zls.types;
+const offsets = zls.offsets;
+const ast = zls.ast;
+
+const allocator = std.testing.allocator;
+
+test "nodesAtLoc" {
+    try testNodesAtLoc(
+        \\<outer>const<inner> foo<inner> = 5<outer>;
+    );
+    try testNodesAtLoc(
+        \\<outer>const f<inner>oo = 5;
+        \\var bar = <inner>2<outer>;
+    );
+    try testNodesAtLoc(
+        \\const foo = <outer>5<inner> +<inner> 2<outer>;
+    );
+    try testNodesAtLoc(
+        \\<outer><inner>fn foo(alpha: u32) void {}
+        \\const _ = foo(5);<inner><outer>
+    );
+}
+
+fn testNodesAtLoc(source: []const u8) !void {
+    var ccp = try helper.collectClearPlaceholders(allocator, source);
+    defer ccp.deinit(allocator);
+
+    const old_locs = ccp.locations.items(.old);
+    const locs = ccp.locations.items(.new);
+
+    std.debug.assert(ccp.locations.len == 4);
+    std.debug.assert(std.mem.eql(u8, offsets.locToSlice(source, old_locs[0]), "<outer>"));
+    std.debug.assert(std.mem.eql(u8, offsets.locToSlice(source, old_locs[1]), "<inner>"));
+    std.debug.assert(std.mem.eql(u8, offsets.locToSlice(source, old_locs[2]), "<inner>"));
+    std.debug.assert(std.mem.eql(u8, offsets.locToSlice(source, old_locs[3]), "<outer>"));
+
+    const inner_loc = offsets.Loc{ .start = locs[1].start, .end = locs[2].start };
+    const outer_loc = offsets.Loc{ .start = locs[0].start, .end = locs[3].end };
+
+    const new_source = try allocator.dupeZ(u8, ccp.new_source);
+    defer allocator.free(new_source);
+
+    var tree = try std.zig.parse(allocator, new_source);
+    defer tree.deinit(allocator);
+
+    const nodes = try ast.nodesAtLoc(allocator, tree, inner_loc);
+    defer allocator.free(nodes);
+
+    const actual_loc = offsets.Loc{
+        .start = offsets.nodeToLoc(tree, nodes[0]).start,
+        .end = offsets.nodeToLoc(tree, nodes[nodes.len - 1]).end,
+    };
+
+    var error_builder = ErrorBuilder.init(allocator, new_source);
+    defer error_builder.deinit();
+    errdefer error_builder.writeDebug();
+
+    if (outer_loc.start != actual_loc.start) {
+        try error_builder.msgAtIndex("actual start here", actual_loc.start, .err, .{});
+        try error_builder.msgAtIndex("expected start here", outer_loc.start, .err, .{});
+        return error.LocStartMismatch;
+    }
+
+    if (outer_loc.end != actual_loc.end) {
+        try error_builder.msgAtIndex("actual end here", actual_loc.end, .err, .{});
+        try error_builder.msgAtIndex("expected end here", outer_loc.end, .err, .{});
+        return error.LocEndMismatch;
+    }
+}