1 files changed, 109 insertions, 14 deletions
diff --git a/src/syntax.zig b/src/syntax.zig
index 9e31669f..e637bf67 100644
--- a/src/syntax.zig
+++ b/src/syntax.zig
@@ -376,28 +376,58 @@ pub fn supportsPath(path: []const u8) bool {
     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 &.{};
+/// Owned source analysis. Both slices use the allocator passed to analyzeSource.
+/// The parsed tree is released before returning; callers can cache this result.
+pub const SourceAnalysis = struct {
+    declarations: []ContextDeclaration = &.{},
+    colors: []u8 = &.{},
+
+    pub fn deinit(self: *SourceAnalysis, gpa: std.mem.Allocator) void {
+        gpa.free(self.declarations);
+        gpa.free(self.colors);
+        self.* = .{};
+    }
+};
+
+pub fn analyzeSource(gpa: std.mem.Allocator, path: []const u8, content: []const u8, want_declarations: bool, want_colors: bool) !SourceAnalysis {
+    if (!enabled or content.len == 0 or (!want_declarations and !want_colors)) 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: SourceAnalysis = .{};
+    errdefer result.deinit(gpa);
+    if (want_declarations) result.declarations = try declarationsFromTree(gpa, content, tree);
+    if (want_colors) {
+        result.colors = try gpa.alloc(u8, content.len);
+        @memset(result.colors, 0);
+        paintTree(result.colors, content, selected, tree);
+    }
+    return result;
+}
+
+pub fn contextDeclarations(gpa: std.mem.Allocator, path: []const u8, content: []const u8) ![]ContextDeclaration {
+    const analysis = try analyzeSource(gpa, path, content, true, false);
+    return analysis.declarations;
+}
+
+fn declarationsFromTree(gpa: std.mem.Allocator, content: []const u8, tree: *ts.Tree) ![]ContextDeclaration {
     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();
+    // The cursor retains its ancestor stack: sibling visits do not repeatedly
+    // reconstruct parents in broad syntax trees. Preorder preserves source order.
+    var cursor = tree.rootNode().walk();
+    defer cursor.destroy();
     while (true) {
-        if (contextSpan(node, content)) |span| {
+        const node = cursor.node();
+        if (node.isNamed()) 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 (cursor.gotoFirstChild()) continue;
+        while (!cursor.gotoNextSibling()) {
+            if (!cursor.gotoParent()) return result.toOwnedSlice(gpa);
         }
-        if (node.namedChild(0)) |child| {
-            node = child;
-            continue;
-        }
-        while (node.nextNamedSibling() == null) node = node.parent() orelse return result.toOwnedSlice(gpa);
-        node = node.nextNamedSibling().?;
     }
 }
 
@@ -492,6 +522,10 @@ fn contextSpan(node: ts.Node, source: []const u8) ?ContextDeclaration {
 fn paint(styles: []u8, source: []const u8, selected: Selected) void {
     const tree = selected.parser.parseString(source, null) orelse return;
     defer tree.destroy();
+    paintTree(styles, source, selected, tree);
+}
+
+fn paintTree(styles: []u8, source: []const u8, selected: Selected, tree: *ts.Tree) void {
     runQuery(styles, selected, tree, 0);
     if (std.mem.eql(u8, selected.name, "markdown")) {
         inject(styles, source, tree.rootNode(), true);
@@ -1279,3 +1313,64 @@ test "syntax stacked addresses do not split multiline source groups" {
         for (styles[at..][0..label.len]) |style| try std.testing.expectEqual(@intFromEnum(Syn.none), style);
     }
 }
+
+test "syntax source analysis preserves declarations and injected colors" {
+    if (!enabled) return;
+    const gpa = std.testing.allocator;
+    start(gpa);
+    defer stop();
+    const Fixture = struct { path: []const u8, source: []const u8 };
+    for ([_]Fixture{
+        .{ .path = "a.zig", .source = "const Outer = struct {\n // comment\n const Inner = struct {\n pub fn run(\n x: u32,\n ) u32 { return x; }\n };\n};\n" },
+        .{ .path = "a.cpp", .source = "namespace Outer {\nstruct Inner {\nint run(\n int x\n) { return x; }\n};\n}\n" },
+        .{ .path = "a.rs", .source = "mod outer {\nstruct Inner {}\nimpl Inner {\nfn run(\n &self\n) {}\n}\n}\n" },
+        .{ .path = "a.js", .source = "function outer() {\nclass Inner {\nrun(\n x\n) { return x; }\n}\n}\n" },
+        .{ .path = "a.py", .source = "class Outer:\n    class Inner:\n        def run(\n            self, x\n        ):\n            return x\n" },
+        .{ .path = "a.md", .source = "# Heading\n\n```zig\npub fn run() void {}\n```\n" },
+        .{ .path = "a.typst", .source = "= Heading\n\n```zig\npub fn run() void {}\n```\n" },
+    }) |fixture| {
+        if (!supportsPath(fixture.path)) continue;
+        var analysis = try analyzeSource(gpa, fixture.path, fixture.source, true, true);
+        defer analysis.deinit(gpa);
+        const colors = try highlightFileRange(gpa, fixture.path, fixture.source, 0, fixture.source.len);
+        defer gpa.free(colors);
+        try std.testing.expectEqualSlices(u8, colors, analysis.colors);
+
+        // Compare the previous named-node walk to the cursor walk, including
+        // source ordering, wrapper deduplication and inclusive declaration ends.
+        const selected = (try forExt(std.fs.path.extension(fixture.path))).?;
+        const tree = selected.parser.parseString(fixture.source, null).?;
+        defer tree.destroy();
+        var reference: std.ArrayList(ContextDeclaration) = .empty;
+        defer reference.deinit(gpa);
+        var node = tree.rootNode();
+        walk: while (true) {
+            if (contextSpan(node, fixture.source)) |span| {
+                if (reference.items.len == 0 or reference.items[reference.items.len - 1].start_line != span.start_line or
+                    reference.items[reference.items.len - 1].end_line != span.end_line)
+                    try reference.append(gpa, span);
+            }
+            if (node.namedChild(0)) |child| {
+                node = child;
+                continue;
+            }
+            while (node.nextNamedSibling() == null) node = node.parent() orelse break :walk;
+            node = node.nextNamedSibling().?;
+        }
+        try std.testing.expectEqual(reference.items.len, analysis.declarations.len);
+        for (reference.items, analysis.declarations) |old, new| try std.testing.expect(std.meta.eql(old, new));
+
+        var declarations_only = try analyzeSource(gpa, fixture.path, fixture.source, true, false);
+        defer declarations_only.deinit(gpa);
+        try std.testing.expectEqual(@as(usize, 0), declarations_only.colors.len);
+        try std.testing.expectEqual(analysis.declarations.len, declarations_only.declarations.len);
+        var colors_only = try analyzeSource(gpa, fixture.path, fixture.source, false, true);
+        defer colors_only.deinit(gpa);
+        try std.testing.expectEqual(@as(usize, 0), colors_only.declarations.len);
+        try std.testing.expectEqualSlices(u8, analysis.colors, colors_only.colors);
+    }
+    var unsupported = try analyzeSource(gpa, "a.unknown", "text", true, true);
+    unsupported.deinit(gpa);
+    try std.testing.expectEqual(@as(usize, 0), unsupported.declarations.len);
+    try std.testing.expectEqual(@as(usize, 0), unsupported.colors.len);
+}