Zig netbpm implementation

netbpm.zig

const std = @import("std");
const fs = std.fs;

const Allocator = std.mem.Allocator;

pub const MagicNumber = enum(u3) {
    P1 = 1,
    P2 = 2,
    P3 = 3,
    P4 = 4,
    P5 = 5,
    P6 = 6,
};

pub fn channelCount(magic_number: MagicNumber) u3 {
    return switch (magic_number) {
        .P1, .P4 => 1, // binary
        .P2, .P5 => 1, // grayscale
        .P3, .P6 => 3, // rgb
    };
}

pub const Image = struct {
    magic_number: MagicNumber = .P1,
    width: usize = undefined,
    height: usize = undefined,
    buf: []u8 = undefined,
    head: usize = undefined,

    pub fn init(
        allocator: Allocator,
        magic_number: MagicNumber,
        w: usize,
        h: usize,
    ) !Image {
        var image = Image{};
        image.width = w;
        image.height = h;
        image.magic_number = magic_number;
        var buffer_size: usize = 60;
        buffer_size += (image.width * (channelCount(magic_number)) * 9) * image.height;
        // std.debug.print("\nBuffer size: {d}", .{buffer_size});
        image.buf = try allocator.alloc(
            u8,
            buffer_size,
        );
        try image.writeHeader();
        return image;
    }

    fn writeHeader(self: *Image) !void {
        const magic = try std.fmt.bufPrint(self.buf[self.head..], "P{d}\n", .{@intFromEnum(
            self.magic_number,
        )});
        self.head += magic.len;
        const specs = try std.fmt.bufPrint(self.buf[self.head..], "{d} {d}\n255\n", .{
            self.width,
            self.height,
        });
        self.head += specs.len;
    }

    fn stride(self: *Image, x: usize, y: usize) usize {
        return y * self.width + x;
    }

    pub fn set(
        self: *Image,
        comptime T: MagicNumber,
        x: usize,
        y: usize,
        rgb: @Vector(channelCount(T), u8),
    ) !void {
        _ = y;
        _ = x;
        switch (T) {
            .P1,
            .P2,
            => {
                const printed = try std.fmt.bufPrint(self.buf[self.head..], "{d}", .{
                    rgb[0],
                });
                self.head += printed.len;
            },
            .P3 => {
                const printed = try std.fmt.bufPrint(self.buf[self.head..], "{d} {d} {d} ", .{
                    rgb[0],
                    rgb[1],
                    rgb[2],
                });
                self.head += printed.len;
            },
            .P4, .P5, .P6 => {
                @panic("Not yet implemented!");
            },
        }
    }

    pub fn write(self: *Image, file_path: []const u8) !void {
        if (std.fs.path.dirname(file_path)) |dirname| {
            fs.cwd().makeDir(dirname) catch {
                std.debug.print("\n{s} already exists", .{dirname});
            };
        }
        var file = try fs.cwd().createFile(file_path, .{});
        defer file.close();
        try file.writeAll(self.buf[0..self.head]);
    }
};