From cfb0b023ad83e63074f4761ac11fd470678f49f2 Mon Sep 17 00:00:00 2001
From: Alex Kladov <aleksey.kladov@gmail.com>
Date: Sat, 3 Dec 2022 15:23:13 +0000
Subject: [PATCH] fix #801, IOOB in foldingRanges (#802)

* Add smoke tests for folding ranges

* fix index out of bounds in foldingRanges

closes #801

For invalid syntax trees, zig's parser seems to return bogus data where
startToken > endToken, which then causes everything else to crash.

This seems like a deeper issue, which needs to be fixed "properly", but
let's just paper over it here.
---
 src/Server.zig                       |  2 +-
 tests/lsp_features/folding_range.zig | 66 ++++++++++++++++++++++++++++
 tests/tests.zig                      |  5 ++-
 3 files changed, 70 insertions(+), 3 deletions(-)
 create mode 100644 tests/lsp_features/folding_range.zig

diff --git a/src/Server.zig b/src/Server.zig
index fcf6102..7aabb99 100644
--- a/src/Server.zig
+++ b/src/Server.zig
@@ -2454,7 +2454,7 @@ fn foldingRangeHandler(server: *Server, writer: anytype, id: types.RequestId, re
             end: Ast.TokenIndex,
             end_reach: Inclusivity,
         ) std.mem.Allocator.Error!bool {
-            const can_add = !tree.tokensOnSameLine(start, end);
+            const can_add = start < end and !tree.tokensOnSameLine(start, end);
             if (can_add) {
                 try addTokRange(p_ranges, tree, start, end, end_reach);
             }
diff --git a/tests/lsp_features/folding_range.zig b/tests/lsp_features/folding_range.zig
new file mode 100644
index 0000000..956d4f1
--- /dev/null
+++ b/tests/lsp_features/folding_range.zig
@@ -0,0 +1,66 @@
+const std = @import("std");
+const zls = @import("zls");
+const builtin = @import("builtin");
+
+const Context = @import("../context.zig").Context;
+
+const types = zls.types;
+const requests = zls.requests;
+
+const allocator: std.mem.Allocator = std.testing.allocator;
+
+test "foldingRange - empty" {
+    try testFoldingRange("", "[]");
+}
+
+test "foldingRange - smoke" {
+    try testFoldingRange(
+        \\fn main() u32 {
+        \\    return 1 + 1;
+        \\}
+    ,
+        \\[{"startLine":0,"endLine":1}]
+    );
+}
+
+test "foldingRange - #801" {
+    try testFoldingRange(
+        \\fn score(c: u8) !u32 {
+        \\    return switch(c) {
+        \\        'a'...'z' => c - 'a',
+        \\        'A'...'Z' => c - 'A',
+        \\        _ => error
+        \\    };
+        \\}
+    ,
+        \\[]
+    );
+}
+
+fn testFoldingRange(source: []const u8, expect: []const u8) !void {
+    var ctx = try Context.init();
+    defer ctx.deinit();
+
+    const test_uri: []const u8 = switch (builtin.os.tag) {
+        .windows => "file:///C:\\test.zig",
+        else => "file:///test.zig",
+    };
+
+    try ctx.requestDidOpen(test_uri, source);
+
+    const request = requests.FoldingRange{ .params = .{ .textDocument = .{ .uri = test_uri } } };
+
+    const response = try ctx.requestGetResponse(?[]types.FoldingRange, "textDocument/foldingRange", request);
+    defer response.deinit();
+
+    var actual = std.ArrayList(u8).init(allocator);
+    defer actual.deinit();
+
+    try std.json.stringify(response.result, .{}, actual.writer());
+    try expectEqualJson(expect, actual.items);
+}
+
+fn expectEqualJson(expect: []const u8, actual: []const u8) !void {
+    // TODO: Actually compare strings as JSON values.
+    return std.testing.expectEqualStrings(expect, actual);
+}
diff --git a/tests/tests.zig b/tests/tests.zig
index 1951513..d85c824 100644
--- a/tests/tests.zig
+++ b/tests/tests.zig
@@ -10,11 +10,12 @@ comptime {
     // TODO Document Synchronization
 
     // LSP features
-    _ = @import("lsp_features/semantic_tokens.zig");
+    _ = @import("lsp_features/completion.zig");
+    _ = @import("lsp_features/folding_range.zig");
     _ = @import("lsp_features/inlay_hints.zig");
     _ = @import("lsp_features/references.zig");
-    _ = @import("lsp_features/completion.zig");
     _ = @import("lsp_features/selection_range.zig");
+    _ = @import("lsp_features/semantic_tokens.zig");
 
     // Language features
     _ = @import("language_features/cimport.zig");