replace ArrayLists with bounded storage

1 files changed, 171 insertions, 56 deletions

diff --git a/src/lsp/lsp_zls.zig b/src/lsp/lsp_zls.zig
index e1624df0..76c5594c 100644
--- a/src/lsp/lsp_zls.zig
+++ b/src/lsp/lsp_zls.zig
@@ -60,6 +60,9 @@ const enc: offsets.Encoding = .@"utf-8";
 /// unbounded tree walk because someone opened a file in `/`.
 const max_files = 512;
 const max_rows = 2000;
+const max_output_bytes = max_rows * 4096;
+const max_trace_bytes = 32 * 1024;
+const max_code_action_bytes = max_rows * 1024;
 
 pub const supports: std.EnumSet(lsp.Kind) = .initMany(&.{
     .definition,
@@ -171,12 +174,18 @@ fn nowUs() u64 {
 /// find out what actually happened.
 const Trace = struct {
     on: bool = false,
-    arena: std.mem.Allocator,
-    buf: std.ArrayList(u8) = .empty,
+    buf: [max_trace_bytes]u8 = undefined,
+    len: usize = 0,
 
     fn note(t: *Trace, comptime fmt: []const u8, args: anytype) void {
-        if (!t.on) return;
-        t.buf.print(t.arena, fmt ++ "\n", args) catch {};
+        if (!t.on or t.len == t.buf.len) return;
+        var w: std.Io.Writer = .fixed(t.buf[t.len..]);
+        w.print(fmt ++ "\n", args) catch {};
+        t.len += w.buffered().len;
+    }
+
+    fn written(t: *const Trace) []const u8 {
+        return t.buf[0..t.len];
     }
 };
 
@@ -189,33 +198,34 @@ pub fn query(gpa: std.mem.Allocator, arena: std.mem.Allocator, req: lsp.Req, out
     // 24-entry ring every time you looked at it.
     if (req.kind == .status) return status(arena, req, out) catch {};
 
-    // Rows land in a scratch buffer first, for two reasons a plain Writer
-    // cannot serve: the log wants an exact row count, and `explain` throws the
-    // rows away and prints the narration in their place — neither a count nor
-    // a rewind exists on a Writer you do not own.
-    var scratch: std.Io.Writer.Allocating = .init(gpa);
-    defer scratch.deinit();
+    // Rows land in bounded scratch storage first, for two reasons a plain
+    // Writer cannot serve: the log wants an exact row count, and `explain`
+    // throws the rows away and prints the narration in their place. A query
+    // that exceeds the row budget's byte allowance keeps its ordered prefix.
+    const scratch_buf = gpa.alloc(u8, max_output_bytes) catch return;
+    defer gpa.free(scratch_buf);
+    var scratch: std.Io.Writer = .fixed(scratch_buf);
 
-    var tr: Trace = .{ .arena = arena, .on = req.kind == .explain };
+    var tr: Trace = .{ .on = req.kind == .explain };
     const t0 = nowUs();
     var err_name: []const u8 = "";
-    run(gpa, arena, req, &scratch.writer, &tr) catch |e| {
+    run(gpa, arena, req, &scratch, &tr) catch |e| {
         err_name = @errorName(e);
         tr.note("ERROR: {s} — the query threw; the editor shows this as 'no result'", .{err_name});
     };
     const us = nowUs() -| t0;
-    const rows = std.mem.count(u8, scratch.written(), "\n");
+    const rows = std.mem.count(u8, scratch.buffered(), "\n");
     record(req, us, rows, err_name);
 
     if (req.kind != .explain) {
-        out.writeAll(scratch.written()) catch {};
+        out.writeAll(scratch.buffered()) catch {};
         return;
     }
     // the question was never "where is it", it was "what did you do"
     out.print("lsp explain — the definition query at byte {d} of {s}\n\n", .{
         req.offset, if (req.path.len == 0) "(no file)" else std.fs.path.basename(req.path),
     }) catch {};
-    out.writeAll(tr.buf.items) catch {};
+    out.writeAll(tr.written()) catch {};
     if (hideTime())
         out.print("\n{d} row(s)\n", .{rows}) catch {}
     else
@@ -470,10 +480,16 @@ const Lines = struct {
     src: []const u8,
 
     fn build(arena: std.mem.Allocator, src: []const u8) !Lines {
-        var starts: std.ArrayList(u32) = .empty;
-        try starts.append(arena, 0);
-        for (src, 0..) |c, i| if (c == '\n') try starts.append(arena, @intCast(i + 1));
-        return .{ .starts = starts.items, .src = src };
+        const count = std.mem.countScalar(u8, src, '\n') + 1;
+        const starts = try arena.alloc(u32, count);
+        starts[0] = 0;
+        var n: usize = 1;
+        for (src, 0..) |c, i| {
+            if (c != '\n') continue;
+            starts[n] = @intCast(i + 1);
+            n += 1;
+        }
+        return .{ .starts = starts, .src = src };
     }
 
     fn line(l: Lines, n: usize) []const u8 {
@@ -529,7 +545,8 @@ fn goto(
     const ctx = try Analyser.getPositionContext(arena, tree, off, true);
     tr.note("position context: .{s}", .{@tagName(ctx)});
 
-    var decls: std.ArrayList(Analyser.DeclWithHandle) = .empty;
+    const decls = try arena.alloc(Analyser.DeclWithHandle, max_rows);
+    var decl_count: usize = 0;
     switch (ctx) {
         .var_access, .test_doctest_name => {
             const name_loc = offsets.identifierLocFromIndex(tree, off) orelse {
@@ -550,7 +567,8 @@ fn goto(
                 tr.note("STOP: lookupSymbolGlobal found no `{s}` in scope here", .{name});
                 return;
             };
-            try decls.append(arena, d);
+            decls[decl_count] = d;
+            decl_count += 1;
         },
         .field_access => |loc| {
             const name_loc = offsets.identifierLocFromIndex(tree, off) orelse {
@@ -576,19 +594,22 @@ fn goto(
                 explainFieldMiss(tr, offsets.locToSlice(tree.source, loc));
                 return;
             }
-            try decls.appendSlice(arena, found);
+            decl_count = @min(found.len, decls.len);
+            @memcpy(decls[0..decl_count], found[0..decl_count]);
         },
         .label_access, .label_decl => {
             const name_loc = offsets.identifierLocFromIndex(tree, off) orelse return;
             const name = offsets.locToSlice(tree.source, name_loc);
             const d = try Analyser.lookupLabel(handle, name, off) orelse return;
-            try decls.append(arena, d);
+            decls[decl_count] = d;
+            decl_count += 1;
         },
         .enum_literal => {
             const name_loc = offsets.identifierLocFromIndex(tree, off) orelse return;
             const name = offsets.locToSlice(tree.source, name_loc);
             const d = try analyser.getSymbolEnumLiteral(handle, off, name) orelse return;
-            try decls.append(arena, d);
+            decls[decl_count] = d;
+            decl_count += 1;
         },
         // `gd` on `@import("foo.zig")` opens foo.zig. Only relative imports
         // resolve — a package name needs the build graph we do not run.
@@ -626,10 +647,10 @@ fn goto(
         },
     }
 
-    tr.note("resolved to {d} declaration(s); asking each for its {s} token", .{ decls.items.len, @tagName(kind) });
+    tr.note("resolved to {d} declaration(s); asking each for its {s} token", .{ decl_count, @tagName(kind) });
     // One decl that fails to resolve must not swallow the others: a field
     // access can name several and a partial answer beats none.
-    for (decls.items) |d| {
+    for (decls[0..decl_count]) |d| {
         switch (kind) {
             .declaration => try rowForToken(arena, base, d.definitionToken(analyser, false) catch continue, out),
             .definition => try rowForToken(arena, base, d.definitionToken(analyser, true) catch continue, out),
@@ -730,6 +751,86 @@ fn completionSource(arena: std.mem.Allocator, req: lsp.Req, tr: *Trace) [:0]cons
     return best;
 }
 
+/// ZLS's collector takes a growable list. This equivalent writes directly into
+/// the query's row-bounded slice and stops in declaration order when it fills.
+fn collectContainerDecls(
+    analyser: *Analyser,
+    container_type: Analyser.Type,
+    original_handle: *DocumentStore.Handle,
+    instance_access: bool,
+    out: []Analyser.DeclWithHandle,
+    len: *usize,
+) !void {
+    const info = switch (container_type.data) {
+        .container => |info| info,
+        .either => |entries| {
+            for (entries) |entry| {
+                if (len.* == out.len) return;
+                const ty: Analyser.Type = .{
+                    .data = entry.type_data,
+                    .is_type_val = container_type.is_type_val,
+                };
+                try collectContainerDecls(analyser, ty, original_handle, instance_access, out, len);
+            }
+            return;
+        },
+        else => return,
+    };
+    const container_scope = info.scope_handle;
+    const handle = container_scope.handle;
+    const tree = &handle.tree;
+    const document_scope = try handle.getDocumentScope();
+    const main_token = tree.nodeMainToken(container_scope.toNode());
+    const is_enum = tree.tokenTag(main_token) == .keyword_enum;
+
+    for (document_scope.getScopeDeclarationsConst(container_scope.scope)) |decl_index| {
+        if (len.* == out.len) return;
+        const decl = document_scope.declarations.get(@intFromEnum(decl_index));
+        const candidate: Analyser.DeclWithHandle = .{
+            .decl = decl,
+            .handle = handle,
+            .container_type = container_type,
+        };
+        if (handle != original_handle and !candidate.isPublic()) continue;
+
+        switch (decl) {
+            .ast_node => |node| switch (tree.nodeTag(node)) {
+                .container_field_init, .container_field_align, .container_field => {
+                    if (is_enum) {
+                        if (instance_access) continue;
+                        const field_name = offsets.tokenToSlice(tree, tree.nodeMainToken(node));
+                        if (std.mem.eql(u8, field_name, "_")) continue;
+                    } else if (!instance_access) continue;
+                },
+                .fn_proto,
+                .fn_proto_multi,
+                .fn_proto_one,
+                .fn_proto_simple,
+                .fn_decl,
+                .global_var_decl,
+                .local_var_decl,
+                .simple_var_decl,
+                .aligned_var_decl,
+                => {
+                    if (instance_access) {
+                        const alias_type = try candidate.resolveType(analyser) orelse continue;
+                        const func_ty = try analyser.resolveFuncProtoOfCallable(alias_type) orelse continue;
+                        if (!analyser.firstParamIs(func_ty, .{
+                            .data = .{ .container = info },
+                            .is_type_val = true,
+                        })) continue;
+                    }
+                },
+                else => unreachable,
+            },
+            .label => continue,
+            else => {},
+        }
+        out[len.*] = candidate;
+        len.* += 1;
+    }
+}
+
 /// Tab after a `.`. NOT an autocomplete popup: the seam answers with
 /// locations, so the question it asks is "what could go here, and where is
 /// each of those DEFINED" — one `+Search` row per candidate, pointing at its
@@ -796,7 +897,7 @@ fn completion(
             try analyser.resolveDerefType(ty) orelse
             try analyser.resolveOptionalUnwrap(ty) orelse break;
     }
-    // `instance_access` is collectDeclarationsOfContainer's one knob and it
+    // `instance_access` is the collector's one knob and it
     // means different things per container kind: it is what hides an enum's
     // members and what reveals a struct's fields. So: a field access asks
     // whichever side the left-hand expression already was; a field named
@@ -806,12 +907,13 @@ fn completion(
     // whose fields are also its tags, so it wants the instance side wherever
     // it appears.
     const value_pos = ctx == .enum_literal and !in_init;
-    var decls: std.ArrayList(Analyser.DeclWithHandle) = .empty;
-    try analyser.collectDeclarationsOfContainer(ty, handle, switch (ctx) {
+    const decls = try arena.alloc(Analyser.DeclWithHandle, max_rows);
+    var decl_count: usize = 0;
+    try collectContainerDecls(analyser, ty, handle, switch (ctx) {
         .field_access => !ty.is_type_val,
         else => in_init or ty.isUnionType(),
-    }, &decls);
-    tr.note("{d} candidate(s) in scope", .{decls.items.len});
+    }, decls, &decl_count);
+    tr.note("{d} candidate(s) in scope", .{decl_count});
     // What completionSource spliced in at `off`, which the ROWS have to take
     // back out. The tree being read is the REPAIRED copy, so for a candidate
     // declared in this same file two things are wrong with the naive row: its
@@ -824,7 +926,7 @@ fn completion(
     const pad = tree.source.len - orig.len;
     const dot = lsp.lineCol(orig, off);
     var n: usize = 0;
-    for (decls.items) |d| {
+    for (decls[0..decl_count]) |d| {
         if (n >= max_rows) return;
         // In a value position the answer is the container's MEMBERS plus its
         // DECL LITERALS (`.empty`, `.init`) — never a method, an unrelated
@@ -959,10 +1061,9 @@ fn workspaceSymbols(
     out: *std.Io.Writer,
 ) !void {
     if (req.arg.len == 0) return;
-    var files: std.ArrayList([]const u8) = .empty;
-    try collectZigFiles(arena, io, req.root, &files);
+    const files = try collectZigFiles(arena, io, req.root);
     var n: usize = 0;
-    for (files.items) |path| {
+    for (files) |path| {
         if (n >= max_rows) return;
         const src = readFileZ(arena, io, path) catch continue;
         var tree: Ast = Ast.parse(arena, src, .zig) catch continue;
@@ -1147,14 +1248,18 @@ fn treeDiagnostics(
         const lines: Lines = try .build(arena, tree.source);
         for (tree.errors) |e| {
             if (n >= max_rows) return n;
-            var buf: std.ArrayList(u8) = .empty;
-            var w: std.Io.Writer.Allocating = .fromArrayList(arena, &buf);
-            tree.renderError(e, &w.writer) catch continue;
+            var count_buf: [256]u8 = undefined;
+            var counting: std.Io.Writer.Discarding = .init(&count_buf);
+            tree.renderError(e, &counting.writer) catch continue;
+            const size = std.math.cast(usize, counting.fullCount()) orelse continue;
+            const rendered = try arena.alloc(u8, size);
+            var w: std.Io.Writer = .fixed(rendered);
+            tree.renderError(e, &w) catch continue;
             const at = tree.tokenStart(e.token);
             const lc = lsp.lineCol(tree.source, at);
             n += 1;
             lsp.row(out, path, lc.line, lc.col, try std.fmt.allocPrint(arena, "error: {s}   {s}", .{
-                w.written(), std.mem.trim(u8, lines.line(lc.line), " \t"),
+                w.buffered(), std.mem.trim(u8, lines.line(lc.line), " \t"),
             }));
         }
         return n;
@@ -1187,17 +1292,16 @@ fn workspaceDiagnostics(
     req: lsp.Req,
     out: *std.Io.Writer,
 ) !void {
-    var files: std.ArrayList([]const u8) = .empty;
-    try collectZigFiles(arena, io, req.root, &files);
+    const files = try collectZigFiles(arena, io, req.root);
     var total: usize = 0;
-    for (files.items) |path| {
+    for (files) |path| {
         const src = readFileZ(arena, io, path) catch continue;
         var tree: Ast = Ast.parse(arena, src, .zig) catch continue;
         defer tree.deinit(arena);
         total += treeDiagnostics(gpa, arena, lsp.rel(req.root, path), &tree, out) catch continue;
         if (total >= max_rows) return;
     }
-    if (total == 0) lsp.row(out, lsp.rel(req.root, req.path), 0, 0, try std.fmt.allocPrint(arena, "no diagnostics in {d} file(s)", .{files.items.len}));
+    if (total == 0) lsp.row(out, lsp.rel(req.root, req.path), 0, 0, try std.fmt.allocPrint(arena, "no diagnostics in {d} file(s)", .{files.len}));
 }
 
 // ----------------------------------------------------------- code actions
@@ -1225,8 +1329,12 @@ fn codeActions(
     var bundle = try astCheck(gpa, "", tree);
     defer bundle.deinit(gpa);
 
+    // The external builder owns growable collections internally. Confine all
+    // of them and their payloads to one preallocated query-local region.
+    const action_storage = try arena.alloc(u8, max_code_action_bytes);
+    var action_fba: std.heap.FixedBufferAllocator = .init(action_storage);
     var builder: zls.code_actions.Builder = .{
-        .arena = arena,
+        .arena = action_fba.allocator(),
         .analyser = analyser,
         .handle = handle,
         .offset_encoding = enc,
@@ -1236,7 +1344,7 @@ fn codeActions(
     const at = offsets.locToRange(tree.source, .{ .start = off, .end = off }, enc);
     builder.generateCodeActionsInRange(at) catch {};
 
-    for (builder.actions.items) |a| {
+    for (builder.actions.items[0..@min(builder.actions.items.len, max_rows)]) |a| {
         out.print("{s}\n", .{a.title}) catch return;
     }
     if (builder.actions.items.len == 0) out.print("no code actions\n", .{}) catch {};
@@ -1257,10 +1365,14 @@ fn formatQuery(arena: std.mem.Allocator, req: lsp.Req, out: *std.Io.Writer) !voi
         lsp.row(out, path, 0, 0, "cannot format: file does not parse");
         return;
     }
-    var buf: std.ArrayList(u8) = .empty;
-    var w: std.Io.Writer.Allocating = .fromArrayList(arena, &buf);
-    tree.render(arena, &w.writer, .{}) catch return;
-    const formatted = w.written();
+    var count_buf: [4096]u8 = undefined;
+    var counting: std.Io.Writer.Discarding = .init(&count_buf);
+    tree.render(arena, &counting.writer, .{}) catch return;
+    const size = std.math.cast(usize, counting.fullCount()) orelse return;
+    const render_buf = try arena.alloc(u8, size);
+    var w: std.Io.Writer = .fixed(render_buf);
+    tree.render(arena, &w, .{}) catch return;
+    const formatted = w.buffered();
     if (std.mem.eql(u8, formatted, req.source)) {
         lsp.row(out, path, 0, 0, "already formatted");
         return;
@@ -1287,13 +1399,15 @@ fn formatQuery(arena: std.mem.Allocator, req: lsp.Req, out: *std.Io.Writer) !voi
 
 /// Every `.zig` under `root`, skipping the directories a source walk has no
 /// business in. Capped, because this runs on a keypress.
-fn collectZigFiles(arena: std.mem.Allocator, io: std.Io, root: []const u8, out: *std.ArrayList([]const u8)) !void {
-    if (root.len == 0) return;
-    var dir = std.Io.Dir.cwd().openDir(io, root, .{ .iterate = true }) catch return;
+fn collectZigFiles(arena: std.mem.Allocator, io: std.Io, root: []const u8) ![]const []const u8 {
+    const files = try arena.alloc([]const u8, max_files);
+    var len: usize = 0;
+    if (root.len == 0) return files[0..0];
+    var dir = std.Io.Dir.cwd().openDir(io, root, .{ .iterate = true }) catch return files[0..0];
     defer dir.close(io);
-    var w = dir.walkSelectively(arena) catch return;
+    var w = dir.walkSelectively(arena) catch return files[0..0];
     defer w.deinit();
-    while (out.items.len < max_files) {
+    while (len < files.len) {
         const e = (w.next(io) catch break) orelse break;
         switch (e.kind) {
             .directory => {
@@ -1304,12 +1418,13 @@ fn collectZigFiles(arena: std.mem.Allocator, io: std.Io, root: []const u8, out:
             },
             .file => {
                 if (!std.mem.endsWith(u8, e.basename, ".zig")) continue;
-                const p = std.fs.path.join(arena, &.{ root, e.path }) catch continue;
-                try out.append(arena, p);
+                files[len] = std.fs.path.join(arena, &.{ root, e.path }) catch continue;
+                len += 1;
             },
             else => {},
         }
     }
+    return files[0..len];
 }
 
 fn readFileZ(arena: std.mem.Allocator, io: std.Io, path: []const u8) ![:0]u8 {