Palette from Image — Zig source

// Palette from Image — median-cut quantization over RGBA pixels.
//
// Language: Zig (0.13+), standard library only
// CosmoDev polyglot showcase port of the `palette-from-image` tool.
// Ported from src/lib/palette-extract.ts — display source, part of CosmoDev's
// polyglot tool pages.
//
// Deterministic: stable sorts only, widest-channel median split, buckets
// average into swatches. Transparent pixels are skipped.

const std = @import("std");
const Allocator = std.mem.Allocator;

pub const Swatch = struct { r: u8, g: u8, b: u8, population: usize };

const Pixel = struct { r: u8, g: u8, b: u8 };

/// Down-sample so large images quantize in bounded time (TS: MAX_SAMPLES).
const MAX_SAMPLES = 16_384;

pub fn extractPalette(gpa: Allocator, rgba: []const u8, max_colors: usize) ![]Swatch {
    const total = rgba.len / 4;
    if (total == 0) return &[_]Swatch{};

    var pixels = std.ArrayList(Pixel).init(gpa);
    defer pixels.deinit();
    const stride = @max(1, total / MAX_SAMPLES);
    var i: usize = 0;
    while (i < total) : (i += stride) {
        const o = i * 4;
        if (rgba[o + 3] == 0) continue; // fully transparent — skip
        try pixels.append(.{ .r = rgba[o], .g = rgba[o + 1], .b = rgba[o + 2] });
    }
    if (pixels.items.len == 0) return &[_]Swatch{};

    var buckets = std.ArrayList([]Pixel).init(gpa);
    defer buckets.deinit();
    try buckets.append(pixels.items);

    while (buckets.items.len < max_colors) {
        // Widest-range bucket with more than one distinct value splits.
        var best_idx: ?usize = null;
        var best_range: u8 = 1; // range 1 (exact duplicates) never splits
        var best_channel: u2 = 0;
        for (buckets.items, 0..) |bucket, k| {
            const rc = channelRange(bucket);
            if (rc.range > best_range) {
                best_range = rc.range;
                best_idx = k;
                best_channel = rc.channel;
            }
        }
        if (best_idx == null) break; // every bucket is uniform — done

        const bucket = buckets.orderedRemove(best_idx.?);
        const sorted = try gpa.dupe(Pixel, bucket);
        defer gpa.free(sorted);
        // Zig's sort is UNSTABLE; sort with an index tiebreak to mirror the
        // TS lib's stable sort (deterministic median split).
        const Ctx = struct {
            channel: u2,
            fn lessThan(self: @This(), a: Pixel, b: Pixel) bool {
                return false; // replaced below by indexed sort
            }
        };
        _ = Ctx;
        // Indexed stable sort: pair each pixel with its original position.
        const Idx = struct { px: Pixel, i: usize };
        var idx = try gpa.alloc(Idx, sorted.len);
        defer gpa.free(idx);
        for (sorted, 0..) |p, n| idx[n] = .{ .px = p, .i = n };
        const ch = best_channel;
        std.sort.pdq(Idx, idx, ch, struct {
            fn lessThan(channel: u2, a: Idx, b: Idx) bool {
                const av = switch (channel) { 0 => a.px.r, 1 => a.px.g, else => a.px.b };
                const bv = switch (channel) { 0 => b.px.r, 1 => b.px.g, else => b.px.b };
                return if (av == bv) a.i < b.i else av < bv;
            }
        }.lessThan);
        for (idx, 0..) |e, n| sorted[n] = e.px;

        const mid = sorted.len / 2;
        try buckets.append(sorted[0..mid]);
        try buckets.append(sorted[mid..]);
    }

    var out = try gpa.alloc(Swatch, buckets.items.len);
    errdefer gpa.free(out);
    for (buckets.items, 0..) |b, k| {
        var rsum: u64 = 0;
        var gsum: u64 = 0;
        var bsum: u64 = 0;
        for (b) |p| {
            rsum += p.r;
            gsum += p.g;
            bsum += p.b;
        }
        out[k] = .{
            .r = @intCast(@divTrunc(rsum + b.len / 2, b.len)),
            .g = @intCast(@divTrunc(gsum + b.len / 2, b.len)),
            .b = @intCast(@divTrunc(bsum + b.len / 2, b.len)),
            .population = b.len,
        };
    }
    std.sort.pdq(Swatch, out, {}, struct {
        fn lessThan(_: void, a: Swatch, b: Swatch) bool {
            return a.population > b.population;
        }
    }.lessThan);
    return out;
}

const RangeChannel = struct { range: u8, channel: u2 };

fn channelRange(bucket: []const Pixel) RangeChannel {
    var r_min: u8 = 255;
    var r_max: u8 = 0;
    var g_min: u8 = 255;
    var g_max: u8 = 0;
    var b_min: u8 = 255;
    var b_max: u8 = 0;
    for (bucket) |p| {
        r_min = @min(r_min, p.r); r_max = @max(r_max, p.r);
        g_min = @min(g_min, p.g); g_max = @max(g_max, p.g);
        b_min = @min(b_min, p.b); b_max = @max(b_max, p.b);
    }
    const r = r_max - r_min;
    const g = g_max - g_min;
    const b = b_max - b_min;
    if (r >= g and r >= b) return .{ .range = r, .channel = 0 };
    if (g >= b) return .{ .range = g, .channel = 1 };
    return .{ .range = b, .channel = 2 };
}

pub fn main() !void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const gpa = arena.allocator();

    // Smoke: 2 red + 1 blue pixels → red swatch wins on population.
    const rgba = [_]u8{ 200, 30, 30, 255, 202, 28, 32, 255, 30, 30, 200, 255 };
    const swatches = try extractPalette(gpa, &rgba, 2);
    defer gpa.free(swatches);
    for (swatches) |s|
        std.debug.print("#{X:0>2}{X:0>2}{X:0>2} x{d}\n", .{ s.r, s.g, s.b, s.population });
}