src/png.zig
Ref: Size: 9.9 KiB
const std = @import("std");
pub const Image = struct {
width: u32,
height: u32,
pixels: []u8, // RGBA8, row-major, width*height*4 bytes
pub fn deinit(self: *Image, alloc: std.mem.Allocator) void {
alloc.free(self.pixels);
self.* = undefined;
}
};
pub const EncodeError = error{ OutOfMemory, WriteFailed };
pub const DecodeError = error{
OutOfMemory,
InvalidPng,
UnsupportedPng, // only RGBA8 non-interlaced is supported
CorruptChunk,
};
const signature = [_]u8{ 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A };
fn adler32(data: []const u8) u32 {
var a: u32 = 1;
var b: u32 = 0;
for (data) |byte| {
a = (a + byte) % 65521;
b = (b + a) % 65521;
}
return (b << 16) | a;
}
fn writeChunk(writer: anytype, chunk_type: *const [4]u8, payload: []const u8) EncodeError!void {
writer.writeInt(u32, @intCast(payload.len), .big) catch return error.WriteFailed;
writer.writeAll(chunk_type) catch return error.WriteFailed;
writer.writeAll(payload) catch return error.WriteFailed;
var crc = std.hash.Crc32.init();
crc.update(chunk_type);
crc.update(payload);
writer.writeInt(u32, crc.final(), .big) catch return error.WriteFailed;
}
/// Build a zlib stream wrapping the `filtered` data using DEFLATE stored
/// blocks (type 0, no compression). This is always valid PNG and avoids
/// dependency on the std.compress.flate encoder, which is incomplete in
/// Zig 0.15.
fn buildZlibStored(alloc: std.mem.Allocator, filtered: []const u8) EncodeError![]u8 {
// zlib header: CMF=0x78 (deflate, window=32K), FLG=0x01 (no dict, level=0,
// fcheck makes CMF*256+FLG divisible by 31: 0x7801 % 31 == 0).
const zlib_header = [_]u8{ 0x78, 0x01 };
// DEFLATE stored block layout:
// 1 byte: BFINAL | (BTYPE << 1) — BTYPE=00 for stored
// 2 bytes: LEN (little-endian u16)
// 2 bytes: NLEN (one's complement of LEN, little-endian)
// LEN bytes: data
//
// Maximum single stored block payload is 65535 bytes.
const max_block: usize = 65535;
const actual_blocks: usize = if (filtered.len == 0) 1 else (filtered.len + max_block - 1) / max_block;
// Header per block: 5 bytes. Total deflate stream bytes:
const deflate_len = actual_blocks * 5 + filtered.len;
// Full buffer: zlib_header(2) + deflate + adler32(4)
const total = 2 + deflate_len + 4;
const buf = alloc.alloc(u8, total) catch return error.OutOfMemory;
errdefer alloc.free(buf);
var pos: usize = 0;
buf[pos] = zlib_header[0];
pos += 1;
buf[pos] = zlib_header[1];
pos += 1;
var src_pos: usize = 0;
var block_idx: usize = 0;
while (block_idx < actual_blocks) : (block_idx += 1) {
const remaining = filtered.len - src_pos;
const block_len: u16 = @intCast(@min(remaining, max_block));
const is_final = block_idx == actual_blocks - 1;
const bfinal: u8 = if (is_final) 0x01 else 0x00;
buf[pos] = bfinal; // BFINAL=is_final, BTYPE=00
pos += 1;
std.mem.writeInt(u16, buf[pos..][0..2], block_len, .little);
pos += 2;
const nlen: u16 = ~block_len;
std.mem.writeInt(u16, buf[pos..][0..2], nlen, .little);
pos += 2;
@memcpy(buf[pos..][0..block_len], filtered[src_pos..][0..block_len]);
pos += block_len;
src_pos += block_len;
}
// Adler-32 of the uncompressed (filtered) data, big-endian
std.mem.writeInt(u32, buf[pos..][0..4], adler32(filtered), .big);
pos += 4;
std.debug.assert(pos == total);
return buf;
}
pub fn encode(alloc: std.mem.Allocator, img: Image, writer: anytype) EncodeError!void {
std.debug.assert(img.pixels.len == @as(usize, img.width) * img.height * 4);
writer.writeAll(&signature) catch return error.WriteFailed;
var ihdr: [13]u8 = undefined;
std.mem.writeInt(u32, ihdr[0..4], img.width, .big);
std.mem.writeInt(u32, ihdr[4..8], img.height, .big);
ihdr[8] = 8; // bit depth
ihdr[9] = 6; // colour type = RGBA
ihdr[10] = 0; // compression method
ihdr[11] = 0; // filter method
ihdr[12] = 0; // interlace method = none
try writeChunk(writer, "IHDR", &ihdr);
const row_bytes = @as(usize, img.width) * 4;
const filtered_len = (row_bytes + 1) * img.height;
const filtered = alloc.alloc(u8, filtered_len) catch return error.OutOfMemory;
defer alloc.free(filtered);
// Filter type 0 (None) per row
var y: u32 = 0;
while (y < img.height) : (y += 1) {
const src_off = @as(usize, y) * row_bytes;
const dst_off = @as(usize, y) * (row_bytes + 1);
filtered[dst_off] = 0; // filter byte
@memcpy(filtered[dst_off + 1 ..][0..row_bytes], img.pixels[src_off..][0..row_bytes]);
}
const compressed = try buildZlibStored(alloc, filtered);
defer alloc.free(compressed);
try writeChunk(writer, "IDAT", compressed);
try writeChunk(writer, "IEND", &.{});
}
pub fn decode(alloc: std.mem.Allocator, bytes: []const u8) DecodeError!Image {
if (bytes.len < signature.len + 8) return error.InvalidPng;
if (!std.mem.eql(u8, bytes[0..signature.len], &signature)) return error.InvalidPng;
var cursor: usize = signature.len;
var width: u32 = 0;
var height: u32 = 0;
var idat_accum: std.ArrayList(u8) = .empty;
defer idat_accum.deinit(alloc);
var seen_ihdr = false;
var seen_iend = false;
while (cursor + 8 <= bytes.len and !seen_iend) {
const len = std.mem.readInt(u32, bytes[cursor..][0..4], .big);
cursor += 4;
const ctype = bytes[cursor..][0..4];
cursor += 4;
if (cursor + len + 4 > bytes.len) return error.CorruptChunk;
const payload = bytes[cursor..][0..len];
cursor += len;
cursor += 4; // skip CRC
if (std.mem.eql(u8, ctype, "IHDR")) {
if (payload.len != 13) return error.InvalidPng;
width = std.mem.readInt(u32, payload[0..4], .big);
height = std.mem.readInt(u32, payload[4..8], .big);
// bit depth=8, colour type=6 (RGBA), interlace=0
if (payload[8] != 8 or payload[9] != 6 or payload[12] != 0)
return error.UnsupportedPng;
seen_ihdr = true;
} else if (std.mem.eql(u8, ctype, "IDAT")) {
if (!seen_ihdr) return error.InvalidPng;
idat_accum.appendSlice(alloc, payload) catch return error.OutOfMemory;
} else if (std.mem.eql(u8, ctype, "IEND")) {
seen_iend = true;
}
}
if (!seen_ihdr or !seen_iend) return error.InvalidPng;
// zlib stream: 2-byte header + deflate body + 4-byte adler32
if (idat_accum.items.len < 6) return error.InvalidPng;
// Strip the 2-byte zlib header and 4-byte adler32 footer to get raw deflate
const deflate_data = idat_accum.items[2 .. idat_accum.items.len - 4];
const row_bytes = @as(usize, width) * 4;
const filtered_len = (row_bytes + 1) * @as(usize, height);
const filtered = alloc.alloc(u8, filtered_len) catch return error.OutOfMemory;
defer alloc.free(filtered);
// Decompress using std.compress.flate.Decompress with the new Zig 0.15 API.
// The indirect vtable (used when a window buffer is provided) fills its
// internal buffer on each vtable call and returns 0; the caller must loop,
// draining the buffer on alternate calls.
{
var in_reader: std.Io.Reader = .fixed(deflate_data);
var decomp_buf: [std.compress.flate.max_window_len]u8 = undefined;
var decomp: std.compress.flate.Decompress = .init(&in_reader, .raw, &decomp_buf);
var dst_writer: std.Io.Writer = .fixed(filtered);
var written: usize = 0;
while (written < filtered_len) {
const n = decomp.reader.stream(&dst_writer, .unlimited) catch |err| switch (err) {
error.EndOfStream => break,
else => return error.CorruptChunk,
};
written += n;
if (n == 0 and decomp.reader.seek == decomp.reader.end) break;
}
if (written != filtered_len) return error.CorruptChunk;
}
const pixels = alloc.alloc(u8, @as(usize, width) * height * 4) catch return error.OutOfMemory;
errdefer alloc.free(pixels);
var row: u32 = 0;
while (row < height) : (row += 1) {
const dst_off = @as(usize, row) * row_bytes;
const src_off = @as(usize, row) * (row_bytes + 1);
if (filtered[src_off] != 0) return error.UnsupportedPng; // only filter type 0
@memcpy(pixels[dst_off..][0..row_bytes], filtered[src_off + 1 ..][0..row_bytes]);
}
return .{ .width = width, .height = height, .pixels = pixels };
}
test "encode then decode roundtrip recovers pixels" {
const alloc = std.testing.allocator;
var src_pixels = [_]u8{
0xff, 0x00, 0x00, 0xff, 0x00, 0xff, 0x00, 0xff,
0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
};
const src = Image{ .width = 2, .height = 2, .pixels = &src_pixels };
var buf: std.ArrayList(u8) = .empty;
defer buf.deinit(alloc);
try encode(alloc, src, buf.writer(alloc));
var decoded = try decode(alloc, buf.items[0..]);
defer decoded.deinit(alloc);
try std.testing.expectEqual(@as(u32, 2), decoded.width);
try std.testing.expectEqual(@as(u32, 2), decoded.height);
try std.testing.expectEqualSlices(u8, &src_pixels, decoded.pixels);
}
test "decode rejects RGB (non-alpha) PNGs with UnsupportedPng" {
const alloc = std.testing.allocator;
var bytes: std.ArrayList(u8) = .empty;
defer bytes.deinit(alloc);
try bytes.appendSlice(alloc, &signature);
var ihdr: [13]u8 = undefined;
std.mem.writeInt(u32, ihdr[0..4], 1, .big);
std.mem.writeInt(u32, ihdr[4..8], 1, .big);
ihdr[8] = 8;
ihdr[9] = 2; // colour type 2 = RGB (not RGBA)
ihdr[10] = 0;
ihdr[11] = 0;
ihdr[12] = 0;
try writeChunk(bytes.writer(alloc), "IHDR", &ihdr);
try writeChunk(bytes.writer(alloc), "IEND", &.{});
try std.testing.expectError(error.UnsupportedPng, decode(alloc, bytes.items[0..]));
}