Reuse bounded source analysis and speed up result context traversal

1 files changed, 213 insertions, 0 deletions

diff --git a/src/locations_cache.zig b/src/locations_cache.zig
new file mode 100644
index 00000000..a079226b
--- /dev/null
+++ b/src/locations_cache.zig
@@ -0,0 +1,213 @@
+const std = @import("std");
+const syntax = @import("syntax.zig");
+
+/// A cache hit borrows its analysis; an entry too large to retain transfers
+/// ownership instead. Call deinit after consuming either result.
+pub const Result = struct {
+    analysis: syntax.SourceAnalysis,
+    owned: bool = false,
+
+    pub fn deinit(result: *Result, gpa: std.mem.Allocator) void {
+        if (result.owned) result.analysis.deinit(gpa);
+        result.* = undefined;
+    }
+};
+
+/// Owned source snapshots keep analysis valid across filesystem changes and
+/// borrowed editor buffers. Entries are ordered from least to most recent.
+pub const Cache = struct {
+    pub const max_bytes = 64 * 1024 * 1024;
+    pub const max_entries = 64;
+
+    const Entry = struct {
+        path: []u8,
+        source: []u8,
+        analysis: syntax.SourceAnalysis,
+        declarations_ready: bool,
+        colors_ready: bool,
+
+        fn size(entry: Entry) usize {
+            return entry.path.len + entry.source.len + entry.analysis.colors.len +
+                entry.analysis.declarations.len * @sizeOf(syntax.ContextDeclaration);
+        }
+
+        fn deinit(entry: *Entry, gpa: std.mem.Allocator) void {
+            gpa.free(entry.path);
+            gpa.free(entry.source);
+            entry.analysis.deinit(gpa);
+        }
+    };
+
+    entries: [max_entries]Entry = undefined,
+    len: usize = 0,
+    bytes: usize = 0,
+    analyses: usize = 0,
+    hits: usize = 0,
+
+    pub fn deinit(cache: *Cache, gpa: std.mem.Allocator) void {
+        for (cache.entries[0..cache.len]) |*entry| entry.deinit(gpa);
+        cache.* = .{};
+    }
+
+    /// Borrowed slices remain valid until the next get or cache deinit; an
+    /// owned result remains valid until Result.deinit. All calls for this
+    /// cache must use the same allocator. Unsupported languages
+    /// are cached too, including their empty analysis arrays.
+    pub fn get(cache: *Cache, gpa: std.mem.Allocator, path: []const u8, source: []const u8, want_declarations: bool, want_colors: bool) !Result {
+        return cache.getBounded(gpa, path, source, want_declarations, want_colors, max_bytes);
+    }
+
+    fn getBounded(cache: *Cache, gpa: std.mem.Allocator, path: []const u8, source: []const u8, want_declarations: bool, want_colors: bool, budget: usize) !Result {
+        if (source.len > budget or path.len > budget - source.len) return error.SourceTooLarge;
+        var previous: ?usize = null;
+        var declarations = want_declarations;
+        var colors = want_colors;
+        for (cache.entries[0..cache.len], 0..) |entry, index| {
+            if (!std.mem.eql(u8, entry.path, path)) continue;
+            previous = index;
+            if (!std.mem.eql(u8, entry.source, source)) break;
+            if ((!want_declarations or entry.declarations_ready) and (!want_colors or entry.colors_ready)) {
+                // Moving the small ownership record keeps the underlying
+                // allocations intact and avoids timestamp/overflow state.
+                std.mem.copyForwards(Entry, cache.entries[index .. cache.len - 1], cache.entries[index + 1 .. cache.len]);
+                cache.entries[cache.len - 1] = entry;
+                cache.hits +|= 1;
+                return .{ .analysis = entry.analysis };
+            }
+            declarations = declarations or entry.declarations_ready;
+            colors = colors or entry.colors_ready;
+            break;
+        }
+        // A supported full-source color map needs one byte per source byte.
+        // Reject that known minimum before parsing, so the caller's uncached
+        // path does not repeat an expensive analysis for oversized inputs.
+        const input_bytes = path.len + source.len;
+        if (colors and syntax.supportsPath(path) and source.len > budget - input_bytes) return error.SourceTooLarge;
+        var analysis = try syntax.analyzeSource(gpa, path, source, declarations, colors);
+        errdefer analysis.deinit(gpa);
+        cache.analyses +|= 1;
+        const analysis_bytes = analysis.colors.len + analysis.declarations.len * @sizeOf(syntax.ContextDeclaration);
+        if (analysis_bytes > budget - input_bytes) return .{ .analysis = analysis, .owned = true };
+        const owned_path = try gpa.dupe(u8, path);
+        errdefer gpa.free(owned_path);
+        const owned_source = try gpa.dupe(u8, source);
+        errdefer gpa.free(owned_source);
+        const next: Entry = .{
+            .path = owned_path,
+            .source = owned_source,
+            .analysis = analysis,
+            .declarations_ready = declarations,
+            .colors_ready = colors,
+        };
+        // Commit only after every allocation succeeds. A failed refresh leaves
+        // the previous cached snapshot usable by the next request.
+        if (previous) |index| cache.remove(gpa, index);
+        while (cache.len == max_entries or cache.bytes > budget - next.size()) cache.remove(gpa, 0);
+        cache.entries[cache.len] = next;
+        cache.len += 1;
+        cache.bytes += next.size();
+        return .{ .analysis = next.analysis };
+    }
+
+    fn remove(cache: *Cache, gpa: std.mem.Allocator, index: usize) void {
+        cache.bytes -= cache.entries[index].size();
+        cache.entries[index].deinit(gpa);
+        std.mem.copyForwards(Entry, cache.entries[index .. cache.len - 1], cache.entries[index + 1 .. cache.len]);
+        cache.len -= 1;
+    }
+};
+
+test "locations cache validates exact source and path and upgrades analysis" {
+    const gpa = std.testing.allocator;
+    syntax.start(gpa);
+    defer syntax.stop();
+    var cache: Cache = .{};
+    defer cache.deinit(gpa);
+    const source = "const Thing = struct {\n    value: u32,\n};\n";
+    _ = try cache.get(gpa, "a.zig", source, true, false);
+    try std.testing.expectEqual(@as(usize, 1), cache.analyses);
+    _ = try cache.get(gpa, "a.zig", source, true, false);
+    try std.testing.expectEqual(@as(usize, 1), cache.analyses);
+    try std.testing.expectEqual(@as(usize, 1), cache.hits);
+    _ = try cache.get(gpa, "a.zig", source, false, true);
+    try std.testing.expectEqual(@as(usize, 2), cache.analyses);
+    try std.testing.expect(cache.entries[0].declarations_ready and cache.entries[0].colors_ready);
+    _ = try cache.get(gpa, "a.zig", source, true, true);
+    try std.testing.expectEqual(@as(usize, 2), cache.analyses);
+    _ = try cache.get(gpa, "a.zig", "const Thing = struct {\n    other: u32,\n};\n", true, true);
+    try std.testing.expectEqual(@as(usize, 3), cache.analyses);
+    try std.testing.expectEqual(@as(usize, 1), cache.len);
+    _ = try cache.get(gpa, "a.txt", source, true, true);
+    try std.testing.expectEqual(@as(usize, 4), cache.analyses);
+    try std.testing.expectEqual(@as(usize, 2), cache.len);
+}
+
+test "locations cache bounds memory and entries and evicts least recently used" {
+    const gpa = std.testing.allocator;
+    var cache: Cache = .{};
+    defer cache.deinit(gpa);
+    for (0..Cache.max_entries) |index| {
+        var path: [32]u8 = undefined;
+        _ = try cache.get(gpa, try std.fmt.bufPrint(&path, "{d}.unknown", .{index}), "source", false, false);
+    }
+    _ = try cache.get(gpa, "0.unknown", "source", false, false);
+    _ = try cache.get(gpa, "new.unknown", "source", false, false);
+    try std.testing.expectEqual(Cache.max_entries, cache.len);
+    try std.testing.expectEqualStrings("2.unknown", cache.entries[0].path);
+    const count = cache.analyses;
+    _ = try cache.get(gpa, "0.unknown", "source", false, false);
+    try std.testing.expectEqual(count, cache.analyses);
+    cache.deinit(gpa);
+    _ = try cache.getBounded(gpa, "a.unknown", "source", false, false, 40);
+    _ = try cache.getBounded(gpa, "b.unknown", "source", false, false, 40);
+    _ = try cache.getBounded(gpa, "c.unknown", "source", false, false, 40);
+    try std.testing.expectEqual(@as(usize, 2), cache.len);
+    try std.testing.expect(cache.bytes <= 40);
+    try std.testing.expectEqualStrings("b.unknown", cache.entries[0].path);
+    try std.testing.expectError(error.SourceTooLarge, cache.getBounded(gpa, "large.unknown", "x" ** 41, false, false, 40));
+    try std.testing.expectEqual(@as(usize, 2), cache.len);
+    if (syntax.supportsPath("large.zig")) {
+        const before = cache.analyses;
+        try std.testing.expectError(error.SourceTooLarge, cache.getBounded(gpa, "large.zig", "x" ** 25, true, true, 40));
+        try std.testing.expectEqual(before, cache.analyses);
+    }
+    cache.deinit(gpa);
+    try std.testing.expectEqual(@as(usize, 0), cache.bytes);
+    try std.testing.expectEqual(@as(usize, 0), cache.len);
+}
+
+test "locations cache keeps its previous snapshot after allocation failure" {
+    var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{});
+    const gpa = failing.allocator();
+    var cache: Cache = .{};
+    defer cache.deinit(gpa);
+    _ = try cache.get(gpa, "a.unknown", "old bytes", false, false);
+    const retained = cache.bytes;
+    failing.fail_index = failing.alloc_index;
+    try std.testing.expectError(error.OutOfMemory, cache.get(gpa, "a.unknown", "new bytes", false, false));
+    try std.testing.expectEqual(@as(usize, 1), cache.len);
+    try std.testing.expectEqual(retained, cache.bytes);
+    try std.testing.expectEqualStrings("old bytes", cache.entries[0].source);
+    // This succeeds with allocations still disabled because the old entry
+    // remains valid and the hit path moves ownership records only.
+    _ = try cache.get(gpa, "a.unknown", "old bytes", false, false);
+    try std.testing.expectEqual(@as(usize, 1), cache.hits);
+}
+
+test "locations cache transfers oversized completed analysis without reparsing" {
+    if (!syntax.supportsPath("large.zig")) return error.SkipZigTest;
+    const gpa = std.testing.allocator;
+    syntax.start(gpa);
+    defer syntax.stop();
+    var cache: Cache = .{};
+    defer cache.deinit(gpa);
+    const source = "const Thing = struct {\n    value: u32,\n};\n";
+    const budget = "large.zig".len + source.len;
+    var result = try cache.getBounded(gpa, "large.zig", source, true, false, budget);
+    defer result.deinit(gpa);
+    try std.testing.expect(result.owned);
+    try std.testing.expect(result.analysis.declarations.len > 0);
+    try std.testing.expectEqual(@as(usize, 1), cache.analyses);
+    try std.testing.expectEqual(@as(usize, 0), cache.len);
+    try std.testing.expectEqual(@as(usize, 0), cache.bytes);
+}