Optimize document symbols (#1050)

* optimize document symbol generation

* match folding range index to position conversion and documentation to document symbol's

* skip function decls with no name

* skip document symbol field in opaque type

src/Server.zig

@@ -15,6 +15,7 @@ const semantic_tokens = @import("semantic_tokens.zig");
 const inlay_hints = @import("inlay_hints.zig");
 const code_actions = @import("code_actions.zig");
 const folding_range = @import("folding_range.zig");
+const document_symbol = @import("document_symbol.zig");
 const shared = @import("shared.zig");
 const Ast = std.zig.Ast;
 const tracy = @import("tracy.zig");
@@ -2498,7 +2499,7 @@ pub fn documentSymbolsHandler(server: *Server, request: types.DocumentSymbolPara
 
     const handle = server.document_store.getHandle(request.textDocument.uri) orelse return null;
 
-    return try analysis.getDocumentSymbols(server.arena.allocator(), handle.tree, server.offset_encoding);
+    return try document_symbol.getDocumentSymbols(server.arena.allocator(), handle.tree, server.offset_encoding);
 }
 
 pub fn formattingHandler(server: *Server, request: types.DocumentFormattingParams) Error!?[]types.TextEdit {

src/analysis.zig

@@ -1804,237 +1804,6 @@ pub fn getPositionContext(
     return if (tok.tag == .identifier) PositionContext{ .var_access = tok.loc } else .empty;
 }
 
-fn addOutlineNodes(allocator: std.mem.Allocator, tree: Ast, child: Ast.Node.Index, context: *GetDocumentSymbolsContext) error{OutOfMemory}!void {
-    switch (tree.nodes.items(.tag)[child]) {
-        .string_literal,
-        .number_literal,
-        .builtin_call,
-        .builtin_call_comma,
-        .builtin_call_two,
-        .builtin_call_two_comma,
-        .call,
-        .call_comma,
-        .call_one,
-        .call_one_comma,
-        .async_call,
-        .async_call_comma,
-        .async_call_one,
-        .async_call_one_comma,
-        .identifier,
-        .add,
-        .add_wrap,
-        .array_cat,
-        .array_mult,
-        .assign,
-        .assign_bit_and,
-        .assign_bit_or,
-        .assign_shl,
-        .assign_shr,
-        .assign_bit_xor,
-        .assign_div,
-        .assign_sub,
-        .assign_sub_wrap,
-        .assign_mod,
-        .assign_add,
-        .assign_add_wrap,
-        .assign_mul,
-        .assign_mul_wrap,
-        .bang_equal,
-        .bit_and,
-        .bit_or,
-        .shl,
-        .shr,
-        .bit_xor,
-        .bool_and,
-        .bool_or,
-        .div,
-        .equal_equal,
-        .error_union,
-        .greater_or_equal,
-        .greater_than,
-        .less_or_equal,
-        .less_than,
-        .merge_error_sets,
-        .mod,
-        .mul,
-        .mul_wrap,
-        .field_access,
-        .switch_range,
-        .sub,
-        .sub_wrap,
-        .@"orelse",
-        .address_of,
-        .@"await",
-        .bit_not,
-        .bool_not,
-        .optional_type,
-        .negation,
-        .negation_wrap,
-        .@"resume",
-        .@"try",
-        .array_type,
-        .array_type_sentinel,
-        .ptr_type,
-        .ptr_type_aligned,
-        .ptr_type_bit_range,
-        .ptr_type_sentinel,
-        .slice_open,
-        .slice_sentinel,
-        .deref,
-        .unwrap_optional,
-        .array_access,
-        .@"return",
-        .@"break",
-        .@"continue",
-        .array_init,
-        .array_init_comma,
-        .array_init_dot,
-        .array_init_dot_comma,
-        .array_init_dot_two,
-        .array_init_dot_two_comma,
-        .array_init_one,
-        .array_init_one_comma,
-        .@"switch",
-        .switch_comma,
-        .switch_case,
-        .switch_case_one,
-        .@"for",
-        .for_simple,
-        .enum_literal,
-        .struct_init,
-        .struct_init_comma,
-        .struct_init_dot,
-        .struct_init_dot_comma,
-        .struct_init_dot_two,
-        .struct_init_dot_two_comma,
-        .struct_init_one,
-        .struct_init_one_comma,
-        .@"while",
-        .while_simple,
-        .while_cont,
-        .@"defer",
-        .@"if",
-        .if_simple,
-        .multiline_string_literal,
-        .block,
-        .block_semicolon,
-        .block_two,
-        .block_two_semicolon,
-        .error_set_decl,
-        => return,
-        .container_decl,
-        .container_decl_trailing,
-        .container_decl_arg,
-        .container_decl_arg_trailing,
-        .container_decl_two,
-        .container_decl_two_trailing,
-        .tagged_union,
-        .tagged_union_trailing,
-        .tagged_union_enum_tag,
-        .tagged_union_enum_tag_trailing,
-        .tagged_union_two,
-        .tagged_union_two_trailing,
-        => {
-            var buf: [2]Ast.Node.Index = undefined;
-            const members = tree.fullContainerDecl(&buf, child).?.ast.members;
-            for (members) |member|
-                try addOutlineNodes(allocator, tree, member, context);
-            return;
-        },
-        else => {},
-    }
-    try getDocumentSymbolsInternal(allocator, tree, child, context);
-}
-
-const GetDocumentSymbolsContext = struct {
-    symbols: *std.ArrayListUnmanaged(types.DocumentSymbol),
-    encoding: offsets.Encoding,
-};
-
-fn getDocumentSymbolsInternal(allocator: std.mem.Allocator, tree: Ast, node: Ast.Node.Index, context: *GetDocumentSymbolsContext) error{OutOfMemory}!void {
-    const name = getDeclName(tree, node) orelse return;
-    if (name.len == 0)
-        return;
-
-    const range = offsets.nodeToRange(tree, node, context.encoding);
-
-    const tags = tree.nodes.items(.tag);
-    (try context.symbols.addOne(allocator)).* = .{
-        .name = name,
-        .kind = switch (tags[node]) {
-            .fn_proto,
-            .fn_proto_simple,
-            .fn_proto_multi,
-            .fn_proto_one,
-            .fn_decl,
-            => .Function,
-            .local_var_decl,
-            .global_var_decl,
-            .aligned_var_decl,
-            .simple_var_decl,
-            => .Variable,
-            .container_field,
-            .container_field_align,
-            .container_field_init,
-            .tagged_union_enum_tag,
-            .tagged_union_enum_tag_trailing,
-            .tagged_union,
-            .tagged_union_trailing,
-            .tagged_union_two,
-            .tagged_union_two_trailing,
-            => .Field,
-            else => .Variable,
-        },
-        .range = range,
-        .selectionRange = range,
-        .detail = "",
-        .children = ch: {
-            var children = std.ArrayListUnmanaged(types.DocumentSymbol){};
-
-            var child_context = GetDocumentSymbolsContext{
-                .symbols = &children,
-                .encoding = context.encoding,
-            };
-
-            var buf: [2]Ast.Node.Index = undefined;
-            if (tree.fullContainerDecl(&buf, node)) |container_decl| {
-                for (container_decl.ast.members) |child| {
-                    try addOutlineNodes(allocator, tree, child, &child_context);
-                }
-            } else if (tree.fullVarDecl(node)) |var_decl| {
-                if (var_decl.ast.init_node != 0)
-                    try addOutlineNodes(allocator, tree, var_decl.ast.init_node, &child_context);
-            } else if (tags[node] == .fn_decl) fn_ch: {
-                const fn_decl = tree.nodes.items(.data)[node];
-                var params: [1]Ast.Node.Index = undefined;
-                const fn_proto = tree.fullFnProto(&params, fn_decl.lhs).?;
-                if (!isTypeFunction(tree, fn_proto)) break :fn_ch;
-                const ret_stmt = findReturnStatement(tree, fn_proto, fn_decl.rhs) orelse break :fn_ch;
-                const type_decl = tree.nodes.items(.data)[ret_stmt].lhs;
-                if (type_decl != 0)
-                    try addOutlineNodes(allocator, tree, type_decl, &child_context);
-            }
-            break :ch children.items;
-        },
-    };
-}
-
-pub fn getDocumentSymbols(allocator: std.mem.Allocator, tree: Ast, encoding: offsets.Encoding) ![]types.DocumentSymbol {
-    var symbols = std.ArrayListUnmanaged(types.DocumentSymbol){};
-    try symbols.ensureTotalCapacity(allocator, tree.rootDecls().len);
-
-    var context = GetDocumentSymbolsContext{
-        .symbols = &symbols,
-        .encoding = encoding,
-    };
-
-    for (tree.rootDecls()) |idx| {
-        try getDocumentSymbolsInternal(allocator, tree, idx, &context);
-    }
-
-    return symbols.items;
-}
-
 pub const Declaration = union(enum) {
     /// Index of the ast node
     ast_node: Ast.Node.Index,

src/document_symbol.zig

@@ -0,0 +1,241 @@
+const std = @import("std");
+const types = @import("lsp.zig");
+const offsets = @import("offsets.zig");
+const ast = @import("ast.zig");
+const analysis = @import("analysis.zig");
+const Ast = std.zig.Ast;
+const log = std.log.scoped(.zls_document_symbol);
+const tracy = @import("tracy.zig");
+
+const Symbol = struct {
+    name: []const u8,
+    detail: ?[]const u8 = null,
+    kind: types.SymbolKind,
+    loc: offsets.Loc,
+    selection_loc: offsets.Loc,
+    children: std.ArrayListUnmanaged(Symbol),
+};
+
+pub const Context = struct {
+    arena: std.mem.Allocator,
+    last_var_decl_name: ?[]const u8,
+    parent_node: Ast.Node.Index,
+    parent_symbols: *std.ArrayListUnmanaged(Symbol),
+    total_symbol_count: *usize,
+};
+
+fn callback(ctx: *Context, tree: Ast, node: Ast.Node.Index) error{OutOfMemory}!void {
+    if (node == 0) return;
+
+    const node_tags = tree.nodes.items(.tag);
+    const main_tokens = tree.nodes.items(.main_token);
+    const token_tags = tree.tokens.items(.tag);
+
+    const decl_name_token = analysis.getDeclNameToken(tree, node);
+    const decl_name = analysis.getDeclName(tree, node);
+
+    var new_ctx = ctx.*;
+    const maybe_symbol: ?Symbol = switch (node_tags[node]) {
+        .global_var_decl,
+        .local_var_decl,
+        .simple_var_decl,
+        .aligned_var_decl,
+        => blk: {
+            new_ctx.last_var_decl_name = decl_name;
+            if (!ast.isContainer(tree, ctx.parent_node)) break :blk null;
+
+            const kind: types.SymbolKind = switch (token_tags[main_tokens[node]]) {
+                .keyword_var => .Variable,
+                .keyword_const => .Constant,
+                else => unreachable,
+            };
+
+            break :blk .{
+                .name = decl_name.?,
+                .detail = null,
+                .kind = kind,
+                .loc = offsets.nodeToLoc(tree, node),
+                .selection_loc = offsets.tokenToLoc(tree, decl_name_token.?),
+                .children = .{},
+            };
+        },
+
+        .test_decl,
+        .fn_decl,
+        => |tag| blk: {
+            const kind: types.SymbolKind = switch (tag) {
+                .test_decl => .Method, // there is no SymbolKind that represents a tests
+                .fn_decl => .Function,
+                else => unreachable,
+            };
+
+            var buffer: [1]Ast.Node.Index = undefined;
+            const detail = if (tree.fullFnProto(&buffer, node)) |fn_info| analysis.getFunctionSignature(tree, fn_info) else null;
+
+            break :blk .{
+                .name = decl_name orelse break :blk null,
+                .detail = detail,
+                .kind = kind,
+                .loc = offsets.nodeToLoc(tree, node),
+                .selection_loc = offsets.tokenToLoc(tree, decl_name_token.?),
+                .children = .{},
+            };
+        },
+
+        .container_field_init,
+        .container_field_align,
+        .container_field,
+        => blk: {
+            const kind: types.SymbolKind = switch (node_tags[ctx.parent_node]) {
+                .root => .Field,
+                .container_decl,
+                .container_decl_trailing,
+                .container_decl_arg,
+                .container_decl_arg_trailing,
+                .container_decl_two,
+                .container_decl_two_trailing,
+                => switch (token_tags[main_tokens[ctx.parent_node]]) {
+                    .keyword_struct => .Field,
+                    .keyword_union => .Field,
+                    .keyword_enum => .EnumMember,
+                    .keyword_opaque => break :blk null,
+                    else => unreachable,
+                },
+                .tagged_union,
+                .tagged_union_trailing,
+                .tagged_union_enum_tag,
+                .tagged_union_enum_tag_trailing,
+                .tagged_union_two,
+                .tagged_union_two_trailing,
+                => .Field,
+                else => unreachable,
+            };
+
+            break :blk .{
+                .name = decl_name.?,
+                .detail = ctx.last_var_decl_name,
+                .kind = kind,
+                .loc = offsets.nodeToLoc(tree, node),
+                .selection_loc = offsets.tokenToLoc(tree, decl_name_token.?),
+                .children = .{},
+            };
+        },
+        else => null,
+    };
+
+    new_ctx.parent_node = node;
+    if (maybe_symbol) |symbol| {
+        var symbol_ptr = try ctx.parent_symbols.addOne(ctx.arena);
+        symbol_ptr.* = symbol;
+        new_ctx.parent_symbols = &symbol_ptr.children;
+        ctx.total_symbol_count.* += 1;
+    }
+
+    try ast.iterateChildren(tree, node, &new_ctx, error{OutOfMemory}, callback);
+}
+
+/// a mapping from a source index to a line character pair
+const IndexToPositionEntry = struct {
+    output: *types.Position,
+    source_index: usize,
+
+    const Self = @This();
+
+    fn lessThan(_: void, lhs: Self, rhs: Self) bool {
+        return lhs.source_index < rhs.source_index;
+    }
+};
+
+/// converts `Symbol` to `types.DocumentSymbol`
+fn convertSymbols(
+    arena: std.mem.Allocator,
+    tree: Ast,
+    from: []const Symbol,
+    total_symbol_count: usize,
+    encoding: offsets.Encoding,
+) error{OutOfMemory}![]types.DocumentSymbol {
+    const tracy_zone = tracy.trace(@src());
+    defer tracy_zone.end();
+
+    var symbol_buffer = std.ArrayListUnmanaged(types.DocumentSymbol){};
+    try symbol_buffer.ensureTotalCapacityPrecise(arena, total_symbol_count);
+
+    // instead of converting every `offsets.Loc` to `types.Range` by calling `offsets.locToRange`
+    // we instead store a mapping from source indices to their desired position, sort them by their source index
+    // and then iterate through them which avoids having to re-iterate through the source file to find out the line number
+    // this reduces algorithmic complexity from `O(file_size*symbol_count)` to `O(symbol_count*log(symbol_count))`
+    var mappings = std.ArrayListUnmanaged(IndexToPositionEntry){};
+    try mappings.ensureTotalCapacityPrecise(arena, total_symbol_count * 4);
+
+    const result = convertSymbolsInternal(from, &symbol_buffer, &mappings);
+
+    // sort items based on their source position
+    std.sort.sort(IndexToPositionEntry, mappings.items, {}, IndexToPositionEntry.lessThan);
+
+    var last_index: usize = 0;
+    var last_position: types.Position = .{ .line = 0, .character = 0 };
+    for (mappings.items) |mapping| {
+        const index = mapping.source_index;
+        const position = offsets.advancePosition(tree.source, last_position, last_index, index, encoding);
+        defer last_index = index;
+        defer last_position = position;
+
+        mapping.output.* = position;
+    }
+
+    return result;
+}
+
+fn convertSymbolsInternal(
+    from: []const Symbol,
+    symbol_buffer: *std.ArrayListUnmanaged(types.DocumentSymbol),
+    mappings: *std.ArrayListUnmanaged(IndexToPositionEntry),
+) []types.DocumentSymbol {
+    // aquire storage for exactly `from.len` symbols
+    const prev_len = symbol_buffer.items.len;
+    symbol_buffer.items.len += from.len;
+    const to: []types.DocumentSymbol = symbol_buffer.items[prev_len..];
+
+    for (from, to) |symbol, *out| {
+        out.* = .{
+            .name = symbol.name,
+            .detail = symbol.detail,
+            .kind = symbol.kind,
+            // will be set later through the mapping below
+            .range = undefined,
+            .selectionRange = undefined,
+            .children = convertSymbolsInternal(symbol.children.items, symbol_buffer, mappings),
+        };
+        mappings.appendSliceAssumeCapacity(&[4]IndexToPositionEntry{
+            .{ .output = &out.range.start, .source_index = symbol.loc.start },
+            .{ .output = &out.selectionRange.start, .source_index = symbol.selection_loc.start },
+            .{ .output = &out.selectionRange.end, .source_index = symbol.selection_loc.end },
+            .{ .output = &out.range.end, .source_index = symbol.loc.end },
+        });
+    }
+
+    return to;
+}
+
+pub fn getDocumentSymbols(
+    arena: std.mem.Allocator,
+    tree: Ast,
+    encoding: offsets.Encoding,
+) error{OutOfMemory}![]types.DocumentSymbol {
+    const tracy_zone = tracy.trace(@src());
+    defer tracy_zone.end();
+
+    var root_symbols = std.ArrayListUnmanaged(Symbol){};
+    var total_symbol_count: usize = 0;
+
+    var ctx = Context{
+        .arena = arena,
+        .last_var_decl_name = null,
+        .parent_node = 0, // root-node
+        .parent_symbols = &root_symbols,
+        .total_symbol_count = &total_symbol_count,
+    };
+    try ast.iterateChildren(tree, 0, &ctx, error{OutOfMemory}, callback);
+
+    return try convertSymbols(arena, tree, root_symbols.items, ctx.total_symbol_count.*, encoding);
+}

src/folding_range.zig

@@ -69,53 +69,47 @@ const Builder = struct {
             };
         }
 
-        const Item = struct {
+        // a mapping from a source index to a line character pair
+        const IndexToPositionEntry = struct {
             output: *types.FoldingRange,
-            input: *const FoldingRange,
+            source_index: usize,
             where: enum { start, end },
 
             const Self = @This();
 
-            fn getInputIndex(self: Self) usize {
-                return switch (self.where) {
-                    .start => self.input.loc.start,
-                    .end => self.input.loc.end,
-                };
-            }
-
             fn lessThan(_: void, lhs: Self, rhs: Self) bool {
-                return lhs.getInputIndex() < rhs.getInputIndex();
+                return lhs.source_index < rhs.source_index;
             }
         };
 
-        // one item for every start and end position
+        // one mapping for every start and end position
-        var items = try builder.allocator.alloc(Item, builder.locations.items.len * 2);
+        var mappings = try builder.allocator.alloc(IndexToPositionEntry, builder.locations.items.len * 2);
-        defer builder.allocator.free(items);
+        defer builder.allocator.free(mappings);
 
         for (builder.locations.items, result_locations, 0..) |*folding_range, *result, i| {
-            items[2 * i + 0] = .{ .output = result, .input = folding_range, .where = .start };
+            mappings[2 * i + 0] = .{ .output = result, .source_index = folding_range.loc.start, .where = .start };
-            items[2 * i + 1] = .{ .output = result, .input = folding_range, .where = .end };
+            mappings[2 * i + 1] = .{ .output = result, .source_index = folding_range.loc.end, .where = .end };
         }
 
-        // sort items based on their source position
+        // sort mappings based on their source index
-        std.sort.sort(Item, items, {}, Item.lessThan);
+        std.sort.sort(IndexToPositionEntry, mappings, {}, IndexToPositionEntry.lessThan);
 
         var last_index: usize = 0;
         var last_position: types.Position = .{ .line = 0, .character = 0 };
-        for (items) |item| {
+        for (mappings) |mapping| {
-            const index = item.getInputIndex();
+            const index = mapping.source_index;
             const position = offsets.advancePosition(builder.tree.source, last_position, last_index, index, builder.encoding);
             defer last_index = index;
             defer last_position = position;
 
-            switch (item.where) {
+            switch (mapping.where) {
                 .start => {
-                    item.output.startLine = position.line;
+                    mapping.output.startLine = position.line;
-                    item.output.startCharacter = position.character;
+                    mapping.output.startCharacter = position.character;
                 },
                 .end => {
-                    item.output.endLine = position.line;
+                    mapping.output.endLine = position.line;
-                    item.output.endCharacter = position.character;
+                    mapping.output.endCharacter = position.character;
                 },
             }
         }

tests/lsp_features/document_symbol.zig

@@ -7,34 +7,56 @@ const tres = @import("tres");
 const Context = @import("../context.zig").Context;
 
 const types = zls.types;
-const requests = zls.requests;
 
 const allocator: std.mem.Allocator = std.testing.allocator;
 
-test "documentSymbol - smoke" {
+test "documentSymbol - container decl" {
     try testDocumentSymbol(
         \\const S = struct {
         \\    fn f() void {}
         \\};
     ,
-        \\Variable S
+        \\Constant S
+        \\  Function f
+    );
+    try testDocumentSymbol(
+        \\const S = struct {
+        \\    alpha: u32,
+        \\    fn f() void {}
+        \\};
+    ,
+        \\Constant S
+        \\  Field alpha
         \\  Function f
     );
 }
 
-// FIXME: https://github.com/zigtools/zls/issues/986
+test "documentSymbol - enum" {
+    try testDocumentSymbol(
+        \\const E = enum {
+        \\    alpha,
+        \\    beta,
+        \\};
+    ,
+        \\Constant E
+        \\  EnumMember alpha
+        \\  EnumMember beta
+    );
+}
+
+// https://github.com/zigtools/zls/issues/986
 test "documentSymbol - nested struct with self" {
     try testDocumentSymbol(
         \\const Foo = struct {
         \\    const Self = @This();
         \\    pub fn foo() !Self {}
-        \\    const Bar = struct {};
+        \\    const Bar = union {};
         \\};
     ,
-        \\Variable Foo
+        \\Constant Foo
-        \\  Variable Self
+        \\  Constant Self
         \\  Function foo
-        \\  Variable Bar
+        \\  Constant Bar
     );
 }