selectionRange: include function parameters in the result

zigtools · Jan 14, 2024 · b5c756c · b5c756c
1 parent abe83cf
commit b5c756c
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git a/src/features/selection_range.zig b/src/features/selection_range.zig
@@ -2,6 +2,7 @@ const std = @import("std");
 const Ast = std.zig.Ast;
 
 const DocumentStore = @import("../DocumentStore.zig");
+const ast = @import("../ast.zig");
 const types = @import("../lsp.zig");
 const offsets = @import("../offsets.zig");
 
@@ -24,11 +25,33 @@ pub fn generateSelectionRanges(
         const index = offsets.positionToIndex(handle.tree.source, position, offset_encoding);
 
         locs.clearRetainingCapacity();
-        for (0..handle.tree.nodes.len) |i| {
+        for (0..handle.tree.nodes.len, handle.tree.nodes.items(.tag)) |i, tag| {
             const node = @as(Ast.Node.Index, @intCast(i));
             const loc = offsets.nodeToLoc(handle.tree, node);
-            if (loc.start <= index and index <= loc.end) {
-                try locs.append(arena, loc);
+
+            if (!(loc.start <= index and index <= loc.end)) continue;
+
+            try locs.append(arena, loc);
+            switch (tag) {
+                // Function parameters are not stored in the AST explicitly, iterate over them
+                // manually.
+                .fn_proto, .fn_proto_multi, .fn_proto_one, .fn_proto_simple => {
+                    var buffer: [1]Ast.Node.Index = undefined;
+                    const fn_proto = handle.tree.fullFnProto(&buffer, node).?;
+                    var it = fn_proto.iterate(&handle.tree);
+
+                    while (ast.nextFnParam(&it)) |param| {
+                        const loc = offsets.tokensToLoc(
+                            handle.tree,
+                            ast.paramFirstToken(handle.tree, param),
+                            ast.paramLastToken(handle.tree, param),
+                        );
+                        if (loc.start <= index and index <= loc.end) {
+                            try locs.append(arena, loc);
+                        }
+                    }
+                },
+                else => {},
             }
         }
 

diff --git a/tests/lsp_features/selection_range.zig b/tests/lsp_features/selection_range.zig
@@ -15,14 +15,22 @@ test "selectionRange - empty" {
     try testSelectionRange("<>", &.{});
 }
 
-test "seletionRange - smoke" {
+test "selectionRange - smoke" {
     try testSelectionRange(
         \\fn main() void {
         \\    const x = 1 <>+ 1;
         \\}
     , &.{ "1 + 1", "const x = 1 + 1", "{\n    const x = 1 + 1;\n}" });
 }
 
+test "selectionRange - function parameter" {
+    try testSelectionRange(
+        \\fn f(x: i32, y: <>struct {}, z: f32) void {
+        \\
+        \\}
+    , &.{ "struct {}", "y: struct {}", "fn f(x: i32, y: struct {}, z: f32) void" });
+}
+
 fn testSelectionRange(source: []const u8, want: []const []const u8) !void {
     var phr = try helper.collectClearPlaceholders(allocator, source);
     defer phr.deinit(allocator);