Show result declaration locations only on their first header line

4 files changed, 28 insertions, 11 deletions

diff --git a/docs/config.md b/docs/config.md
index be8a28de..202e58d8 100644
--- a/docs/config.md
+++ b/docs/config.md
@@ -155,8 +155,9 @@ LocationsConfig context:5 tscontext:on tslocations:off
   locations and keep their code aligned with the matching result.
 - `tscontext` includes enclosing tree-sitter declaration headers in source
   order (default `off`).
-- `tslocations` shows locations beside those declaration headers (default
-  `on`). Declaration headers use the same muted color as locations, with or
+- `tslocations` shows a location on the first line of each declaration header
+  (default `on`); continuation lines keep the same alignment without repeating
+  the location. Declaration headers use the same muted color as locations, with or
   without their locations visible.
 
 Result locations share padding in groups of eight matches, so a long path only
diff --git a/src/locations.zig b/src/locations.zig
index 45d72c57..34d1ad14 100644
--- a/src/locations.zig
+++ b/src/locations.zig
@@ -5,6 +5,9 @@ const syntax = @import("syntax.zig");
 const filesystem = @import("fs.zig");
 
 pub const Config = @import("locations_config.zig").Config;
+const neighbor_context: u8 = 1;
+const declaration_context: u8 = 2;
+const declaration_start: u8 = 4;
 /// Legacy marked rows and hidden locations remain recognizable without metadata.
 /// Visible context locations are flush left and use Row.kind instead.
 pub fn isContextLine(line: []const u8) bool {
@@ -194,24 +197,27 @@ pub fn format(p: *pardes.Pardes, dir: []const u8, input: []const u8, anchor: ?us
                     var context_row = lo;
                     while (context_row <= hi) : (context_row += 1) {
                         if (context_row == row) continue;
-                        wanted[context_row] |= 1;
+                        wanted[context_row] |= neighbor_context;
                     }
                     for (declarations) |declaration| {
                         if (declaration.start_line > row) break;
                         if (declaration.end_line < row) continue;
                         match.depth += 1;
                         var header = declaration.start_line;
-                        while (header <= declaration.header_end_line and header < row) : (header += 1)
-                            wanted[header] |= 2;
+                        while (header <= declaration.header_end_line and header < row) : (header += 1) {
+                            wanted[header] |= declaration_context;
+                            if (header == declaration.start_line) wanted[header] |= declaration_start;
+                        }
                     }
                 }
                 for (wanted, 0..) |kind, row| {
                     if (kind > 0) {
                         const line = source_lines.items[row];
                         const offset = @intFromPtr(line.ptr) - @intFromPtr(source.ptr);
-                        const declaration = kind & 2 != 0;
+                        const declaration = kind & declaration_context != 0;
                         const colors = if (!declaration and offset + line.len <= source_colors.len) source_colors[offset..][0..line.len] else &.{};
-                        try appendContext(arena, &pending, path, line, row, declaration, !declaration or !p.locations_config.tslocations, colors);
+                        const hidden = !declaration or !p.locations_config.tslocations or kind & declaration_start == 0;
+                        try appendContext(arena, &pending, path, line, row, declaration, hidden, colors);
                     }
                 }
             }
diff --git a/test/output.zig b/test/output.zig
index 128d6ff5..722f2e44 100644
--- a/test/output.zig
+++ b/test/output.zig
@@ -1323,7 +1323,10 @@ test "declaration context ellipses mark only final header lines before omitted s
                         try std.testing.expect(row.location_end > 0);
                         if (row.at.line == 16) try std.testing.expectEqual(index, source.search_row.?);
                     } else if (row.declaration) {
-                        try std.testing.expectEqual(visible, row.location_end > 0);
+                        const first_line = row.at.line == 1 or row.at.line == 3 or row.at.line == 7 or row.at.line == 8;
+                        try std.testing.expectEqual(visible and first_line, row.location_end > 0);
+                        try std.testing.expectEqualStrings(source.file.?.path, row.path);
+                        if (!first_line) try std.testing.expect(std.mem.indexOf(u8, shown[0..row.code_start], row.path) == null);
                     } else try std.testing.expectEqual(@as(usize, 0), row.location_end);
                 }
                 try std.testing.expectEqual(@as(usize, if (neighbors == 0) 3 else 2), marked);
diff --git a/test/panes.zig b/test/panes.zig
index da06f692..d844f26e 100644
--- a/test/panes.zig
+++ b/test/panes.zig
@@ -2322,6 +2322,7 @@ const OutputTests = struct {
             if (row.kind == .match) {
                 matches += 1;
                 try std.testing.expect(!row.declaration);
+                try std.testing.expect(row.location_end > 0);
                 const location = panes.Output.location(line);
                 try std.testing.expectEqualStrings(path, location.path);
                 try std.testing.expectEqual(row.at, location.at);
@@ -2332,9 +2333,15 @@ const OutputTests = struct {
                 }
             } else if (row.at.line <= 5) {
                 try std.testing.expect(row.declaration);
-                try std.testing.expect(row.location_end > 0);
-                try std.testing.expect(std.mem.startsWith(u8, line, path));
-                try std.testing.expectEqualStrings(path, panes.Output.location(line).path);
+                try std.testing.expectEqualStrings(path, row.path);
+                if (row.at.line <= 3) {
+                    try std.testing.expect(row.location_end > 0);
+                    try std.testing.expect(std.mem.startsWith(u8, line, path));
+                    try std.testing.expectEqualStrings(path, panes.Output.location(line).path);
+                } else {
+                    try std.testing.expectEqual(@as(usize, 0), row.location_end);
+                    try std.testing.expect(std.mem.indexOf(u8, line[0..row.code_start], path) == null);
+                }
             } else {
                 try std.testing.expect(!row.declaration);
                 try std.testing.expectEqual(@as(usize, 0), row.location_end);