Trigram indexing works and it's fast

Co-authored-by: Techatrix <[email protected]>

Author: SuperAuguste

src/DocumentStore.zig

@@ -187,6 +187,7 @@ pub const BuildFile = struct {
 pub const Handle = struct {
     uri: Uri,
     tree: Ast,
+    trigram_store: TrigramStore,
     /// Contains one entry for every import in the document
     import_uris: std.ArrayListUnmanaged(Uri) = .empty,
     /// Contains one entry for every cimport in the document
@@ -270,13 +271,17 @@ pub const Handle = struct {
         var tree = try parseTree(allocator, text, mode);
         errdefer tree.deinit(allocator);
 
-        const cimports = try collectCImports(allocator, tree);
+        var cimports = try collectCImports(allocator, tree);
         errdefer cimports.deinit(allocator);
 
+        var trigram_store: TrigramStore = try .init(allocator, tree, .@"utf-16");
+        errdefer trigram_store.deinit();
+
         return .{
             .uri = uri,
             .tree = tree,
             .cimports = cimports,
+            .trigram_store = trigram_store,
             .impl = .{
                 .status = .init(@bitCast(Status{
                     .lsp_synced = lsp_synced,
@@ -307,6 +312,8 @@ pub const Handle = struct {
         for (self.cimports.items(.source)) |source| allocator.free(source);
         self.cimports.deinit(allocator);
 
+        self.trigram_store.deinit(allocator);
+
         switch (self.impl.associated_build_file) {
             .none, .resolved => {},
             .unresolved => |*payload| payload.deinit(allocator),
@@ -697,31 +704,6 @@ pub fn getOrLoadHandle(store: *DocumentStore, uri: Uri) ?*Handle {
     };
 }
 
-pub fn trigramIndexUri(
-    store: *DocumentStore,
-    uri: Uri,
-    encoding: offsets.Encoding,
-) error{OutOfMemory}!void {
-    const gop = try store.trigram_stores.getOrPut(store.allocator, uri);
-
-    if (gop.found_existing) {
-        return;
-    }
-
-    errdefer {
-        store.allocator.free(gop.key_ptr.*);
-        store.trigram_stores.swapRemoveAt(gop.index);
-    }
-
-    gop.key_ptr.* = try store.allocator.dupe(u8, uri);
-    gop.value_ptr.* = .empty;
-
-    const file_contents = store.readUri(uri) orelse return;
-    defer store.allocator.free(file_contents);
-
-    try gop.value_ptr.fill(store.allocator, file_contents, encoding);
-}
-
 /// **Thread safe** takes a shared lock
 /// This function does not protect against data races from modifying the BuildFile
 pub fn getBuildFile(self: *DocumentStore, uri: Uri) ?*BuildFile {

src/Server.zig

@@ -1936,36 +1936,17 @@ fn selectionRangeHandler(server: *Server, arena: std.mem.Allocator, request: typ
 fn workspaceSymbolHandler(server: *Server, arena: std.mem.Allocator, request: types.WorkspaceSymbolParams) Error!lsp.ResultType("workspace/symbol") {
     if (request.query.len < 3) return null;
 
-    for (server.workspaces.items) |workspace| {
-        const path = Uri.parse(arena, workspace.uri) catch return error.InternalError;
-        var dir = std.fs.cwd().openDir(path, .{ .iterate = true }) catch return error.InternalError;
-        defer dir.close();
-
-        var walker = try dir.walk(arena);
-        defer walker.deinit();
-
-        while (walker.next() catch return error.InternalError) |entry| {
-            if (std.mem.eql(u8, std.fs.path.extension(entry.basename), ".zig")) {
-                const uri = Uri.fromPath(
-                    arena,
-                    std.fs.path.join(arena, &.{ path, entry.path }) catch return error.InternalError,
-                ) catch return error.InternalError;
-
-                server.document_store.trigramIndexUri(
-                    uri,
-                    server.offset_encoding,
-                ) catch return error.InternalError;
-            }
-        }
-    }
+    // TODO: take this and get copy of handle ptrs
+    server.document_store.lock.lock();
+    defer server.document_store.lock.unlock();
 
     var symbols: std.ArrayListUnmanaged(types.WorkspaceSymbol) = .empty;
     var declaration_buffer: std.ArrayListUnmanaged(TrigramStore.Declaration.Index) = .empty;
 
-    for (
-        server.document_store.trigram_stores.keys(),
-        server.document_store.trigram_stores.values(),
-    ) |uri, trigram_store| {
+    for (server.document_store.handles.keys(), server.document_store.handles.values()) |uri, handle| {
+        const trigram_store = &handle.trigram_store;
+
+        declaration_buffer.clearRetainingCapacity();
         try trigram_store.declarationsForQuery(arena, request.query, &declaration_buffer);
 
         const slice = trigram_store.declarations.slice();

src/TrigramStore.zig

@@ -22,51 +22,25 @@ pub const Declaration = struct {
     range: offsets.Range,
 };
 
-pub const empty: TrigramStore = .{
-    .has_filter = false,
-    .filter_buckets = .empty,
-    .trigram_to_declarations = .empty,
-    .declarations = .empty,
-    .names = .empty,
-};
-
 has_filter: bool,
 filter_buckets: std.ArrayListUnmanaged(CuckooFilter.Bucket),
 trigram_to_declarations: std.AutoArrayHashMapUnmanaged(Trigram, std.ArrayListUnmanaged(Declaration.Index)),
 declarations: std.MultiArrayList(Declaration),
 names: std.ArrayListUnmanaged(u8),
 
-pub fn deinit(store: *TrigramStore, allocator: std.mem.Allocator) void {
-    store.filter_buckets.deinit(allocator);
-    for (store.trigram_to_declarations.values()) |*list| {
-        list.deinit(allocator);
-    }
-    store.trigram_to_declarations.deinit(allocator);
-    store.declarations.deinit(allocator);
-    store.names.deinit(allocator);
-    store.* = undefined;
-}
-
-fn clearRetainingCapacity(store: *TrigramStore) void {
-    store.filter_buckets.clearRetainingCapacity();
-    store.has_filter = false;
-    for (store.trigram_to_declarations.values()) |*list| {
-        list.clearRetainingCapacity();
-    }
-    store.declarations.clearRetainingCapacity();
-    store.names.clearRetainingCapacity();
-}
-
-pub fn fill(
-    store: *TrigramStore,
+pub fn init(
     allocator: std.mem.Allocator,
-    source: [:0]const u8,
+    tree: Ast,
     encoding: offsets.Encoding,
-) error{OutOfMemory}!void {
-    store.clearRetainingCapacity();
-
-    var tree = try Ast.parse(allocator, source, .zig);
-    defer tree.deinit(allocator);
+) error{OutOfMemory}!TrigramStore {
+    var store: TrigramStore = .{
+        .has_filter = false,
+        .filter_buckets = .empty,
+        .trigram_to_declarations = .empty,
+        .declarations = .empty,
+        .names = .empty,
+    };
+    errdefer store.deinit(allocator);
 
     const Context = struct {
         allocator: std.mem.Allocator,
@@ -126,15 +100,61 @@ pub fn fill(
         }
     };
 
-    var context = Context{
+    var context: Context = .{
         .allocator = allocator,
-        .store = store,
+        .store = &store,
         .in_function = false,
         .encoding = encoding,
     };
     try ast.iterateChildren(tree, .root, &context, Context.Error, Context.callback);
 
-    try store.finalize(allocator);
+    const lists = store.trigram_to_declarations.values();
+    var index: usize = 0;
+    while (index < lists.len) {
+        if (lists[index].items.len == 0) {
+            lists[index].deinit(allocator);
+            store.trigram_to_declarations.swapRemoveAt(index);
+        } else {
+            index += 1;
+        }
+    }
+
+    const trigrams = store.trigram_to_declarations.keys();
+
+    if (trigrams.len > 0) {
+        var prng = std.Random.DefaultPrng.init(0);
+
+        const filter_capacity = CuckooFilter.capacityForCount(store.trigram_to_declarations.count()) catch unreachable;
+        try store.filter_buckets.ensureTotalCapacityPrecise(allocator, filter_capacity);
+        store.filter_buckets.items.len = filter_capacity;
+
+        const filter: CuckooFilter = .{ .buckets = store.filter_buckets.items };
+        filter.reset();
+        store.has_filter = true;
+
+        for (trigrams) |trigram| {
+            filter.append(prng.random(), trigram) catch |err| switch (err) {
+                error.EvictionFailed => {
+                    // NOTE: This should generally be quite rare.
+                    store.has_filter = false;
+                    break;
+                },
+            };
+        }
+    }
+
+    return store;
+}
+
+pub fn deinit(store: *TrigramStore, allocator: std.mem.Allocator) void {
+    store.filter_buckets.deinit(allocator);
+    for (store.trigram_to_declarations.values()) |*list| {
+        list.deinit(allocator);
+    }
+    store.trigram_to_declarations.deinit(allocator);
+    store.declarations.deinit(allocator);
+    store.names.deinit(allocator);
+    store.* = undefined;
 }
 
 /// Caller must not submit name.len < 3.
@@ -167,53 +187,15 @@ fn appendDeclaration(
     }
 }
 
-/// Must be called before any queries are executed.
-fn finalize(store: *TrigramStore, allocator: std.mem.Allocator) error{OutOfMemory}!void {
-    {
-        const lists = store.trigram_to_declarations.values();
-        var index: usize = 0;
-        while (index < lists.len) {
-            if (lists[index].items.len == 0) {
-                lists[index].deinit(allocator);
-                store.trigram_to_declarations.swapRemoveAt(index);
-            } else {
-                index += 1;
-            }
-        }
-    }
-
-    const trigrams = store.trigram_to_declarations.keys();
-
-    if (trigrams.len > 0) {
-        var prng = std.Random.DefaultPrng.init(0);
-
-        const filter_capacity = CuckooFilter.capacityForCount(store.trigram_to_declarations.count()) catch unreachable;
-        try store.filter_buckets.ensureTotalCapacityPrecise(allocator, filter_capacity);
-        store.filter_buckets.items.len = filter_capacity;
-
-        const filter: CuckooFilter = .{ .buckets = store.filter_buckets.items };
-        filter.reset();
-        store.has_filter = true;
-
-        for (trigrams) |trigram| {
-            filter.append(prng.random(), trigram) catch |err| switch (err) {
-                error.EvictionFailed => {
-                    // NOTE: This should generally be quite rare.
-                    store.has_filter = false;
-                    break;
-                },
-            };
-        }
-    }
-}
-
+/// Asserts query.len >= 3. Asserts declaration_buffer.items.len == 0.
 pub fn declarationsForQuery(
     store: *const TrigramStore,
     allocator: std.mem.Allocator,
     query: []const u8,
     declaration_buffer: *std.ArrayListUnmanaged(Declaration.Index),
 ) error{OutOfMemory}!void {
     assert(query.len >= 3);
+    assert(declaration_buffer.items.len == 0);
 
     const filter: CuckooFilter = .{ .buckets = store.filter_buckets.items };
 
@@ -226,14 +208,9 @@ pub fn declarationsForQuery(
         }
     }
 
-    const first = (store.trigram_to_declarations.get(query[0..3].*) orelse {
-        declaration_buffer.clearRetainingCapacity();
-        return;
-    }).items;
+    const first = (store.trigram_to_declarations.get(query[0..3].*) orelse return).items;
 
-    declaration_buffer.clearRetainingCapacity();
-    try declaration_buffer.ensureTotalCapacity(allocator, first.len * 2);
-    declaration_buffer.items.len = first.len * 2;
+    try declaration_buffer.resize(allocator, first.len * 2);
 
     var len = first.len;
     @memcpy(declaration_buffer.items[0..len], first);
@@ -242,18 +219,18 @@ pub fn declarationsForQuery(
         const trigram = query[index..][0..3].*;
         const old_len = len;
         len = mergeIntersection(
-            (store.trigram_to_declarations.get(trigram[0..3].*) orelse return {
+            (store.trigram_to_declarations.get(trigram[0..3].*) orelse {
                 declaration_buffer.clearRetainingCapacity();
                 return;
             }).items,
             declaration_buffer.items[0..len],
             declaration_buffer.items[len..],
         );
         @memcpy(declaration_buffer.items[0..len], declaration_buffer.items[old_len..][0..len]);
-        declaration_buffer.items.len = len * 2;
+        declaration_buffer.shrinkRetainingCapacity(len * 2);
     }
 
-    declaration_buffer.items.len = declaration_buffer.items.len / 2;
+    declaration_buffer.shrinkRetainingCapacity(declaration_buffer.items.len / 2);
 }
 
 /// Asserts `@min(a.len, b.len) <= out.len`.