zig/lib/std/io/CountingWriter.zig

const std = @import("../std.zig");
const CountingWriter = @This();
const assert = std.debug.assert;
const native_endian = @import("builtin").target.cpu.arch.endian();
const Writer = std.io.Writer;
const testing = std.testing;

/// Underlying stream to passthrough bytes to.
child_writer: Writer,
bytes_written: u64 = 0,

pub fn writer(cw: *CountingWriter) Writer {
    return .{
        .context = cw,
        .vtable = &.{
            .writeSplat = passthru_writeSplat,
            .writeFile = passthru_writeFile,
        },
    };
}

pub fn unbufferedWriter(cw: *CountingWriter) std.io.BufferedWriter {
    return .{
        .buffer = &.{},
        .unbuffered_writer = writer(cw),
    };
}

fn passthru_writeSplat(context: *anyopaque, data: []const []const u8, splat: usize) anyerror!usize {
    const cw: *CountingWriter = @alignCast(@ptrCast(context));
    const n = try cw.child_writer.writeSplat(data, splat);
    cw.bytes_written += n;
    return n;
}

fn passthru_writeFile(
    context: *anyopaque,
    file: std.fs.File,
    offset: u64,
    len: Writer.VTable.FileLen,
    headers_and_trailers: []const []const u8,
    headers_len: usize,
) anyerror!usize {
    const cw: *CountingWriter = @alignCast(@ptrCast(context));
    const n = try cw.child_writer.writeFile(file, offset, len, headers_and_trailers, headers_len);
    cw.bytes_written += n;
    return n;
}

test CountingWriter {
    var cw: CountingWriter = .{ .child_writer = std.io.null_writer };
    var bw = cw.unbufferedWriter();
    const bytes = "yay";
    try bw.writeAll(bytes);
    try testing.expect(cw.bytes_written == bytes.len);
}