Workspace Symbols

build.zig

@@ -68,6 +68,7 @@ pub fn build(b: *Build) !void {
 
     const known_folders_module = b.dependency("known_folders", .{}).module("known-folders");
     const diffz_module = b.dependency("diffz", .{}).module("diffz");
+    const fastfilter_module = b.dependency("fastfilter", .{}).module("fastfilter");
     const tracy_module = getTracyModule(b, .{
         .target = target,
         .optimize = optimize,
@@ -124,6 +125,7 @@ pub fn build(b: *Build) !void {
         .imports = &.{
             .{ .name = "known-folders", .module = known_folders_module },
             .{ .name = "diffz", .module = diffz_module },
+            .{ .name = "fastfilter", .module = fastfilter_module },
             .{ .name = "tracy", .module = tracy_module },
             .{ .name = "build_options", .module = build_options_module },
             .{ .name = "version_data", .module = version_data_module },

build.zig.zon

@@ -14,6 +14,10 @@
             .url = "https://github.com/ziglibs/diffz/archive/ef45c00d655e5e40faf35afbbde81a1fa5ed7ffb.tar.gz",
             .hash = "1220102cb2c669d82184fb1dc5380193d37d68b54e8d75b76b2d155b9af7d7e2e76d",
         },
+        .fastfilter = .{
+            .url = "https://github.com/hexops/fastfilter/archive/b6e46f6d4811da0d8a0f8675ab85407172a207ba.tar.gz",
+            .hash = "1220c2275142456c3eb8152f626092237a7795e93518896581a9358ce857b3ca550e",
+        },
     },
     .paths = .{""},
 }

src/DocumentScope.zig

@@ -10,7 +10,9 @@ const DocumentScope = @This();
 
 scopes: std.MultiArrayList(Scope) = .{},
 declarations: std.MultiArrayList(Declaration) = .{},
-/// used for looking up a child declaration in a given scope
+trigram_decls_mapping_capacity: u32 = 0,
+declarations_that_should_be_trigram_indexed: std.ArrayListUnmanaged(Declaration.Index) = .{},
+/// Used for looking up a child declaration in a given scope
 declaration_lookup_map: DeclarationLookupMap = .{},
 extra: std.ArrayListUnmanaged(u32) = .{},
 /// All identifier token that are in error sets.
@@ -181,6 +183,7 @@ const ScopeContext = struct {
 
         scopes_start: u32,
         declarations_start: u32,
+        should_trigrams_be_indexed: bool,
 
         fn pushDeclaration(
             pushed: PushedScope,
@@ -206,6 +209,11 @@ const ScopeContext = struct {
             try doc_scope.declarations.append(allocator, declaration);
             const declaration_index: Declaration.Index = @enumFromInt(doc_scope.declarations.len - 1);
 
+            if (pushed.should_trigrams_be_indexed and name.len >= 3) {
+                doc_scope.trigram_decls_mapping_capacity += @intCast(name.len - 2);
+                try doc_scope.declarations_that_should_be_trigram_indexed.append(allocator, declaration_index);
+            }
+
             const data = &doc_scope.scopes.items(.data)[@intFromEnum(pushed.scope)];
             const child_declarations = &doc_scope.scopes.items(.child_declarations)[@intFromEnum(pushed.scope)];
 
@@ -298,6 +306,10 @@ const ScopeContext = struct {
             .scope = context.current_scope.unwrap().?,
             .scopes_start = @intCast(context.child_scopes_scratch.items.len),
             .declarations_start = @intCast(context.child_declarations_scratch.items.len),
+            .should_trigrams_be_indexed = switch (tag) {
+                .container, .container_usingnamespace => true,
+                else => false,
+            },
         };
     }
 
@@ -353,6 +365,7 @@ pub fn init(allocator: std.mem.Allocator, tree: Ast) error{OutOfMemory}!Document
 pub fn deinit(scope: *DocumentScope, allocator: std.mem.Allocator) void {
     scope.scopes.deinit(allocator);
     scope.declarations.deinit(allocator);
+    scope.declarations_that_should_be_trigram_indexed.deinit(allocator);
     scope.declaration_lookup_map.deinit(allocator);
     scope.extra.deinit(allocator);
 

src/DocumentStore.zig

@@ -14,6 +14,7 @@ const AstGen = std.zig.AstGen;
 const Zir = std.zig.Zir;
 const InternPool = @import("analyser/InternPool.zig");
 const DocumentScope = @import("DocumentScope.zig");
+const TrigramStore = @import("TrigramStore.zig");
 
 const DocumentStore = @This();
 
@@ -23,6 +24,15 @@ config: Config,
 lock: std.Thread.RwLock = .{},
 thread_pool: if (builtin.single_threaded) void else *std.Thread.Pool,
 handles: std.StringArrayHashMapUnmanaged(*Handle) = .{},
+/// All `Handle`s in the current workspace folders
+/// have a `TrigramStore`, but there will be some
+/// `TrigramStore`s with no associated `Handle`.
+///
+/// This is an optimization that allows closed
+/// or seldom accessed documents to return results
+/// during Workspace Symbol searches without
+/// maintaining an otherwise useless `Handle`.
+trigram_stores: std.StringArrayHashMapUnmanaged(*TrigramStore) = .{},
 build_files: std.StringArrayHashMapUnmanaged(*BuildFile) = .{},
 cimports: std.AutoArrayHashMapUnmanaged(Hash, translate_c.Result) = .{},
 
@@ -47,14 +57,16 @@ pub const Config = struct {
     build_runner_path: ?[]const u8,
     builtin_path: ?[]const u8,
     global_cache_path: ?[]const u8,
+    workspace_folders: []const Uri = &.{},
 
-    pub fn fromMainConfig(config: @import("Config.zig")) Config {
+    pub fn fromMainConfig(config: @import("Config.zig"), workspace_folders: []const Uri) Config {
         return .{
             .zig_exe_path = config.zig_exe_path,
             .zig_lib_path = config.zig_lib_path,
             .build_runner_path = config.build_runner_path,
             .builtin_path = config.builtin_path,
             .global_cache_path = config.global_cache_path,
+            .workspace_folders = workspace_folders,
         };
     }
 };
@@ -642,6 +654,14 @@ pub fn deinit(self: *DocumentStore) void {
         result.deinit(self.allocator);
     }
     self.cimports.deinit(self.allocator);
+
+    for (self.trigram_stores.keys()) |uri| self.allocator.free(uri);
+    for (self.trigram_stores.values()) |trigram_store| {
+        trigram_store.deinit(self.allocator);
+        self.allocator.destroy(trigram_store);
+    }
+    self.trigram_stores.deinit(self.allocator);
+
     self.* = undefined;
 }
 
@@ -689,6 +709,41 @@ pub fn getOrLoadHandle(self: *DocumentStore, uri: Uri) ?*Handle {
     return self.createAndStoreDocument(uri, file_contents, false) catch return null;
 }
 
+/// Returns a handle to the given document
+/// **Thread safe** takes a shared lock
+/// This function does not protect against data races from modifying the Handle
+pub fn getTrigramStore(store: *DocumentStore, uri: Uri) error{OutOfMemory}!?*TrigramStore {
+    store.lock.lockShared();
+    defer store.lock.unlockShared();
+    return store.trigram_stores.get(uri);
+}
+
+pub fn getOrConstructTrigramStore(store: *DocumentStore, uri: Uri) error{OutOfMemory}!?*TrigramStore {
+    if (try store.getTrigramStore(uri)) |trigram_store| return trigram_store;
+
+    if (store.getOrLoadHandle(uri)) |handle| {
+        const duped_uri = try store.allocator.dupe(u8, uri);
+        errdefer store.allocator.free(duped_uri);
+
+        var trigram_store = try store.allocator.create(TrigramStore);
+        errdefer store.allocator.destroy(trigram_store);
+
+        trigram_store.* = TrigramStore.init(store.allocator, handle) catch |err| switch (err) {
+            error.InvalidUtf8 => {
+                store.allocator.free(duped_uri);
+                store.allocator.destroy(trigram_store);
+                return null;
+            },
+            else => |e| return e,
+        };
+        errdefer trigram_store.deinit(store.allocator);
+
+        try store.trigram_stores.put(store.allocator, duped_uri, trigram_store);
+    }
+
+    return null;
+}
+
 /// **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 {
@@ -956,6 +1011,14 @@ pub fn isInStd(uri: Uri) bool {
     return std.mem.indexOf(u8, uri, "/std/") != null;
 }
 
+pub fn isInWorkspaceFolders(store: DocumentStore, uri: Uri) bool {
+    // TODO: Better logic for detecting workspace folders?
+    return for (store.config.workspace_folders) |wf| {
+        // NOTE: terrible but it should work
+        if (std.mem.startsWith(u8, uri, wf)) break true;
+    } else false;
+}
+
 /// looks for a `zls.build.json` file in the build file directory
 /// has to be freed with `json_compat.parseFree`
 fn loadBuildAssociatedConfiguration(allocator: std.mem.Allocator, build_file: BuildFile) !std.json.Parsed(BuildAssociatedConfig) {
@@ -1317,7 +1380,9 @@ fn createAndStoreDocument(self: *DocumentStore, uri: Uri, text: [:0]const u8, op
         self.allocator.destroy(handle_ptr);
     }
 
-    if (isBuildFile(gop.value_ptr.*.uri)) {
+    const is_build_file = isBuildFile(gop.value_ptr.*.uri);
+
+    if (is_build_file) {
         log.debug("Opened document `{s}` (build file)", .{gop.value_ptr.*.uri});
     } else {
         log.debug("Opened document `{s}`", .{gop.value_ptr.*.uri});

src/Server.zig

@@ -6,6 +6,8 @@ const build_options = @import("build_options");
 const Config = @import("Config.zig");
 const configuration = @import("configuration.zig");
 const DocumentStore = @import("DocumentStore.zig");
+const DocumentScope = @import("DocumentScope.zig");
+const TrigramStore = @import("TrigramStore.zig");
 const types = @import("lsp.zig");
 const Analyser = @import("analysis.zig");
 const ast = @import("ast.zig");
@@ -26,12 +28,13 @@ const inlay_hints = @import("features/inlay_hints.zig");
 const code_actions = @import("features/code_actions.zig");
 const folding_range = @import("features/folding_range.zig");
 const document_symbol = @import("features/document_symbol.zig");
+const workspace_symbols = @import("features/workspace_symbols.zig");
 const completions = @import("features/completions.zig");
 const goto = @import("features/goto.zig");
 const hover_handler = @import("features/hover.zig");
 const selection_range = @import("features/selection_range.zig");
 const diagnostics_gen = @import("features/diagnostics.zig");
-
+const URI = @import("uri.zig");
 const log = std.log.scoped(.zls_server);
 
 // public fields
@@ -518,6 +521,7 @@ fn initializeHandler(server: *Server, _: std.mem.Allocator, request: types.Initi
         for (server.client_capabilities.workspace_folders, workspace_folders) |*dest, src| {
             dest.* = try server.allocator.dupe(u8, src.uri);
         }
+        server.document_store.config.workspace_folders = server.client_capabilities.workspace_folders;
     }
 
     if (request.trace) |trace| {
@@ -607,7 +611,7 @@ fn initializeHandler(server: *Server, _: std.mem.Allocator, request: types.Initi
             .documentRangeFormattingProvider = .{ .bool = false },
             .foldingRangeProvider = .{ .bool = true },
             .selectionRangeProvider = .{ .bool = true },
-            .workspaceSymbolProvider = .{ .bool = false },
+            .workspaceSymbolProvider = .{ .bool = true },
             .workspace = .{
                 .workspaceFolders = .{
                     .supported = true,
@@ -778,6 +782,8 @@ fn didChangeWorkspaceFoldersHandler(server: *Server, arena: std.mem.Allocator, n
     }
 
     server.client_capabilities.workspace_folders = try folders.toOwnedSlice(server.allocator);
+
+    server.document_store.config.workspace_folders = server.client_capabilities.workspace_folders;
 }
 
 fn didChangeConfigurationHandler(server: *Server, arena: std.mem.Allocator, notification: types.DidChangeConfigurationParams) Error!void {
@@ -880,7 +886,7 @@ pub fn updateConfiguration(server: *Server, new_config: configuration.Configurat
         }
     }
 
-    server.document_store.config = DocumentStore.Config.fromMainConfig(server.config);
+    server.document_store.config = DocumentStore.Config.fromMainConfig(server.config, server.client_capabilities.workspace_folders);
 
     if (new_zig_exe_path or new_build_runner_path) blk: {
         if (!std.process.can_spawn) break :blk;
@@ -1595,6 +1601,101 @@ fn selectionRangeHandler(server: *Server, arena: std.mem.Allocator, request: typ
     return try selection_range.generateSelectionRanges(arena, handle, request.positions, server.offset_encoding);
 }
 
+fn workspaceSymbolHandler(server: *Server, arena: std.mem.Allocator, request: types.WorkspaceSymbolParams) Error!ResultType("workspace/symbol") {
+    if (request.query.len < 3) return null;
+
+    var trigrams = std.ArrayListUnmanaged(TrigramStore.Trigram){};
+
+    for (server.client_capabilities.workspace_folders) |workspace_folder| {
+        const path = URI.parse(arena, workspace_folder) 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.pathRelative(arena, workspace_folder, entry.path) catch return error.InternalError;
+                _ = try server.document_store.getOrConstructTrigramStore(uri);
+            }
+        }
+    }
+
+    if (std.unicode.Utf8View.init(request.query)) |view| {
+        var iterator = view.iterator();
+        while (iterator.nextCodepoint()) |codepoint_0| {
+            const next_idx = iterator.i;
+            const codepoint_1 = iterator.nextCodepoint() orelse break;
+            const codepoint_2 = iterator.nextCodepoint() orelse break;
+            try trigrams.append(arena, .{
+                .codepoint_0 = codepoint_0,
+                .codepoint_1 = codepoint_1,
+                .codepoint_2 = codepoint_2,
+            });
+            iterator.i = next_idx;
+        }
+    } else |_| return null;
+
+    if (trigrams.items.len == 0) return null;
+
+    var symbols = std.ArrayListUnmanaged(types.WorkspaceSymbol){};
+    var candidate_decls_buffer = std.ArrayListUnmanaged(Analyser.Declaration.Index){};
+
+    doc_loop: for (server.document_store.trigram_stores.keys(), server.document_store.trigram_stores.values()) |uri, trigram_store| {
+        const handle = server.document_store.getOrLoadHandle(uri) orelse continue;
+
+        const tree = handle.tree;
+        const doc_scope = try handle.getDocumentScope();
+
+        for (trigrams.items) |trigram| {
+            if (!trigram_store.filter.contain(@bitCast(trigram))) continue :doc_loop;
+        }
+
+        candidate_decls_buffer.clearRetainingCapacity();
+
+        const first = trigram_store.getDeclarationsForTrigram(trigrams.items[0]) orelse continue;
+
+        try candidate_decls_buffer.resize(arena, first.len * 2);
+
+        var len = first.len;
+
+        @memcpy(candidate_decls_buffer.items[0..len], first);
+        @memcpy(candidate_decls_buffer.items[len..], first);
+
+        for (trigrams.items[1..]) |trigram| {
+            len = workspace_symbols.mergeIntersection(
+                trigram_store.getDeclarationsForTrigram(trigram) orelse continue :doc_loop,
+                candidate_decls_buffer.items[len..],
+                candidate_decls_buffer.items[0..len],
+            );
+            candidate_decls_buffer.items.len = len * 2;
+            @memcpy(candidate_decls_buffer.items[len..], candidate_decls_buffer.items[0..len]);
+        }
+
+        candidate_decls_buffer.items.len = len;
+
+        for (candidate_decls_buffer.items) |decl_idx| {
+            const decl = doc_scope.declarations.get(@intFromEnum(decl_idx));
+            const name_token = decl.nameToken(tree);
+
+            // TODO: integrate with document_symbol.zig for right kind info
+            try symbols.append(arena, .{
+                .name = tree.tokenSlice(name_token),
+                .kind = .Variable,
+                .location = .{
+                    .Location = .{
+                        .uri = handle.uri,
+                        .range = offsets.tokenToRange(tree, name_token, server.offset_encoding),
+                    },
+                },
+            });
+        }
+    }
+
+    return .{ .array_of_WorkspaceSymbol = symbols.items };
+}
+
 const HandledRequestMethods = enum {
     initialize,
     shutdown,
@@ -1617,6 +1718,7 @@ const HandledRequestMethods = enum {
     @"textDocument/codeAction",
     @"textDocument/foldingRange",
     @"textDocument/selectionRange",
+    @"workspace/symbol",
 };
 
 const HandledNotificationMethods = enum {
@@ -1657,6 +1759,7 @@ fn isBlockingMessage(msg: types.Message) bool {
             .@"textDocument/codeAction",
             .@"textDocument/foldingRange",
             .@"textDocument/selectionRange",
+            .@"workspace/symbol",
             => return false,
         },
         .notification => |notification| switch (std.meta.stringToEnum(HandledNotificationMethods, notification.method) orelse return false) {
@@ -1685,7 +1788,7 @@ pub fn create(allocator: std.mem.Allocator) !*Server {
         .config = .{},
         .document_store = .{
             .allocator = allocator,
-            .config = DocumentStore.Config.fromMainConfig(Config{}),
+            .config = DocumentStore.Config.fromMainConfig(Config{}, &.{}),
             .thread_pool = if (zig_builtin.single_threaded) {} else undefined, // set below
         },
         .job_queue = std.fifo.LinearFifo(Job, .Dynamic).init(allocator),
@@ -1820,6 +1923,7 @@ pub fn sendRequestSync(server: *Server, arena: std.mem.Allocator, comptime metho
         .@"textDocument/codeAction" => try server.codeActionHandler(arena, params),
         .@"textDocument/foldingRange" => try server.foldingRangeHandler(arena, params),
         .@"textDocument/selectionRange" => try server.selectionRangeHandler(arena, params),
+        .@"workspace/symbol" => try server.workspaceSymbolHandler(arena, params),
     };
 }