rewrite folding range (#954)

src/Server.zig

@@ -14,6 +14,7 @@ const offsets = @import("offsets.zig");
 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 shared = @import("shared.zig");
 const Ast = std.zig.Ast;
 const tracy = @import("tracy.zig");
@@ -2579,250 +2580,10 @@ fn foldingRangeHandler(server: *Server, request: types.FoldingRangeParams) Error
     const tracy_zone = tracy.trace(@src());
     defer tracy_zone.end();
 
-    const Token = std.zig.Token;
-    const Node = Ast.Node;
-    const allocator = server.arena.allocator();
     const handle = server.document_store.getHandle(request.textDocument.uri) orelse return null;
+    const allocator = server.arena.allocator();
 
-    const helper = struct {
+    return try folding_range.generateFoldingRanges(allocator, handle.tree, server.offset_encoding);
-        const Inclusivity = enum { inclusive, exclusive };
-
-        fn addTokRange(
-            p_ranges: *std.ArrayList(types.FoldingRange),
-            tree: Ast,
-            start: Ast.TokenIndex,
-            end: Ast.TokenIndex,
-            end_reach: Inclusivity,
-        ) std.mem.Allocator.Error!void {
-            if (tree.tokensOnSameLine(start, end)) return;
-            std.debug.assert(start <= end);
-
-            const start_index = offsets.tokenToIndex(tree, start);
-            const end_index = offsets.tokenToIndex(tree, end);
-
-            const start_line = std.mem.count(u8, tree.source[0..start_index], "\n");
-            const end_line = start_line + std.mem.count(u8, tree.source[start_index..end_index], "\n");
-
-            try p_ranges.append(.{
-                .startLine = @intCast(u32, start_line),
-                .endLine = @intCast(u32, end_line) - @boolToInt(end_reach == .exclusive),
-            });
-        }
-    };
-
-    // Used to store the result
-    var ranges = std.ArrayList(types.FoldingRange).init(allocator);
-
-    const token_tags: []const Token.Tag = handle.tree.tokens.items(.tag);
-    const node_tags: []const Node.Tag = handle.tree.nodes.items(.tag);
-
-    if (token_tags.len == 0) return null;
-    if (token_tags[0] == .container_doc_comment) {
-        var tok: Ast.TokenIndex = 1;
-        while (tok < token_tags.len) : (tok += 1) {
-            if (token_tags[tok] != .container_doc_comment) {
-                break;
-            }
-        }
-        if (tok > 1) { // each container doc comment has its own line, so each one counts for a line
-            try ranges.append(.{
-                .startLine = 0,
-                .endLine = tok - 1,
-            });
-        }
-    }
-
-    for (node_tags) |node_tag, i| {
-        const node = @intCast(Node.Index, i);
-
-        switch (node_tag) {
-            .root => continue,
-            // only fold the expression pertaining to the if statement, and the else statement, each respectively.
-            // TODO: Should folding multiline condition expressions also be supported? Ditto for the other control flow structures.
-            .@"if",
-            .if_simple,
-            => {
-                const if_full = ast.fullIf(handle.tree, node).?;
-
-                const start_tok_1 = ast.lastToken(handle.tree, if_full.ast.cond_expr);
-                const end_tok_1 = ast.lastToken(handle.tree, if_full.ast.then_expr);
-                try helper.addTokRange(&ranges, handle.tree, start_tok_1, end_tok_1, .inclusive);
-
-                if (if_full.ast.else_expr == 0) continue;
-
-                const start_tok_2 = if_full.else_token;
-                const end_tok_2 = ast.lastToken(handle.tree, if_full.ast.else_expr);
-
-                try helper.addTokRange(&ranges, handle.tree, start_tok_2, end_tok_2, .inclusive);
-            },
-
-            // same as if/else
-            .@"for",
-            .for_simple,
-            .@"while",
-            .while_cont,
-            .while_simple,
-            => {
-                const loop_full = ast.fullWhile(handle.tree, node).?;
-
-                const start_tok_1 = ast.lastToken(handle.tree, loop_full.ast.cond_expr);
-                const end_tok_1 = ast.lastToken(handle.tree, loop_full.ast.then_expr);
-                try helper.addTokRange(&ranges, handle.tree, start_tok_1, end_tok_1, .inclusive);
-
-                if (loop_full.ast.else_expr == 0) continue;
-
-                const start_tok_2 = loop_full.else_token;
-                const end_tok_2 = ast.lastToken(handle.tree, loop_full.ast.else_expr);
-                try helper.addTokRange(&ranges, handle.tree, start_tok_2, end_tok_2, .inclusive);
-            },
-
-            .global_var_decl,
-            .simple_var_decl,
-            .aligned_var_decl,
-            .container_field_init,
-            .container_field_align,
-            .container_field,
-            .fn_proto,
-            .fn_proto_multi,
-            .fn_proto_one,
-            .fn_proto_simple,
-            .fn_decl,
-            => decl_node_blk: {
-                doc_comment_range: {
-                    const first_tok: Ast.TokenIndex = handle.tree.firstToken(node);
-                    if (first_tok == 0) break :doc_comment_range;
-
-                    const end_doc_tok = first_tok - 1;
-                    if (token_tags[end_doc_tok] != .doc_comment) break :doc_comment_range;
-
-                    var start_doc_tok = end_doc_tok;
-                    while (start_doc_tok != 0) {
-                        if (token_tags[start_doc_tok - 1] != .doc_comment) break;
-                        start_doc_tok -= 1;
-                    }
-
-                    try helper.addTokRange(&ranges, handle.tree, start_doc_tok, end_doc_tok, .inclusive);
-                }
-
-                // Function prototype folding regions
-                var buffer: [1]Node.Index = undefined;
-                const fn_proto = handle.tree.fullFnProto(&buffer, node) orelse
-                    break :decl_node_blk;
-
-                const list_start_tok: Ast.TokenIndex = fn_proto.lparen;
-                const list_end_tok: Ast.TokenIndex = ast.lastToken(handle.tree, fn_proto.ast.proto_node);
-
-                if (handle.tree.tokensOnSameLine(list_start_tok, list_end_tok)) break :decl_node_blk;
-                try helper.addTokRange(&ranges, handle.tree, list_start_tok, list_end_tok, .exclusive);
-
-                var it = fn_proto.iterate(&handle.tree);
-                while (ast.nextFnParam(&it)) |param| {
-                    const doc_start_tok = param.first_doc_comment orelse continue;
-                    var doc_end_tok = doc_start_tok;
-
-                    while (token_tags[doc_end_tok + 1] == .doc_comment)
-                        doc_end_tok += 1;
-
-                    try helper.addTokRange(&ranges, handle.tree, doc_start_tok, doc_end_tok, .inclusive);
-                }
-            },
-
-            .@"catch",
-            .@"orelse",
-            .multiline_string_literal,
-            // TODO: Similar to condition expressions in control flow structures, should folding multiline grouped expressions be enabled?
-            // .grouped_expression,
-            => {
-                const start_tok = handle.tree.firstToken(node);
-                const end_tok = ast.lastToken(handle.tree, node);
-                try helper.addTokRange(&ranges, handle.tree, start_tok, end_tok, .inclusive);
-            },
-
-            // most other trivial cases can go through here.
-            else => {
-                switch (node_tag) {
-                    .array_init,
-                    .array_init_one,
-                    .array_init_dot_two,
-                    .array_init_one_comma,
-                    .array_init_dot_two_comma,
-                    .array_init_dot,
-                    .array_init_dot_comma,
-                    .array_init_comma,
-
-                    .struct_init,
-                    .struct_init_one,
-                    .struct_init_one_comma,
-                    .struct_init_dot_two,
-                    .struct_init_dot_two_comma,
-                    .struct_init_dot,
-                    .struct_init_dot_comma,
-                    .struct_init_comma,
-
-                    .@"switch",
-                    .switch_comma,
-                    => {},
-
-                    else => disallow_fold: {
-                        if (ast.isBlock(handle.tree, node))
-                            break :disallow_fold;
-
-                        if (ast.isCall(handle.tree, node))
-                            break :disallow_fold;
-
-                        if (ast.isBuiltinCall(handle.tree, node))
-                            break :disallow_fold;
-
-                        if (ast.isContainer(handle.tree, node) and node_tag != .root)
-                            break :disallow_fold;
-
-                        continue; // no conditions met, continue iterating without adding this potential folding range
-                    },
-                }
-
-                const start_tok = handle.tree.firstToken(node);
-                const end_tok = ast.lastToken(handle.tree, node);
-                try helper.addTokRange(&ranges, handle.tree, start_tok, end_tok, .exclusive);
-            },
-        }
-    }
-
-    // Iterate over the source code and look for code regions with #region #endregion
-    {
-        // We add opened folding regions to a stack as we go and pop one off when we find a closing brace.
-        // As an optimization we start with a reasonable capacity, which should work well in most cases since
-        // people will almost never have nesting that deep.
-        var stack = try std.ArrayList(u32).initCapacity(allocator, 10);
-
-        var i: usize = 0;
-        var lines_count: u32 = 0;
-        while (i < handle.tree.source.len) : (i += 1) {
-            const slice = handle.tree.source[i..];
-
-            if (slice[0] == '\n') {
-                lines_count += 1;
-            }
-
-            if (std.mem.startsWith(u8, slice, "//#region")) {
-                try stack.append(lines_count);
-            }
-
-            if (std.mem.startsWith(u8, slice, "//#endregion") and stack.items.len > 0) {
-                const start_line = stack.pop();
-                const end_line = lines_count;
-
-                // Add brace pairs but discard those from the same line, no need to waste memory on them
-                if (start_line != end_line) {
-                    try ranges.append(.{
-                        .startLine = start_line,
-                        .endLine = end_line,
-                    });
-                }
-            }
-        }
-    }
-
-    return ranges.items;
 }
 
 pub const SelectionRange = struct {

src/ast.zig

@@ -976,32 +976,6 @@ pub fn isBuiltinCall(tree: Ast, node: Ast.Node.Index) bool {
     };
 }
 
-pub fn isCall(tree: Ast, node: Ast.Node.Index) bool {
-    return switch (tree.nodes.items(.tag)[node]) {
-        .call,
-        .call_comma,
-        .call_one,
-        .call_one_comma,
-        .async_call,
-        .async_call_comma,
-        .async_call_one,
-        .async_call_one_comma,
-        => true,
-        else => false,
-    };
-}
-
-pub fn isBlock(tree: Ast, node: Ast.Node.Index) bool {
-    return switch (tree.nodes.items(.tag)[node]) {
-        .block_two,
-        .block_two_semicolon,
-        .block,
-        .block_semicolon,
-        => true,
-        else => false,
-    };
-}
-
 /// returns a list of parameters
 pub fn builtinCallParams(tree: Ast, node: Ast.Node.Index, buf: *[2]Ast.Node.Index) ?[]const Node.Index {
     const node_data = tree.nodes.items(.data);

src/folding_range.zig

@@ -0,0 +1,311 @@
+const std = @import("std");
+const ast = @import("ast.zig");
+const types = @import("lsp.zig");
+const offsets = @import("offsets.zig");
+const Ast = std.zig.Ast;
+
+const FoldingRange = struct {
+    loc: offsets.Loc,
+    kind: ?types.FoldingRangeKind = null,
+};
+
+const Inclusivity = enum { inclusive, exclusive };
+
+const Builder = struct {
+    allocator: std.mem.Allocator,
+    locations: std.ArrayListUnmanaged(FoldingRange),
+    tree: Ast,
+    encoding: offsets.Encoding,
+
+    pub fn deinit(builder: *Builder) void {
+        builder.locations.deinit(builder.allocator);
+    }
+
+    pub fn add(
+        builder: *Builder,
+        kind: ?types.FoldingRangeKind,
+        start: Ast.TokenIndex,
+        end: Ast.TokenIndex,
+        start_reach: Inclusivity,
+        end_reach: Inclusivity,
+    ) error{OutOfMemory}!void {
+        if (builder.tree.tokensOnSameLine(start, end)) return;
+        std.debug.assert(start <= end);
+        const start_loc = offsets.tokenToLoc(builder.tree, start);
+        const end_loc = offsets.tokenToLoc(builder.tree, end);
+
+        try builder.locations.append(builder.allocator, .{
+            .loc = .{
+                .start = if (start_reach == .exclusive) start_loc.end else start_loc.start,
+                .end = if (end_reach == .exclusive) end_loc.start else end_loc.end,
+            },
+            .kind = kind,
+        });
+    }
+
+    pub fn addNode(
+        builder: *Builder,
+        kind: ?types.FoldingRangeKind,
+        node: Ast.Node.Index,
+        start_reach: Inclusivity,
+        end_reach: Inclusivity,
+    ) error{OutOfMemory}!void {
+        try builder.add(kind, builder.tree.firstToken(node), ast.lastToken(builder.tree, node), start_reach, end_reach);
+    }
+
+    pub fn getRanges(builder: Builder) error{OutOfMemory}![]types.FoldingRange {
+        var result = try builder.allocator.alloc(types.FoldingRange, builder.locations.items.len);
+        errdefer builder.allocator.free(result);
+
+        for (result) |*r, i| {
+            r.* = .{
+                .startLine = undefined,
+                .endLine = undefined,
+                .kind = builder.locations.items[i].kind,
+            };
+        }
+
+        const Item = struct {
+            output: *types.FoldingRange,
+            input: *const FoldingRange,
+            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();
+            }
+        };
+
+        // one item for every start and end position
+        var items = try builder.allocator.alloc(Item, builder.locations.items.len * 2);
+        defer builder.allocator.free(items);
+
+        for (builder.locations.items) |*folding_range, i| {
+            items[2 * i + 0] = .{ .output = &result[i], .input = folding_range, .where = .start };
+            items[2 * i + 1] = .{ .output = &result[i], .input = folding_range, .where = .end };
+        }
+
+        // sort items based on their source position
+        std.sort.sort(Item, items, {}, Item.lessThan);
+
+        var last_index: usize = 0;
+        var last_position: types.Position = .{ .line = 0, .character = 0 };
+        for (items) |item| {
+            const index = item.getInputIndex();
+            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) {
+                .start => {
+                    item.output.startLine = position.line;
+                    item.output.startCharacter = position.character;
+                },
+                .end => {
+                    item.output.endLine = position.line;
+                    item.output.endCharacter = position.character;
+                },
+            }
+        }
+
+        return result;
+    }
+};
+
+pub fn generateFoldingRanges(allocator: std.mem.Allocator, tree: Ast, encoding: offsets.Encoding) error{OutOfMemory}![]types.FoldingRange {
+    var builder = Builder{
+        .allocator = allocator,
+        .locations = .{},
+        .tree = tree,
+        .encoding = encoding,
+    };
+    defer builder.deinit();
+
+    const token_tags = tree.tokens.items(.tag);
+    const node_tags = tree.nodes.items(.tag);
+    const main_tokens = tree.nodes.items(.main_token);
+
+    var start_doc_comment: ?Ast.TokenIndex = null;
+    var end_doc_comment: ?Ast.TokenIndex = null;
+    for (token_tags) |tag, i| {
+        const token = @intCast(Ast.TokenIndex, i);
+        switch (tag) {
+            .doc_comment,
+            .container_doc_comment,
+            => {
+                if (start_doc_comment == null) {
+                    start_doc_comment = token;
+                    end_doc_comment = token;
+                } else {
+                    end_doc_comment = token;
+                }
+            },
+            else => {
+                if (start_doc_comment != null and end_doc_comment != null) {
+                    try builder.add(.comment, start_doc_comment.?, end_doc_comment.?, .inclusive, .inclusive);
+                    start_doc_comment = null;
+                    end_doc_comment = null;
+                }
+            },
+        }
+    }
+
+    // TODO add folding range normal comments
+
+    // TODO add folding range for top level `@Import()`
+
+    for (node_tags) |node_tag, i| {
+        const node = @intCast(Ast.Node.Index, i);
+
+        switch (node_tag) {
+            .root => continue,
+            // TODO: Should folding multiline condition expressions also be supported? Ditto for the other control flow structures.
+
+            .fn_proto,
+            .fn_proto_multi,
+            .fn_proto_one,
+            .fn_proto_simple,
+            // .fn_decl
+            => {
+                var buffer: [1]Ast.Node.Index = undefined;
+                const fn_proto = tree.fullFnProto(&buffer, node).?;
+
+                const list_start_tok = fn_proto.lparen;
+                const list_end_tok = ast.lastToken(tree, node) -| 1;
+
+                try builder.add(null, list_start_tok, list_end_tok, .exclusive, .exclusive);
+            },
+
+            .block_two,
+            .block_two_semicolon,
+            .block,
+            .block_semicolon,
+            => {
+                try builder.addNode(null, node, .exclusive, .exclusive);
+            },
+            .@"switch",
+            .switch_comma,
+            => {
+                const lhs = tree.nodes.items(.data)[node].lhs;
+                const start_tok = ast.lastToken(tree, lhs) + 2; // lparen + rbrace
+                const end_tok = ast.lastToken(tree, node);
+                try builder.add(null, start_tok, end_tok, .exclusive, .exclusive);
+            },
+
+            .switch_case_one,
+            .switch_case_inline_one,
+            .switch_case,
+            .switch_case_inline,
+            => {
+                const switch_case = tree.fullSwitchCase(node).?.ast;
+                if (switch_case.values.len >= 4) {
+                    const first_value = tree.firstToken(switch_case.values[0]);
+                    const last_value = ast.lastToken(tree, switch_case.values[switch_case.values.len - 1]);
+                    try builder.add(null, first_value, last_value, .inclusive, .inclusive);
+                }
+            },
+
+            .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_two,
+            .tagged_union_two_trailing,
+            .tagged_union_enum_tag,
+            .tagged_union_enum_tag_trailing,
+            => {
+                var buffer: [2]Ast.Node.Index = undefined;
+                const container_decl = tree.fullContainerDecl(&buffer, node).?;
+                if (container_decl.ast.members.len != 0) {
+                    const first_member = container_decl.ast.members[0];
+                    const start_tok = tree.firstToken(first_member) -| 1;
+                    const end_tok = ast.lastToken(tree, node);
+                    try builder.add(null, start_tok, end_tok, .exclusive, .exclusive);
+                }
+            },
+
+            .call,
+            .call_comma,
+            .call_one,
+            .call_one_comma,
+            .async_call,
+            .async_call_comma,
+            .async_call_one,
+            .async_call_one_comma,
+            => {
+                const lparen = main_tokens[node];
+                try builder.add(null, lparen, ast.lastToken(tree, node), .exclusive, .exclusive);
+            },
+
+            // everything after here is mostly untested
+            .array_init,
+            .array_init_one,
+            .array_init_dot_two,
+            .array_init_one_comma,
+            .array_init_dot_two_comma,
+            .array_init_dot,
+            .array_init_dot_comma,
+            .array_init_comma,
+
+            .struct_init,
+            .struct_init_one,
+            .struct_init_one_comma,
+            .struct_init_dot_two,
+            .struct_init_dot_two_comma,
+            .struct_init_dot,
+            .struct_init_dot_comma,
+            .struct_init_comma,
+
+            .builtin_call,
+            .builtin_call_comma,
+            .builtin_call_two,
+            .builtin_call_two_comma,
+
+            .multiline_string_literal,
+            .error_set_decl,
+            .test_decl,
+            => {
+                try builder.addNode(null, node, .inclusive, .inclusive);
+            },
+
+            else => {},
+        }
+    }
+
+    // We add opened folding regions to a stack as we go and pop one off when we find a closing brace.
+    var stack = std.ArrayListUnmanaged(usize){};
+
+    var i: usize = 0;
+    while (std.mem.indexOfPos(u8, tree.source, i, "//#")) |possible_region| {
+        defer i = possible_region + "//#".len;
+        if (std.mem.startsWith(u8, tree.source[possible_region..], "//#region")) {
+            try stack.append(allocator, possible_region);
+        } else if (std.mem.startsWith(u8, tree.source[possible_region..], "//#endregion")) {
+            const start_index = stack.popOrNull() orelse break; // null means there are more endregions than regions
+            const end_index = offsets.lineLocAtIndex(tree.source, possible_region).end;
+            const is_same_line = std.mem.indexOfScalar(u8, tree.source[start_index..end_index], '\n') == null;
+            if (is_same_line) continue;
+            try builder.locations.append(allocator, .{
+                .loc = .{
+                    .start = start_index,
+                    .end = end_index,
+                },
+                .kind = .region,
+            });
+        }
+    }
+
+    return try builder.getRanges();
+}

tests/lsp_features/folding_range.zig

@@ -11,34 +11,197 @@ const types = zls.types;
 const allocator: std.mem.Allocator = std.testing.allocator;
 
 test "foldingRange - empty" {
-    try testFoldingRange("", "[]");
+    try testFoldingRange("", &.{});
 }
 
-test "foldingRange - smoke" {
+test "foldingRange - doc comment" {
+    try testFoldingRange(
+        \\/// hello
+        \\/// world
+        \\var foo = 5;
+    , &.{
+        .{ .startLine = 0, .startCharacter = 0, .endLine = 1, .endCharacter = 9, .kind = .comment },
+    });
+}
+
+test "foldingRange - region" {
+    try testFoldingRange(
+        \\const foo = 0;
+        \\//#region
+        \\const bar = 1;
+        \\//#endregion
+        \\const baz = 2;
+    , &.{
+        .{ .startLine = 1, .startCharacter = 0, .endLine = 3, .endCharacter = 12, .kind = .region },
+    });
+    try testFoldingRange(
+        \\//#region
+        \\const foo = 0;
+        \\//#region
+        \\const bar = 1;
+        \\//#endregion
+        \\const baz = 2;
+        \\//#endregion
+    , &.{
+        .{ .startLine = 2, .startCharacter = 0, .endLine = 4, .endCharacter = 12, .kind = .region },
+        .{ .startLine = 0, .startCharacter = 0, .endLine = 6, .endCharacter = 12, .kind = .region },
+    });
+}
+
+test "foldingRange - if" {
+    try testFoldingRange(
+        \\const foo = if (false) {
+        \\
+        \\};
+    , &.{
+        .{ .startLine = 0, .startCharacter = 24, .endLine = 2, .endCharacter = 0 },
+    });
+    try testFoldingRange(
+        \\const foo = if (false) {
+        \\
+        \\} else {
+        \\
+        \\};
+    , &.{
+        .{ .startLine = 0, .startCharacter = 24, .endLine = 2, .endCharacter = 0 },
+        .{ .startLine = 2, .startCharacter = 8, .endLine = 4, .endCharacter = 0 },
+    });
+}
+
+test "foldingRange - for/while" {
+    try testFoldingRange(
+        \\const foo = for ("") |_| {
+        \\
+        \\};
+    , &.{
+        .{ .startLine = 0, .startCharacter = 26, .endLine = 2, .endCharacter = 0 },
+    });
+    try testFoldingRange(
+        \\const foo = while (true) {
+        \\
+        \\};
+    , &.{
+        .{ .startLine = 0, .startCharacter = 26, .endLine = 2, .endCharacter = 0 },
+    });
+}
+
+test "foldingRange - switch" {
+    try testFoldingRange(
+        \\const foo = switch (5) {
+        \\  0 => {},
+        \\  1 => {}
+        \\};
+    , &.{
+        .{ .startLine = 0, .startCharacter = 24, .endLine = 3, .endCharacter = 0 },
+    });
+    try testFoldingRange(
+        \\const foo = switch (5) {
+        \\  0 => {},
+        \\  1 => {},
+        \\};
+    , &.{
+        .{ .startLine = 0, .startCharacter = 24, .endLine = 3, .endCharacter = 0 },
+    });
+}
+
+test "foldingRange - function" {
     try testFoldingRange(
         \\fn main() u32 {
         \\    return 1 + 1;
         \\}
-    ,
+    , &.{
-        \\[{"startLine":0,"endLine":1}]
+        .{ .startLine = 0, .startCharacter = 15, .endLine = 2, .endCharacter = 0 },
-    );
+    });
-}
-
-test "foldingRange - #801" {
     try testFoldingRange(
-        \\fn score(c: u8) !u32 {
+        \\fn main(
-        \\    return switch(c) {
+        \\  a: ?u32,
-        \\        'a'...'z' => c - 'a',
+        \\) u32 {
-        \\        'A'...'Z' => c - 'A',
+        \\    return 1 + 1;
-        \\        _ => error
-        \\    };
         \\}
-    ,
+    , &.{
-        \\[{"startLine":1,"endLine":4},{"startLine":0,"endLine":5}]
+        .{ .startLine = 0, .startCharacter = 8, .endLine = 2, .endCharacter = 0 },
-    );
+        .{ .startLine = 2, .startCharacter = 7, .endLine = 4, .endCharacter = 0 },
+    });
 }
 
-fn testFoldingRange(source: []const u8, expect: []const u8) !void {
+test "foldingRange - function with doc comment" {
+    try testFoldingRange(
+        \\/// this is
+        \\/// a function
+        \\fn foo(
+        \\    /// this is a parameter
+        \\    a: u32,
+        \\    ///
+        \\    /// this is another parameter
+        \\    b: u32,
+        \\) void {}
+    , &.{
+        .{ .startLine = 0, .startCharacter = 0, .endLine = 1, .endCharacter = 14, .kind = .comment },
+        .{ .startLine = 5, .startCharacter = 4, .endLine = 6, .endCharacter = 33, .kind = .comment },
+        .{ .startLine = 2, .startCharacter = 7, .endLine = 8, .endCharacter = 0 },
+    });
+}
+
+test "foldingRange - container decl" {
+    try testFoldingRange(
+        \\const Foo = struct {
+        \\  alpha: u32,
+        \\  beta: []const u8,
+        \\};
+    , &.{
+        .{ .startLine = 0, .startCharacter = 20, .endLine = 3, .endCharacter = 0 },
+    });
+    try testFoldingRange(
+        \\const Foo = packed struct(u32) {
+        \\  alpha: u16,
+        \\  beta: u16,
+        \\};
+    , &.{
+        // .{ .startLine = 0, .startCharacter = 32, .endLine = 3, .endCharacter = 0 }, // TODO
+        .{ .startLine = 0, .startCharacter = 32, .endLine = 2, .endCharacter = 11 },
+    });
+    try testFoldingRange(
+        \\const Foo = union {
+        \\  alpha: u32,
+        \\  beta: []const u8,
+        \\};
+    , &.{
+        .{ .startLine = 0, .startCharacter = 19, .endLine = 3, .endCharacter = 0 },
+    });
+    try testFoldingRange(
+        \\const Foo = union(enum) {
+        \\  alpha: u32,
+        \\  beta: []const u8,
+        \\};
+    , &.{
+        .{ .startLine = 0, .startCharacter = 25, .endLine = 3, .endCharacter = 0 },
+    });
+}
+
+test "foldingRange - call" {
+    try testFoldingRange(
+        \\extern fn foo(a: bool, b: ?usize) void;
+        \\const result = foo(
+        \\    false,
+        \\    null,  
+        \\);
+    , &.{
+        .{ .startLine = 1, .startCharacter = 19, .endLine = 4, .endCharacter = 0 },
+    });
+}
+
+test "foldingRange - multi-line string literal" {
+    try testFoldingRange(
+        \\const foo =
+        \\    \\hello
+        \\    \\world
+        \\;
+    , &.{
+        .{ .startLine = 1, .startCharacter = 4, .endLine = 3, .endCharacter = 0 },
+    });
+}
+
+fn testFoldingRange(source: []const u8, expect: []const types.FoldingRange) !void {
     var ctx = try Context.init();
     defer ctx.deinit();
 
@@ -53,16 +216,16 @@ fn testFoldingRange(source: []const u8, expect: []const u8) !void {
 
     const response = try ctx.requestGetResponse(?[]types.FoldingRange, "textDocument/foldingRange", params);
 
-    var actual = std.ArrayList(u8).init(allocator);
+    var actual = std.ArrayListUnmanaged(u8){};
-    defer actual.deinit();
+    defer actual.deinit(allocator);
 
-    try tres.stringify(response.result, .{
+    var expected = std.ArrayListUnmanaged(u8){};
-        .emit_null_optional_fields = false,
+    defer expected.deinit(allocator);
-    }, actual.writer());
+
-    try expectEqualJson(expect, actual.items);
+    const options = std.json.StringifyOptions{ .emit_null_optional_fields = false, .whitespace = .{ .indent = .None } };
-}
+    try tres.stringify(response.result, options, actual.writer(allocator));
+    try tres.stringify(expect, options, expected.writer(allocator));
 
-fn expectEqualJson(expect: []const u8, actual: []const u8) !void {
     // TODO: Actually compare strings as JSON values.
-    return std.testing.expectEqualStrings(expect, actual);
+    try std.testing.expectEqualStrings(expected.items, actual.items);
 }