a73x

src/capture.zig

Ref:   Size: 11.8 KiB

//! `--capture <script> <output.png>` mode.
//!
//! Renders a VT script to a single PNG frame for golden-image testing.
//!
//! 1. Stand up a Wayland window + Vulkan context at a forced 80x24 grid,
//!    buffer scale = 1 (so renders are deterministic across multi-monitor
//!    setups).
//! 2. Wait up to 3s for the window to become visible.
//! 3. Pipe the script through a PTY via `/bin/cat`, then drain remaining
//!    output after cat exits.
//! 4. Snapshot the terminal, build a flat Instance list for every cell,
//!    render a single frame to an offscreen VkImage, read the BGRA bytes
//!    back, convert to RGBA and write a PNG.
//!
//! The window itself is never committed/presented — the offscreen target
//! is its own framebuffer. We still need the Wayland surface so Vulkan
//! can allocate a swapchain (required by the current Context.init path)
//! and so the compositor hands us a real configure event.

const std = @import("std");
const vt = @import("vt");
const pty = @import("pty");
const wayland_client = @import("wayland-client");
const renderer = @import("renderer");
const font = @import("font");
const config = @import("config");
const png = @import("png");
const vk = @import("vulkan");

pub const CaptureError = error{
    MissingArgs,
    ScriptNotFound,
    OutputPathUnwritable,
    WindowNotVisible,
    WindowSizeMismatch,
    PngEncodeFailed,
};

const cell_instance = @import("cell_instance");
const appendCellInstances = cell_instance.appendCellInstances;
const glyphTopOffset = cell_instance.glyphTopOffset;

const CAPTURE_COLS: u16 = 80;
const CAPTURE_ROWS: u16 = 24;
const VISIBILITY_TIMEOUT_NS: i128 = 3 * std.time.ns_per_s;

/// Entry point. `argv[0]` is `--capture`; argv[1] = script path, argv[2] = out path.
pub fn run(alloc: std.mem.Allocator, argv: []const [:0]const u8) !void {
    if (argv.len < 3) {
        std.debug.print("usage: waystty --capture <script.vt> <output.png>\n", .{});
        return CaptureError.MissingArgs;
    }
    const script_path = argv[1];
    const out_path = argv[2];

    // Probe script path up-front so we fail fast with a clean error rather
    // than having cat silently print a "No such file" diagnostic onto the
    // captured image.
    std.fs.cwd().access(script_path, .{}) catch |err| {
        std.debug.print("capture: cannot read script {s}: {t}\n", .{ script_path, err });
        return CaptureError.ScriptNotFound;
    };

    // === font + cell metrics (scale=1, same lookup as runTerminal) ===
    var font_lookup = try font.lookupConfiguredFont(alloc);
    defer font_lookup.deinit(alloc);

    const font_size: u32 = config.font_size_px;
    var face = try font.Face.init(alloc, font_lookup.path, font_lookup.index, font_size);
    defer face.deinit();

    const cell_w: u32 = face.cellWidth();
    const cell_h: u32 = face.cellHeight();
    const baseline: u32 = face.baseline();

    const px_w: u32 = @as(u32, CAPTURE_COLS) * cell_w;
    const px_h: u32 = @as(u32, CAPTURE_ROWS) * cell_h;

    // === wayland ===
    const conn = try wayland_client.Connection.init(alloc);
    defer conn.deinit();

    const window = try conn.createWindow(alloc, "waystty-capture");
    defer window.deinit();

    window.width = px_w;
    window.height = px_h;
    _ = conn.display.roundtrip();

    // === vulkan context (swapchain matches requested px size) ===
    var ctx = try renderer.Context.init(
        alloc,
        @ptrCast(conn.display),
        @ptrCast(window.surface),
        px_w,
        px_h,
    );
    defer ctx.deinit();

    // === offscreen render target (separate framebuffer; renders don't present) ===
    var offscreen = try renderer.createOffscreen(
        ctx.vki,
        ctx.vkd,
        ctx.physical_device,
        ctx.device,
        ctx.render_pass,
        ctx.swapchain_format,
        px_w,
        px_h,
    );
    defer renderer.destroyOffscreen(ctx.vkd, ctx.device, offscreen);

    // === glyph atlas + printable ASCII warm-up (matches runTerminal) ===
    var atlas = try font.Atlas.init(alloc, 1024, 1024);
    defer atlas.deinit();

    for (32..127) |cp| {
        _ = atlas.getOrInsert(&face, @intCast(cp)) catch |err| switch (err) {
            error.AtlasFull => break,
            else => return err,
        };
    }
    try ctx.uploadAtlas(atlas.pixels);
    atlas.last_uploaded_y = atlas.cursor_y;
    atlas.needs_full_upload = false;
    atlas.dirty = false;

    // === terminal ===
    var term = try vt.Terminal.init(alloc, .{
        .cols = CAPTURE_COLS,
        .rows = CAPTURE_ROWS,
        .max_scrollback = 1000,
    });
    defer term.deinit();
    term.setReportedSize(.{
        .rows = CAPTURE_ROWS,
        .columns = CAPTURE_COLS,
        .cell_width = cell_w,
        .cell_height = cell_h,
    });

    // === visibility wait + size check ===
    try waitUntilVisible(conn, window);

    if (window.width != px_w or window.height != px_h) {
        std.debug.print(
            "capture: window size mismatch (got {d}x{d}, expected {d}x{d})\n",
            .{ window.width, window.height, px_w, px_h },
        );
        return CaptureError.WindowSizeMismatch;
    }

    // === play script through /bin/cat ===
    try playScript(term, script_path);

    // === snapshot + build instances ===
    try term.snapshot();

    var instances: std.ArrayListUnmanaged(renderer.Instance) = .empty;
    defer instances.deinit(alloc);

    try buildInstancesForSnapshot(
        alloc,
        &instances,
        term,
        &face,
        &atlas,
        cell_w,
        cell_h,
        baseline,
    );

    // If the script needed glyphs that weren't in the ASCII warm-up set,
    // the atlas pixels are newer than the GPU copy. Re-upload the full
    // atlas so the render samples valid texels.
    if (atlas.dirty) {
        try ctx.uploadAtlas(atlas.pixels);
        atlas.dirty = false;
        atlas.last_uploaded_y = atlas.cursor_y;
    }

    // === render one frame to offscreen ===
    const push = renderer.PushConstants{
        .viewport_size = .{ @floatFromInt(px_w), @floatFromInt(px_h) },
        .cell_size = .{ @floatFromInt(cell_w), @floatFromInt(cell_h) },
        .coverage_params = renderer.coverageVariantParams(.baseline),
    };

    try ctx.renderToOffscreen(&offscreen, instances.items, push);

    // === readback BGRA->RGBA ===
    const rgba = try alloc.alloc(u8, @as(usize, px_w) * px_h * 4);
    defer alloc.free(rgba);
    try ctx.readbackOffscreen(&offscreen, rgba);

    // === encode PNG ===
    try writePng(alloc, out_path, px_w, px_h, rgba);

    std.debug.print("capture: wrote {s} ({d}x{d})\n", .{ out_path, px_w, px_h });
}

/// Wait up to VISIBILITY_TIMEOUT_NS for the Wayland compositor to `configure`
/// the surface. For `--capture` we don't need the surface to actually be
/// mapped onto an output (which would require committing a presentable
/// buffer via the swapchain — we deliberately skip that since rendering is
/// offscreen). A configured surface is enough to know our fixed 80x24
/// geometry was accepted.
fn waitUntilVisible(conn: *wayland_client.Connection, window: *wayland_client.Window) !void {
    const deadline = @as(i128, std.time.nanoTimestamp()) + VISIBILITY_TIMEOUT_NS;
    while (std.time.nanoTimestamp() < deadline) {
        _ = conn.display.roundtrip();
        if (window.state.configured) return;
        std.Thread.sleep(10 * std.time.ns_per_ms);
    }
    std.debug.print(
        "capture: window never configured within 3s\n",
        .{},
    );
    return CaptureError.WindowNotVisible;
}

/// Spawn `/bin/cat <script>` on a PTY; feed all its output into `term`.
/// Returns once the child has exited AND two consecutive 20 ms polls
/// produce no new bytes (drain).
///
/// Note: spawns cat with script as argv rather than piping stdin+^D —
/// avoids EOF-signalling races with VT escape sequences.
fn playScript(
    term: *vt.Terminal,
    script_path: [:0]const u8,
) !void {
    var p = try pty.Pty.spawn(.{
        .cols = CAPTURE_COLS,
        .rows = CAPTURE_ROWS,
        .shell = "/bin/cat",
        .shell_args = &.{script_path},
    });
    defer p.deinit();

    var buf: [4096]u8 = undefined;
    var consecutive_empty: u32 = 0;

    // Loop until child exited AND we saw two empty polls in a row (to make
    // sure any straggler bytes in the master buffer have been drained).
    while (true) {
        var pfd = [_]std.posix.pollfd{
            .{ .fd = p.master_fd, .events = std.posix.POLL.IN, .revents = 0 },
        };
        _ = std.posix.poll(&pfd, 20) catch 0;

        var saw_bytes = false;
        while (true) {
            const n = p.read(&buf) catch |err| switch (err) {
                error.WouldBlock => break,
                // EIO on Linux after slave fd closes is the normal signal
                // that cat exited. Break out — the child reaper below will
                // notice.
                error.InputOutput => break,
                else => return err,
            };
            if (n == 0) break;
            term.write(buf[0..n]);
            saw_bytes = true;
        }

        const alive = p.isChildAlive();
        if (saw_bytes) {
            consecutive_empty = 0;
        } else if (!alive) {
            consecutive_empty += 1;
            if (consecutive_empty >= 2) break;
        }
    }

    // VT parser settle — give any delayed effects (timers, etc) a beat.
    std.Thread.sleep(50 * std.time.ns_per_ms);
}

/// Build a flat Instance list covering every cell in the current snapshot.
/// Does not do dirty tracking — this is a one-shot full rebuild. Mirrors
/// the per-cell logic in `main.zig:rebuildRowInstances`, minus the
/// selection/cursor overlay.
fn buildInstancesForSnapshot(
    alloc: std.mem.Allocator,
    instances: *std.ArrayListUnmanaged(renderer.Instance),
    term: *vt.Terminal,
    face: *font.Face,
    atlas: *font.Atlas,
    cell_w: u32,
    cell_h: u32,
    baseline: u32,
) !void {
    const default_bg = term.backgroundColor();
    const bg_uv = atlas.cursorUV();

    const term_rows = term.render_state.row_data.items(.cells);
    var row_idx: u32 = 0;
    while (row_idx < term_rows.len) : (row_idx += 1) {
        const row_cells = term_rows[row_idx];
        const raw_cells = row_cells.items(.raw);
        var col_idx: u32 = 0;
        while (col_idx < raw_cells.len) : (col_idx += 1) {
            const cp = raw_cells[col_idx].codepoint();
            const colors = term.cellColors(row_cells.get(col_idx));
            const glyph_uv = if (cp == 0 or cp == ' ')
                null
            else
                atlas.getOrInsert(face, @intCast(cp)) catch null;

            try appendCellInstances(
                alloc,
                instances,
                row_idx,
                col_idx,
                cell_w,
                cell_h,
                baseline,
                glyph_uv,
                bg_uv,
                colors,
                default_bg,
            );
        }
    }
}

/// Encode `rgba` as a PNG to a brand-new file at `path`. Buffers the full
/// encoded byte stream in memory (fine for 80x24@16px: under 200 KB) and
/// writes it in one shot.
fn writePng(
    alloc: std.mem.Allocator,
    path: [:0]const u8,
    width: u32,
    height: u32,
    rgba: []u8,
) !void {
    var buf: std.ArrayList(u8) = .empty;
    defer buf.deinit(alloc);

    const img: png.Image = .{
        .width = width,
        .height = height,
        .pixels = rgba,
    };
    png.encode(alloc, img, buf.writer(alloc)) catch |err| {
        std.debug.print("capture: PNG encode failed: {t}\n", .{err});
        return CaptureError.PngEncodeFailed;
    };

    const file = std.fs.cwd().createFile(path, .{ .truncate = true }) catch |err| {
        std.debug.print("capture: cannot open output {s}: {t}\n", .{ path, err });
        return CaptureError.OutputPathUnwritable;
    };
    defer file.close();

    file.writeAll(buf.items) catch |err| {
        std.debug.print("capture: write failed for {s}: {t}\n", .{ path, err });
        return CaptureError.OutputPathUnwritable;
    };
}