path: root/src/syntax.zig

//! Tree-sitter syntax highlighting: one style byte per content byte, filled by
//! running each grammar's highlights.scm query (slurped at build time into the
//! ts_queries options module). Grammar set is tiered: `zig`; `minimal` (c, cpp,
//! zig); and `full`, which adds ~23 languages lazily on first use.
const std = @import("std");
const config = @import("pardes_config");
const tracy = @import("tracy.zig");
const grammar_manifest = @import("grammar_manifest.zig");

pub const enabled = config.syntax_highlighting;
const zig_grammar = config.syntax_zig_grammar;
const minimal_grammars = config.syntax_minimal_grammars;
const full_grammars = config.syntax_full_grammars;

const ts = if (enabled) @import("tree-sitter") else struct {
    pub const Language = opaque {};
    pub const Query = opaque {};
};
const ts_queries = if (enabled) @import("ts_queries") else struct {};

const LanguageFn = *const fn () callconv(.c) *const ts.Language;

pub const Syn = enum(u8) { none, keyword, string, number, comment };

const Spec = struct {
    exts: []const []const u8,
    language: *const fn () callconv(.c) *const ts.Language,
    query_src: []const u8,
    compiled_query: ?*ts.Query = null,
};

fn grammarSelected(comptime g: grammar_manifest.Grammar) bool {
    return switch (g.tier) {
        .zig => zig_grammar,
        .minimal => minimal_grammars,
        .full => full_grammars,
    };
}

fn specCount() comptime_int {
    var count = 0;
    for (grammar_manifest.all) |g| {
        if (grammarSelected(g)) count += 1;
    }
    return count;
}

fn initSpecs() [specCount()]Spec {
    var out: [specCount()]Spec = undefined;
    var i = 0;
    inline for (grammar_manifest.all) |g| {
        if (grammarSelected(g)) {
            out[i] = .{
                .exts = g.exts,
                .language = @extern(LanguageFn, .{ .name = "tree_sitter_" ++ g.name }),
                .query_src = @field(ts_queries, g.name ++ "_highlights"),
            };
            i += 1;
        }
    }
    return out;
}

var specs = initSpecs();

const allocation_header_size = 16;
var syntax_allocator: std.mem.Allocator = undefined;
var syntax_started = false;

extern fn ts_set_allocator(
    new_malloc: ?*const fn (size: usize) callconv(.c) ?*anyopaque,
    new_calloc: ?*const fn (nmemb: usize, size: usize) callconv(.c) ?*anyopaque,
    new_realloc: ?*const fn (ptr: ?*anyopaque, size: usize) callconv(.c) ?*anyopaque,
    new_free: ?*const fn (ptr: ?*anyopaque) callconv(.c) void,
) void;

fn syntaxAlloc(size_arg: usize) callconv(.c) ?*anyopaque {
    const size = @max(size_arg, 1);
    const total = std.math.add(usize, allocation_header_size, size) catch return null;
    const bytes = syntax_allocator.alignedAlloc(u8, .@"16", total) catch return null;
    const header: *align(16) usize = @ptrCast(bytes.ptr);
    header.* = total;
    return @ptrCast(bytes.ptr + allocation_header_size);
}

fn syntaxCalloc(count: usize, size: usize) callconv(.c) ?*anyopaque {
    const len = std.math.mul(usize, count, size) catch return null;
    const pointer = syntaxAlloc(len) orelse return null;
    const bytes: [*]u8 = @ptrCast(pointer);
    @memset(bytes[0..len], 0);
    return pointer;
}

fn syntaxRealloc(ptr: ?*anyopaque, new_size: usize) callconv(.c) ?*anyopaque {
    const pointer = ptr orelse return syntaxAlloc(new_size);
    if (new_size == 0) {
        syntaxFree(pointer);
        return null;
    }
    const user: [*]u8 = @ptrCast(pointer);
    const header: *align(16) usize = @ptrCast(@alignCast(user - allocation_header_size));
    const old_total = header.*;
    const old_bytes: []align(16) u8 = @as([*]align(16) u8, @ptrCast(header))[0..old_total];
    const new_total = std.math.add(usize, allocation_header_size, new_size) catch return null;
    const new_bytes = syntax_allocator.realloc(old_bytes, new_total) catch return null;
    const new_header: *align(16) usize = @ptrCast(new_bytes.ptr);
    new_header.* = new_total;
    return @ptrCast(new_bytes.ptr + allocation_header_size);
}

fn syntaxFree(ptr: ?*anyopaque) callconv(.c) void {
    const pointer = ptr orelse return;
    const user: [*]u8 = @ptrCast(pointer);
    const header: *align(16) usize = @ptrCast(@alignCast(user - allocation_header_size));
    const total = header.*;
    const bytes: []align(16) u8 = @as([*]align(16) u8, @ptrCast(header))[0..total];
    syntax_allocator.free(bytes);
}

pub fn start(gpa: std.mem.Allocator) void {
    if (comptime enabled) {
        std.debug.assert(!syntax_started);
        syntax_allocator = gpa;
        syntax_started = true;
        ts_set_allocator(syntaxAlloc, syntaxCalloc, syntaxRealloc, syntaxFree);
    }
}

pub fn stop() void {
    if (comptime enabled) {
        std.debug.assert(syntax_started);
        for (&specs) |*spec| {
            if (spec.compiled_query) |query| query.destroy();
            spec.compiled_query = null;
        }
        ts_set_allocator(null, null, null, null);
        syntax_started = false;
        syntax_allocator = undefined;
    }
}

const Selected = struct { lang: *const ts.Language, query: *ts.Query };

// NOTE: don't lang.destroy() — the tree_sitter_*() languages are static
// singletons reused on every open; destroying one use-after-frees the next.
fn ensure(spec: *Spec) !Selected {
    const lang = spec.language();
    if (spec.compiled_query) |query| return .{ .lang = lang, .query = query };
    var error_offset: u32 = 0;
    const query = try ts.Query.create(lang, spec.query_src, &error_offset);
    spec.compiled_query = query;
    return .{ .lang = lang, .query = query };
}

fn forExt(ext: []const u8) !?Selected {
    for (&specs) |*spec| {
        for (spec.exts) |choice| {
            if (std.ascii.eqlIgnoreCase(ext, choice)) return try ensure(spec);
        }
    }
    return null;
}

// Queries compile on first use (`ensure`), never at startup. Pre-compiling the
// compact tier in Pardes.init cost EVERY boot ~120ms of ts_query__perform_analysis
// (55% of a Debug startup) to save ~40ms on the first .zig/.c/.cpp open — a pane
// of prose or a shell paid for a language it never opened. Grammar availability
// is unchanged; only the timing moved.

fn synFor(name: []const u8) Syn {
    for ([_]struct { []const u8, Syn }{
        .{ "comment", .comment },
        .{ "string", .string },
        .{ "character", .string },
        .{ "number", .number },
        .{ "float", .number },
        .{ "boolean", .number },
        .{ "keyword", .keyword },
        .{ "include", .keyword },
        .{ "conditional", .keyword },
        .{ "repeat", .keyword },
    }) |m| {
        if (std.mem.indexOf(u8, name, m[0]) != null) return m[1];
    }
    return .none;
}

/// One Syn byte per content byte in [start, end). Caller frees.
pub fn highlightFileRange(gpa: std.mem.Allocator, path: []const u8, content: []const u8, start_byte_raw: usize, end_byte_raw: usize) ![]u8 {
    const tz = tracy.zone(@src(), "highlightFileRange");
    defer tz.end();
    if (!enabled) return &.{};
    const ext = std.fs.path.extension(path);
    const selected = (forExt(ext) catch return &.{}) orelse return &.{};

    const start_byte = @min(start_byte_raw, content.len);
    const end_byte = @max(start_byte, @min(end_byte_raw, content.len));
    const source = content[start_byte..end_byte];
    const styles = try gpa.alloc(u8, source.len);
    errdefer gpa.free(styles);
    @memset(styles, 0);

    const parser = ts.Parser.create();
    defer parser.destroy();
    parser.setLanguage(selected.lang) catch return styles;
    const tree = parser.parseString(source, null) orelse return styles;
    defer tree.destroy();

    const cursor = ts.QueryCursor.create();
    defer cursor.destroy();
    cursor.exec(selected.query, tree.rootNode());
    while (cursor.nextMatch()) |match| {
        for (match.captures) |cap| {
            const syn = synFor(selected.query.captureNameForId(cap.index) orelse "");
            if (syn == .none) continue;
            var b: usize = cap.node.startByte();
            const end = @min(@as(usize, cap.node.endByte()), styles.len);
            while (b < end) : (b += 1) styles[b] = @intFromEnum(syn);
        }
    }
    return styles;
}

test "tree-sitter allocator callbacks preserve and free exact allocations" {
    syntax_allocator = std.testing.allocator;
    defer syntax_allocator = undefined;

    var live: ?*anyopaque = syntaxCalloc(4, 1) orelse return error.OutOfMemory;
    defer if (live) |pointer| syntaxFree(pointer);
    const original: [*]u8 = @ptrCast(live.?);
    try std.testing.expectEqualSlices(u8, &.{ 0, 0, 0, 0 }, original[0..4]);
    @memcpy(original[0..4], "data");

    live = syntaxRealloc(live, 32) orelse return error.OutOfMemory;
    const grown: [*]u8 = @ptrCast(live.?);
    try std.testing.expectEqualSlices(u8, "data", grown[0..4]);

    try std.testing.expect(syntaxCalloc(std.math.maxInt(usize), 2) == null);
    try std.testing.expect(syntaxRealloc(live, 0) == null);
    live = null;
}