Add optional tree-sitter declaration context

1 files changed, 217 insertions, 0 deletions

diff --git a/src/syntax.zig b/src/syntax.zig
index 77ef80d7..69a69bda 100644
--- a/src/syntax.zig
+++ b/src/syntax.zig
@@ -355,6 +355,140 @@ pub fn highlightFileRange(gpa: std.mem.Allocator, path: []const u8, content: []c
     return styles;
 }
 
+/// Source spans are independent of editor rows. Locations output uses the same
+/// declarations as sticky editor headers. Byte ends are exclusive; rows are
+/// zero-based and inclusive, including the last signature/opening-brace row.
+pub const ContextDeclaration = struct {
+    start_byte: usize,
+    end_byte: usize,
+    start_line: usize,
+    end_line: usize,
+    header_end_byte: usize,
+    header_end_line: usize,
+};
+
+pub fn supportsPath(path: []const u8) bool {
+    if (!enabled) return false;
+    const ext = std.fs.path.extension(path);
+    for (specs) |spec| for (spec.exts) |choice| {
+        if (std.ascii.eqlIgnoreCase(ext, choice)) return true;
+    };
+    return false;
+}
+
+pub fn contextDeclarations(gpa: std.mem.Allocator, path: []const u8, content: []const u8) ![]ContextDeclaration {
+    if (!enabled or content.len == 0) return &.{};
+    const selected = (try forExt(std.fs.path.extension(path))) orelse return &.{};
+    const tree = selected.parser.parseString(content, null) orelse return &.{};
+    defer tree.destroy();
+    var result: std.ArrayList(ContextDeclaration) = .empty;
+    errdefer result.deinit(gpa);
+    // Iterative preorder keeps declaration order and avoids recursive traversal
+    // on deeply nested, partly edited source.
+    var node = tree.rootNode();
+    while (true) {
+        if (contextSpan(node, content)) |span| {
+            if (result.items.len == 0 or result.items[result.items.len - 1].start_line != span.start_line or
+                result.items[result.items.len - 1].end_line != span.end_line)
+                try result.append(gpa, span);
+        }
+        if (node.namedChild(0)) |child| {
+            node = child;
+            continue;
+        }
+        while (node.nextNamedSibling() == null) node = node.parent() orelse return result.toOwnedSlice(gpa);
+        node = node.nextNamedSibling().?;
+    }
+}
+
+fn contextSpan(node: ts.Node, source: []const u8) ?ContextDeclaration {
+    const kind = node.kind();
+    var body: ?ts.Node = null;
+    var accepted = false;
+    if (std.mem.eql(u8, kind, "Decl")) {
+        // Zig's FnProto excludes its body: the surrounding Decl owns both.
+        accepted = childOfKind(node, "FnProto") != null;
+        body = childOfKind(node, "Block");
+        if (body == null) return null;
+    } else if (std.mem.eql(u8, kind, "ContainerDecl") or std.mem.eql(u8, kind, "TestDecl")) {
+        accepted = true;
+        body = childOfKind(node, "Block");
+    } else {
+        const kinds = [_][]const u8{
+            "function_definition",   "function_declaration", "function_item",         "function",                "method_definition",
+            "method_declaration",    "method",               "singleton_method",      "constructor_declaration", "destructor_definition",
+            "class_definition",      "class_declaration",    "class_specifier",       "class",                   "singleton_class",
+            "object_declaration",    "struct_item",          "struct_specifier",      "struct_declaration",      "union_specifier",
+            "union_item",            "union_declaration",    "enum_item",             "enum_specifier",          "enum_declaration",
+            "interface_declaration", "trait_item",           "trait_definition",      "impl_item",               "mod_item",
+            "module_definition",     "module_declaration",   "module",                "namespace_definition",    "namespace_declaration",
+            "record_declaration",    "type_declaration",     "extension_declaration", "protocol_declaration",
+        };
+        for (kinds) |candidate| if (std.mem.eql(u8, kind, candidate)) {
+            accepted = true;
+            break;
+        };
+        body = node.childByFieldName("body");
+    }
+    if (!accepted) return null;
+    var owner = node;
+    if (std.mem.eql(u8, kind, "ContainerDecl")) {
+        // A multiline `const Name = struct` starts at the binding, not at the
+        // struct token. Stop at real scopes so returned anonymous containers
+        // do not inherit an enclosing function's signature.
+        var ancestor = node.parent();
+        while (ancestor) |parent| {
+            const parent_kind = parent.kind();
+            if (std.mem.eql(u8, parent_kind, "VarDecl")) {
+                owner = parent;
+                break;
+            }
+            if (std.mem.eql(u8, parent_kind, "Block") or std.mem.eql(u8, parent_kind, "Decl") or
+                std.mem.eql(u8, parent_kind, "ContainerDecl") or std.mem.eql(u8, parent_kind, "FnProto") or
+                std.mem.eql(u8, parent_kind, "ContainerField") or std.mem.eql(u8, parent_kind, "ParamDecl")) break;
+            ancestor = parent.parent();
+        }
+    }
+    const beginning = owner.startPoint();
+    const end = node.endPoint();
+    if (beginning.row >= end.row) return null;
+    var header_end: usize = node.startByte();
+    if (body) |block| {
+        header_end = block.startByte();
+        // Indentation grammars start their body at its first statement.
+        if (header_end < source.len and source[header_end] == '{') {
+            header_end += 1;
+        } else {
+            while (header_end > node.startByte() and std.ascii.isWhitespace(source[header_end - 1])) header_end -= 1;
+        }
+    } else {
+        // Containers often expose their braces directly, without a body node.
+        var i: u32 = 0;
+        while (i < node.childCount()) : (i += 1) {
+            const child = node.child(i) orelse continue;
+            if (std.mem.eql(u8, child.kind(), "{")) {
+                header_end = child.endByte();
+                break;
+            }
+        }
+        // Ruby and similar grammars have a body_statement without a field.
+        if (header_end == node.startByte()) {
+            header_end = std.mem.indexOfScalarPos(u8, source, header_end, '\n') orelse source.len;
+        }
+    }
+    const line_start = @as(usize, owner.startByte()) - beginning.column;
+    const header_line = beginning.row + std.mem.count(u8, source[line_start..@min(header_end, source.len)], "\n");
+    const header_line_end = std.mem.indexOfScalarPos(u8, source, @min(header_end, source.len), '\n') orelse source.len;
+    return .{
+        .start_byte = line_start,
+        .end_byte = node.endByte(),
+        .start_line = beginning.row,
+        .end_line = end.row - @as(usize, if (end.column == 0) 1 else 0),
+        .header_end_byte = header_line_end,
+        .header_end_line = header_line,
+    };
+}
+
 fn paint(styles: []u8, source: []const u8, selected: Selected) void {
     const tree = selected.parser.parseString(source, null) orelse return;
     defer tree.destroy();
@@ -366,6 +500,89 @@ fn paint(styles: []u8, source: []const u8, selected: Selected) void {
     }
 }
 
+test "syntax context nested Zig containers and multiline function signatures" {
+    if (!enabled) return;
+    const gpa = std.testing.allocator;
+    start(gpa);
+    defer stop();
+    const source =
+        "pub const Outer = struct {\n" ++
+        "    const Inner = struct {\n" ++
+        "        pub fn run(\n" ++
+        "            value: u32,\n" ++
+        "        ) void {\n" ++
+        "            const ignored = .{\n" ++
+        "                value,\n" ++
+        "            };\n" ++
+        "        }\n" ++
+        "    };\n" ++
+        "};\n";
+    const declarations = try contextDeclarations(gpa, "nested.zig", source);
+    defer gpa.free(declarations);
+    try std.testing.expectEqual(@as(usize, 3), declarations.len);
+    try std.testing.expectEqual(@as(usize, 0), declarations[0].start_line);
+    try std.testing.expectEqual(@as(usize, 10), declarations[0].end_line);
+    try std.testing.expectEqual(@as(usize, 1), declarations[1].start_line);
+    try std.testing.expectEqual(@as(usize, 2), declarations[2].start_line);
+    try std.testing.expectEqual(@as(usize, 4), declarations[2].header_end_line);
+    try std.testing.expectEqual(@as(usize, 8), declarations[2].end_line);
+    try std.testing.expectEqualStrings("        pub fn run(\n            value: u32,\n        ) void {", source[declarations[2].start_byte..declarations[2].header_end_byte]);
+    const unsupported = try contextDeclarations(gpa, "notes.unknown", source);
+    defer gpa.free(unsupported);
+    try std.testing.expectEqual(@as(usize, 0), unsupported.len);
+}
+
+test "syntax context Python class and function body excludes first statement" {
+    if (!enabled or !full_grammars) return;
+    const gpa = std.testing.allocator;
+    start(gpa);
+    defer stop();
+    const source = "class Outer:\n    def run(\n        self, value,\n    ):\n        return value\n";
+    const declarations = try contextDeclarations(gpa, "nested.py", source);
+    defer gpa.free(declarations);
+    try std.testing.expectEqual(@as(usize, 2), declarations.len);
+    try std.testing.expectEqual(@as(usize, 0), declarations[0].header_end_line);
+    try std.testing.expectEqual(@as(usize, 3), declarations[1].header_end_line);
+}
+
+test "syntax context multiline bindings and C++ namespaces preserve declaration starts" {
+    if (!enabled) return;
+    const gpa = std.testing.allocator;
+    start(gpa);
+    defer stop();
+    const zig_source = "const Outer =\n    struct {\n        field: u8,\n    };\n";
+    const zig_declarations = try contextDeclarations(gpa, "nested.zig", zig_source);
+    defer gpa.free(zig_declarations);
+    try std.testing.expectEqual(@as(usize, 1), zig_declarations.len);
+    try std.testing.expectEqual(@as(usize, 0), zig_declarations[0].start_line);
+    try std.testing.expectEqual(@as(usize, 1), zig_declarations[0].header_end_line);
+    if (!minimal_grammars and !full_grammars) return;
+    const cpp_source = "namespace example {\nstruct Outer {\n    int run() {\n        return 1;\n    }\n};\n}\n";
+    const cpp_declarations = try contextDeclarations(gpa, "nested.cpp", cpp_source);
+    defer gpa.free(cpp_declarations);
+    try std.testing.expectEqual(@as(usize, 3), cpp_declarations.len);
+    for (cpp_declarations, 0..) |declaration, row| {
+        try std.testing.expectEqual(row, declaration.start_line);
+        try std.testing.expectEqual(row, declaration.header_end_line);
+    }
+}
+
+test "syntax context Rust modules impls and methods remain nested" {
+    if (!enabled or !full_grammars) return;
+    const gpa = std.testing.allocator;
+    start(gpa);
+    defer stop();
+    const source = "mod outer {\n    impl Example {\n        fn run(&self) {\n            work();\n        }\n    }\n}\n";
+    const declarations = try contextDeclarations(gpa, "nested.rs", source);
+    defer gpa.free(declarations);
+    try std.testing.expectEqual(@as(usize, 3), declarations.len);
+    for (declarations, 0..) |declaration, row| {
+        try std.testing.expectEqual(row, declaration.start_line);
+        try std.testing.expectEqual(row, declaration.header_end_line);
+        try std.testing.expectEqual(@as(usize, 6) - row, declaration.end_line);
+    }
+}
+
 pub fn highlightLocations(gpa: std.mem.Allocator, content: []const u8, start_byte_raw: usize, end_byte_raw: usize) ![]u8 {
     const tz = tracy.zone(@src(), "highlightLocations");
     defer tz.end();