Highlight exact source occurrences in LSP references and goto results

5 files changed, 89 insertions, 18 deletions

diff --git a/src/lsp/lsp.zig b/src/lsp/lsp.zig
index ca19009a..9dd1da5a 100644
--- a/src/lsp/lsp.zig
+++ b/src/lsp/lsp.zig
@@ -85,10 +85,15 @@ pub fn row(out: *std.Io.Writer, path: []const u8, line: usize, col: usize, text:
 
 // Input positions are zero-based and end-exclusive; displayed spans are one-based and inclusive.
 pub fn spanRow(out: *std.Io.Writer, path: []const u8, line: usize, col: usize, end_line: usize, end_col: usize, text: []const u8) std.Io.Writer.Error!void {
-    if (end_line != line or end_col <= col) return row(out, path, line, col, text);
-    try out.print("{s}:{d}:{d}-{d} {s}\n", .{
-        path, line + 1, col + 1, end_col, std.mem.trim(u8, text, " \t\r\n"),
-    });
+    return sourceSpanRow(out, path, line, col, end_line, end_col, std.mem.trim(u8, text, " \t\r\n"));
+}
+
+/// Source bytes keep their indentation and trailing spaces so byte columns
+/// still identify the same token in the displayed snippet.
+pub fn sourceSpanRow(out: *std.Io.Writer, path: []const u8, line: usize, col: usize, end_line: usize, end_col: usize, text: []const u8) std.Io.Writer.Error!void {
+    try out.print("{s}:{d}:{d}", .{ path, line + 1, col + 1 });
+    if (end_line == line and end_col > col) try out.print("-{d}", .{end_col});
+    try out.print(" {s}\n", .{std.mem.trimEnd(u8, text, "\r\n")});
 }
 
 pub fn edit(out: *std.Io.Writer, start: usize, end: usize) std.Io.Writer.Error!void {
@@ -114,9 +119,10 @@ pub fn lineCol(source: []const u8, offset: usize) struct { line: usize, col: usi
 }
 
 test "LSP encoders report every insufficient output capacity" {
-    const cases = [_]struct { kind: enum { row, span, edit, put }, expected: []const u8 }{
+    const cases = [_]struct { kind: enum { row, span, source_span, edit, put }, expected: []const u8 }{
         .{ .kind = .row, .expected = "file:1:3 hi\n" },
         .{ .kind = .span, .expected = "file:1:3-5 hi\n" },
+        .{ .kind = .source_span, .expected = "file:1:3-5 \t hi \t\n" },
         .{ .kind = .edit, .expected = "@edit 1 3\n" },
         .{ .kind = .put, .expected = "@put 1 3 hé%20%25%0A\n" },
     };
@@ -127,6 +133,7 @@ test "LSP encoders report every insufficient output capacity" {
             const result = switch (case.kind) {
                 .row => row(&out, "file", 0, 2, " hi \n"),
                 .span => spanRow(&out, "file", 0, 2, 0, 5, " hi \n"),
+                .source_span => sourceSpanRow(&out, "file", 0, 2, 0, 5, "\t hi \t\r\n"),
                 .edit => edit(&out, 1, 3),
                 .put => put(&out, 1, 3, "hé %\n"),
             };
diff --git a/src/lsp/lsp_client.zig b/src/lsp/lsp_client.zig
index 8c6a0da9..c3762f08 100644
--- a/src/lsp/lsp_client.zig
+++ b/src/lsp/lsp_client.zig
@@ -948,15 +948,13 @@ fn emitRange(cx: *Cx, enc: Enc, uri: []const u8, range: ?std.json.Value, note: [
 
     const col = colBytes(lntext, r.sc, enc);
     const shown = lsp.rel(cx.base, path);
-    const rowtext = if (note.len > 0) note else lntext;
-    if (r.el == r.sl and r.ec > r.sc) {
-        // spanRow wants the protocol's EXCLUSIVE end as a 1-based inclusive
-        // byte column; converting the exclusive utf-16 end unit yields the
-        // exclusive byte column, which is the same number.
-        const end_col = colBytes(lntext, r.ec, enc);
-        try lsp.spanRow(cx.out, shown, r.sl, col, r.el, end_col, rowtext);
+    // A same-line exclusive end byte offset is also the displayed inclusive
+    // one-based column. Multi-line ranges display only their starting address.
+    const end_col = if (r.el == r.sl and r.ec > r.sc) colBytes(lntext, r.ec, enc) else col;
+    if (note.len == 0) {
+        try lsp.sourceSpanRow(cx.out, shown, r.sl, col, r.el, end_col, lntext);
     } else {
-        try lsp.row(cx.out, shown, r.sl, col, rowtext);
+        try lsp.spanRow(cx.out, shown, r.sl, col, r.el, end_col, note);
     }
     cx.rows += 1;
 }
@@ -2208,3 +2206,26 @@ test "the transport this host actually gives us is a pair, and both ends are clo
     try std.testing.expectEqual(@as(isize, msg.len), libc.read(sv[1], &got, got.len));
     try std.testing.expectEqualStrings(msg, got[0..msg.len]);
 }
+
+test "LSP source snippets preserve indentation with UTF16 byte ranges and prose stays trimmed" {
+    const gpa = std.testing.allocator;
+    const source = "\t  const 𝕏 = target; \t\r\n";
+    for ([_]bool{ false, true }) |multiline| {
+        var arena: std.heap.ArenaAllocator = .init(gpa);
+        defer arena.deinit();
+        const json = try std.fmt.allocPrint(arena.allocator(), "{{\"start\":{{\"line\":0,\"character\":14}},\"end\":{{\"line\":{d},\"character\":20}}}}", .{@as(u8, if (multiline) 1 else 0)});
+        const parsed = try std.json.parseFromSlice(std.json.Value, arena.allocator(), json, .{});
+        for ([_]bool{ false, true }) |prose| {
+            var buffer: [256]u8 = undefined;
+            var out: std.Io.Writer = .fixed(&buffer);
+            var cx: Cx = .{ .arena = arena.allocator(), .base = "/", .cur_path = "/file.c", .cur_src = source, .out = &out };
+            try emitRange(&cx, .utf16, "file:///file.c", parsed.value, if (prose) "  explanatory label \t" else "");
+            const expected = try std.fmt.allocPrint(arena.allocator(), "file.c:1:17{s} {s}\n", .{
+                if (multiline) "" else "-22",
+                if (prose) "explanatory label" else "\t  const 𝕏 = target; \t",
+            });
+            try std.testing.expectEqualStrings(expected, out.buffered());
+            try std.testing.expectEqual(@as(usize, 1), cx.rows);
+        }
+    }
+}
diff --git a/src/lsp/lsp_zls.zig b/src/lsp/lsp_zls.zig
index 47e6a165..68538125 100644
--- a/src/lsp/lsp_zls.zig
+++ b/src/lsp/lsp_zls.zig
@@ -512,7 +512,7 @@ fn rowForToken(arena: std.mem.Allocator, base: []const u8, th: Analyser.TokenWit
     if (th.token >= tree.tokens.len) return;
     const r = offsets.tokenToRange(tree, th.token, enc);
     const path = lsp.rel(base, th.handle.uri.toFsPath(arena) catch return);
-    try lsp.spanRow(out, path, r.start.line, r.start.character, r.end.line, r.end.character, lineAt(tree.source, r.start.line));
+    try lsp.sourceSpanRow(out, path, r.start.line, r.start.character, r.end.line, r.end.character, lineAt(tree.source, r.start.line));
 }
 
 // ---------------------------------------------------------------- goto
@@ -667,7 +667,7 @@ fn goto(
                 const t2 = &nd.handle.tree;
                 const rr = offsets.nodeToRange(t2, nd.node, enc);
                 const path = lsp.rel(base, nd.handle.uri.toFsPath(arena) catch continue);
-                try lsp.spanRow(out, path, rr.start.line, rr.start.character, rr.end.line, rr.end.character, lineAt(t2.source, rr.start.line));
+                try lsp.sourceSpanRow(out, path, rr.start.line, rr.start.character, rr.end.line, rr.end.character, lineAt(t2.source, rr.start.line));
             },
         }
     }
@@ -1165,7 +1165,7 @@ fn references(
             try lsp.edit(out, at, at + want.len);
         } else {
             const r = offsets.tokenToRange(tree, tok, enc);
-            try lsp.spanRow(out, path, r.start.line, r.start.character, r.end.line, r.end.character, lines.?.line(r.start.line));
+            try lsp.sourceSpanRow(out, path, r.start.line, r.start.character, r.end.line, r.end.character, lines.?.line(r.start.line));
         }
     }
 }
@@ -1474,3 +1474,31 @@ test "LSP ZLS rename refuses a partial edit set beyond its row budget" {
         }
     }
 }
+
+test "LSP source snippets retain native definition and reference indentation" {
+    const source = "const Outer = struct {\n" ++
+        "\t  const value: u32 = 1;\n" ++
+        "\t  fn use() void {\n" ++
+        "\t\t_ = value;\n" ++
+        "\t  }\n" ++
+        "};\n";
+    const gpa = std.testing.allocator;
+    for ([_]lsp.Kind{ .definition, .declaration, .references }) |kind| {
+        var arena: std.heap.ArenaAllocator = .init(gpa);
+        defer arena.deinit();
+        var out: std.Io.Writer.Allocating = .init(gpa);
+        defer out.deinit();
+        try query(gpa, arena.allocator(), .{
+            .kind = kind,
+            .path = "/source-snippet.zig",
+            .root = "/",
+            .source = source,
+            .offset = @intCast(std.mem.lastIndexOf(u8, source, "value").?),
+        }, &out.writer);
+        const definition = "source-snippet.zig:2:10-14 \t  const value: u32 = 1;\n";
+        try std.testing.expectEqualStrings(if (kind == .references)
+            definition ++ "source-snippet.zig:4:7-11 \t\t_ = value;\n"
+        else
+            definition, out.written());
+    }
+}
diff --git a/src/main.zig b/src/main.zig
index 2d7b6538..2ee295e1 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -379,6 +379,8 @@ test {
     _ = @import("host_io.zig").Shell;
     _ = @import("host_io.zig").Lsp;
     _ = @import("lsp/lsp_client.zig");
+    _ = @import("lsp/lsp.zig");
+    if (comptime @import("pardes_config").zls_backend) _ = @import("lsp/lsp_zls.zig");
     if (comptime pardes.platform == .tty) {
         _ = @import("tty/tty.zig");
         _ = @import("tty/panel_compositor.zig");
diff --git a/src/panes.zig b/src/panes.zig
index 315b0b0f..7855866d 100644
--- a/src/panes.zig
+++ b/src/panes.zig
@@ -2205,12 +2205,16 @@ pub const Output = struct {
         const source_start = if (metadata) |m| m.code_start else target.code_start;
         if (source_start >= line.len) return row;
         if (metadata) |m| if (m.kind == .context) return row;
-        const is_search = switch (origin) {
+        const mark_occurrence = switch (origin) {
             .search => true,
             .cmd => |cmd| cmd == .Grep,
+            .query => |kind| switch (kind) {
+                .definition, .declaration, .type_definition, .implementation, .references => true,
+                else => false,
+            },
             else => false,
         };
-        if (is_search and spot.col > 0 and spot.end_line == spot.line and
+        if (mark_occurrence and spot.col > 0 and spot.end_line == spot.line and
             spot.end_col >= spot.col)
         {
             row.mark_start = source_start +| (spot.col - 1);
@@ -2243,6 +2247,15 @@ pub const Output = struct {
         try std.testing.expectEqualStrings("warning:", "src/a.rs:3:5-7 warning: unused"[diagnostic.mark_start..diagnostic.mark_end]);
         const malformed = rowDecoration(.{ .from = .search }, "src/a.rs:3:99999999999999999999-99999999999999999999 text");
         try std.testing.expect(malformed.mark_end <= "src/a.rs:3:99999999999999999999-99999999999999999999 text".len);
+        const snippet = "src/a.rs:3:14-16 \t\tfoo + 界 + foo";
+        for ([_]lsp.Kind{ .definition, .declaration, .type_definition, .implementation, .references }) |kind| {
+            const occurrence = rowDecoration(.{ .from = .{ .query = kind } }, snippet);
+            try std.testing.expectEqualStrings("foo", snippet[occurrence.mark_start..occurrence.mark_end]);
+        }
+        for ([_]lsp.Kind{ .document_symbols, .workspace_symbols, .completion, .incoming_calls, .outgoing_calls, .supertypes, .subtypes, .rename }) |kind| {
+            const label = rowDecoration(.{ .from = .{ .query = kind } }, snippet);
+            try std.testing.expectEqual(@as(usize, 0), label.mark_end);
+        }
     }
 
     test "plain output does not opt into location decoration" {