Cache Breakpoint Planner — Zig source

// Cache Breakpoint Planner — find the blocks a set of prompts share and
// place cache breakpoints where they pay.
//
// Language: Zig (0.13+, stdlib only)
// Port of src/lib/cacheBreakpointPlanner.ts (the canonical TypeScript
// implementation). Plan strings/slices reference the INPUT sessions (no
// copies) — keep the inputs alive while using the plan. Reason strings are
// formatted into fixed buffers.
//
// Tool page: https://dev.cosmolabs.org/tools/cache-breakpoint-planner

const std = @import("std");

/// Cached reads bill at ~0.1x — the saving on the cached share is ~90%.
pub const cache_read_discount: f64 = 0.1;

/// One prompt session: an id plus its ordered blocks.
pub const Session = struct {
    id: []const u8,
    blocks: []const []const u8,
};

/// Place the cache breakpoint AFTER this block index (0-based).
pub const Breakpoint = struct {
    after_block: i64, // -1 = terminal (the shared tail sits at the end)
    label: []const u8,
    reason: [256]u8,
    reason_len: usize,
    cached_tokens: i64,

    pub fn reasonText(self: *const Breakpoint) []const u8 {
        return self.reason[0..self.reason_len];
    }
};

pub const Row = struct {
    id: []const u8,
    total_tokens: i64,
    unique_tokens: i64,
    cached_ratio: f64,
};

/// Raise these caps for larger display examples; the planner saturates.
pub const MAX_BLOCKS = 64;
pub const MAX_SESSIONS = 32;
pub const MAX_BREAKPOINTS = 2;
pub const MAX_WARNINGS = 4;

pub const Plan = struct {
    prefix_blocks: []const []const u8 = &.{},
    prefix_tokens: i64 = 0,
    suffix_blocks: []const []const u8 = &.{},
    suffix_tokens: i64 = 0,
    breakpoints: [MAX_BREAKPOINTS]Breakpoint = undefined,
    breakpoint_count: usize = 0,
    rows: [MAX_SESSIONS]Row = undefined,
    row_count: usize = 0,
    /// Estimated cost saving across the sessions vs no caching (0-1).
    estimated_savings: f64 = 0,
    warnings: [MAX_WARNINGS][192]u8 = undefined,
    warning_lens: [MAX_WARNINGS]usize = undefined,
    warning_count: usize = 0,

    pub fn warningText(self: *const Plan, i: usize) []const u8 {
        return self.warnings[i][0..self.warning_lens[i]];
    }
};

/// The `type: 'prose'` path of the tokenEstimator, inlined: every non-empty
/// line costs max(1, round(length / 4)) tokens; empty text is 0.
fn tok(text: []const u8) i64 {
    if (text.len == 0) return 0;
    var tokens: i64 = 0;
    var it = std.mem.splitScalar(u8, text, '\n');
    while (it.next()) |line| {
        if (line.len > 0) {
            const per: i64 = @intFromFloat(@as(f64, @floatFromInt(line.len)) / 4.0 + 0.5);
            tokens += @max(1, per);
        }
    }
    return tokens;
}

fn tokBlocks(blocks: []const []const u8) i64 {
    var total: i64 = 0;
    for (blocks) |b| total += tok(b);
    return total;
}

fn warn(plan: *Plan, comptime fmt: []const u8, args: anytype) void {
    if (plan.warning_count >= MAX_WARNINGS) return;
    const i = plan.warning_count;
    const out = std.fmt.bufPrint(&plan.warnings[i], fmt, args) catch
        std.fmt.bufPrint(&plan.warnings[i], "{s}", .{"(warning truncated)"}) catch return;
    plan.warning_lens[i] = out.len;
    plan.warning_count += 1;
}

/// Plan cache breakpoints for a set of prompt sessions: find the common
/// leading/trailing blocks across every session and place breakpoints where
/// the cache pays.
pub fn planBreakpoints(sessions: []const Session) Plan {
    var plan: Plan = .{};

    if (sessions.len == 0) {
        warn(&plan, "No sessions given — paste at least two prompts to compare.", .{});
        return plan;
    }
    if (sessions.len == 1) {
        warn(&plan, "Only one session — a prefix needs at least two prompts to detect.", .{});
    }

    // Common leading blocks by position.
    var shortest: usize = sessions[0].blocks.len;
    for (sessions[1..]) |s| {
        shortest = @min(shortest, s.blocks.len);
    }
    var prefix_end: usize = 0;
    while (prefix_end < shortest) {
        var shared = true;
        for (sessions) |s| {
            if (!std.mem.eql(u8, s.blocks[prefix_end], sessions[0].blocks[prefix_end])) {
                shared = false;
                break;
            }
        }
        if (!shared) break;
        prefix_end += 1;
    }

    // Common trailing blocks, matched from each session's own tail, never
    // overlapping the prefix.
    var suffix_len: usize = 0;
    while (suffix_len < shortest - prefix_end) {
        var shared = true;
        for (sessions) |s| {
            const a = s.blocks[s.blocks.len - 1 - suffix_len];
            const b = sessions[0].blocks[sessions[0].blocks.len - 1 - suffix_len];
            if (!std.mem.eql(u8, a, b)) {
                shared = false;
                break;
            }
        }
        if (!shared) break;
        suffix_len += 1;
    }

    plan.prefix_blocks = sessions[0].blocks[0..@min(prefix_end, MAX_BLOCKS)];
    if (suffix_len > 0) {
        const n = sessions[0].blocks.len;
        plan.suffix_blocks = sessions[0].blocks[n - suffix_len .. n];
    }
    plan.prefix_tokens = tokBlocks(plan.prefix_blocks);
    plan.suffix_tokens = tokBlocks(plan.suffix_blocks);

    if (plan.prefix_blocks.len > 0) {
        const bp = &plan.breakpoints[plan.breakpoint_count];
        bp.* = .{
            .after_block = @as(i64, @intCast(prefix_end)) - 1,
            .label = "after the shared prefix",
            .reason = undefined,
            .reason_len = 0,
            .cached_tokens = plan.prefix_tokens,
        };
        const out = std.fmt.bufPrint(&bp.reason, "{d} block(s) identical across every session — cache once, hit on every request.", .{plan.prefix_blocks.len}) catch "";
        bp.reason_len = out.len;
        plan.breakpoint_count += 1;
    }
    if (suffix_len > 0) {
        const bp = &plan.breakpoints[plan.breakpoint_count];
        bp.* = .{
            .after_block = -1, // terminal: the shared tail sits at the end
            .label = "shared tail",
            .reason = undefined,
            .reason_len = 0,
            .cached_tokens = plan.suffix_tokens,
        };
        const out = std.fmt.bufPrint(&bp.reason, "{d} trailing block(s) also identical — extend the cache segment or accept the re-read.", .{suffix_len}) catch "";
        bp.reason_len = out.len;
        plan.breakpoint_count += 1;
    }
    if (plan.breakpoint_count == 0) {
        warn(&plan, "No shared leading or trailing blocks — nothing to cache across these sessions.", .{});
    }

    for (sessions, 0..) |s, i| {
        if (i >= MAX_SESSIONS) break;
        const row = &plan.rows[plan.row_count];
        const total = tokBlocks(s.blocks);
        const unique = @max(total - plan.prefix_tokens - plan.suffix_tokens, 0);
        row.* = .{
            .id = s.id,
            .total_tokens = total,
            .unique_tokens = unique,
            .cached_ratio = if (total > 0)
                @min(@as(f64, @floatFromInt(plan.prefix_tokens + plan.suffix_tokens)) / @as(f64, @floatFromInt(total)), 1.0)
            else
                0.0,
        };
        plan.row_count += 1;
    }

    var sum: i64 = 0;
    for (plan.rows[0..plan.row_count]) |r| sum += r.total_tokens;
    const avg_total: f64 = if (plan.row_count > 0)
        @as(f64, @floatFromInt(sum)) / @as(f64, @floatFromInt(plan.row_count))
    else
        0.0;
    const cached_share: f64 = if (avg_total > 0)
        @min(@as(f64, @floatFromInt(plan.prefix_tokens + plan.suffix_tokens)) / avg_total, 1.0)
    else
        0.0;
    plan.estimated_savings = cached_share * (1.0 - cache_read_discount);

    return plan;
}