zig/lib/std/io/FixedBufferStream.zig

//! This turns a const byte buffer into an `io.Reader`, or `io.SeekableStream`.

const std = @import("../std.zig");
const io = std.io;
const testing = std.testing;
const mem = std.mem;
const assert = std.debug.assert;
const FixedBufferStream = @This();

buffer: []const u8,
pos: usize = 0,

pub const ReadError = error{};
pub const SeekError = error{};
pub const GetSeekPosError = error{};

pub const Reader = io.Reader(*Self, ReadError, read);

pub const SeekableStream = io.SeekableStream(
    *Self,
    SeekError,
    GetSeekPosError,
    seekTo,
    seekBy,
    getPos,
    getEndPos,
);

const Self = @This();

pub fn reader(self: *Self) Reader {
    return .{ .context = self };
}

pub fn seekableStream(self: *Self) SeekableStream {
    return .{ .context = self };
}

pub fn read(self: *Self, dest: []u8) ReadError!usize {
    const size = @min(dest.len, self.buffer.len - self.pos);
    const end = self.pos + size;

    @memcpy(dest[0..size], self.buffer[self.pos..end]);
    self.pos = end;

    return size;
}

pub fn seekTo(self: *Self, pos: u64) SeekError!void {
    self.pos = @min(std.math.lossyCast(usize, pos), self.buffer.len);
}

pub fn seekBy(self: *Self, amt: i64) SeekError!void {
    if (amt < 0) {
        const abs_amt = @abs(amt);
        const abs_amt_usize = std.math.cast(usize, abs_amt) orelse std.math.maxInt(usize);
        if (abs_amt_usize > self.pos) {
            self.pos = 0;
        } else {
            self.pos -= abs_amt_usize;
        }
    } else {
        const amt_usize = std.math.cast(usize, amt) orelse std.math.maxInt(usize);
        const new_pos = std.math.add(usize, self.pos, amt_usize) catch std.math.maxInt(usize);
        self.pos = @min(self.buffer.len, new_pos);
    }
}

pub fn getEndPos(self: *Self) GetSeekPosError!u64 {
    return self.buffer.len;
}

pub fn getPos(self: *Self) GetSeekPosError!u64 {
    return self.pos;
}

pub fn getWritten(self: Self) []const u8 {
    return self.buffer[0..self.pos];
}

pub fn reset(self: *Self) void {
    self.pos = 0;
}

test "output" {
    var buf: [255]u8 = undefined;
    var fbs: FixedBufferStream = .{ .buffer = &buf };
    const stream = fbs.writer();

    try stream.print("{s}{s}!", .{ "Hello", "World" });
    try testing.expectEqualSlices(u8, "HelloWorld!", fbs.getWritten());
}

test "output at comptime" {
    comptime {
        var buf: [255]u8 = undefined;
        var fbs: FixedBufferStream = .{ .buffer = &buf };
        const stream = fbs.writer();

        try stream.print("{s}{s}!", .{ "Hello", "World" });
        try testing.expectEqualSlices(u8, "HelloWorld!", fbs.getWritten());
    }
}

test "output 2" {
    var buffer: [10]u8 = undefined;
    var fbs: FixedBufferStream = .{ .buffer = &buffer };

    try fbs.writer().writeAll("Hello");
    try testing.expect(mem.eql(u8, fbs.getWritten(), "Hello"));

    try fbs.writer().writeAll("world");
    try testing.expect(mem.eql(u8, fbs.getWritten(), "Helloworld"));

    try testing.expectError(error.NoSpaceLeft, fbs.writer().writeAll("!"));
    try testing.expect(mem.eql(u8, fbs.getWritten(), "Helloworld"));

    fbs.reset();
    try testing.expect(fbs.getWritten().len == 0);

    try testing.expectError(error.NoSpaceLeft, fbs.writer().writeAll("Hello world!"));
    try testing.expect(mem.eql(u8, fbs.getWritten(), "Hello worl"));

    try fbs.seekTo((try fbs.getEndPos()) + 1);
    try testing.expectError(error.NoSpaceLeft, fbs.writer().writeAll("H"));
}

test "input" {
    const bytes = [_]u8{ 1, 2, 3, 4, 5, 6, 7 };
    var fbs: FixedBufferStream = .{ .buffer = &bytes };

    var dest: [4]u8 = undefined;

    var amt_read = try fbs.reader().read(&dest);
    try testing.expect(amt_read == 4);
    try testing.expect(mem.eql(u8, dest[0..4], bytes[0..4]));

    amt_read = try fbs.reader().read(&dest);
    try testing.expect(amt_read == 3);
    try testing.expect(mem.eql(u8, dest[0..3], bytes[4..7]));

    amt_read = try fbs.reader().read(&dest);
    try testing.expect(amt_read == 0);

    try fbs.seekTo((try fbs.getEndPos()) + 1);
    amt_read = try fbs.reader().read(&dest);
    try testing.expect(amt_read == 0);
}