From 4b0da6f6aec80471d812a06028a546c395f114ab Mon Sep 17 00:00:00 2001
From: Techarix <19954306+Techatrix@users.noreply.github.com>
Date: Sat, 25 Feb 2023 01:35:00 +0100
Subject: [PATCH] fix folding range on multi line function declaration

---
 src/folding_range.zig                | 14 +++++++++-----
 tests/lsp_features/folding_range.zig | 11 +++++++----
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/src/folding_range.zig b/src/folding_range.zig
index 857fac1..8b95792 100644
--- a/src/folding_range.zig
+++ b/src/folding_range.zig
@@ -178,12 +178,16 @@ pub fn generateFoldingRanges(allocator: std.mem.Allocator, tree: Ast, encoding:
                 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;
+                if (fn_proto.ast.params.len != 0) {
+                    const list_start_tok = fn_proto.lparen;
+                    const last_param = fn_proto.ast.params[fn_proto.ast.params.len - 1];
+                    const last_param_tok = ast.lastToken(tree, last_param);
+                    const param_has_comma = last_param_tok + 1 < tree.tokens.len and token_tags[last_param_tok + 1] == .comma;
+                    const list_end_tok = last_param_tok + @boolToInt(param_has_comma);
 
-                if (list_start_tok > list_end_tok) continue; // Incomplete, ie `fn a()`
-
-                try builder.add(null, list_start_tok, list_end_tok, .exclusive, .exclusive);
+                    if (list_start_tok > list_end_tok) continue; // Incomplete, ie `fn a()`
+                    try builder.add(null, list_start_tok, list_end_tok, .exclusive, .inclusive);
+                }
             },
 
             .block_two,
diff --git a/tests/lsp_features/folding_range.zig b/tests/lsp_features/folding_range.zig
index 8b6b63b..8d57336 100644
--- a/tests/lsp_features/folding_range.zig
+++ b/tests/lsp_features/folding_range.zig
@@ -115,12 +115,15 @@ test "foldingRange - function" {
     try testFoldingRange(
         \\fn main(
         \\  a: ?u32,
-        \\) u32 {
+        \\) !u32 {
         \\    return 1 + 1;
         \\}
     , &.{
-        .{ .startLine = 0, .startCharacter = 8, .endLine = 2, .endCharacter = 0 },
-        .{ .startLine = 2, .startCharacter = 7, .endLine = 4, .endCharacter = 0 },
+        .{ .startLine = 0, .startCharacter = 8, .endLine = 1, .endCharacter = 10 },
+        .{ .startLine = 2, .startCharacter = 8, .endLine = 4, .endCharacter = 0 },
+        // TODO investigate why VS-Code doesn't show the second folding range
+        // .{ .startLine = 0, .startCharacter = 8, .endLine = 2, .endCharacter = 0 },
+        // .{ .startLine = 2, .startCharacter = 8, .endLine = 4, .endCharacter = 0 },
     });
 }
 
@@ -138,7 +141,7 @@ test "foldingRange - function with doc comment" {
     , &.{
         .{ .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 },
+        .{ .startLine = 2, .startCharacter = 7, .endLine = 7, .endCharacter = 11 },
     });
 }