用 Zig 优化 yes 命令

[jiacai2050]

起因是看到一篇文章，作者介绍了如何用 Rust 优化 yes 命令，第一个 buffer 的版本还比较好懂，第二个复用 buffer 的就没有那么直接了。想了下用 Zig 实现会是怎么样？于是就有了下面的测试：

PS：测试时会用的 pv 这个命令

Buffer Write 版本

const std = @import("std");

pub fn main() !void {
    const argv = std.os.argv;

    const output: []const u8 = if (argv.len > 1)
        argv.ptr[1][0..std.mem.len(argv.ptr[1])]
    else
        "y";

    const out = std.io.getStdOut();
    var f = std.io.bufferedWriter(out.writer());
    while ((try f.write(output)) > 0) {
        _ = try f.write("\n");
    }
}

编译运行

zig build-exe yes.zig -Drelease-safe
./yes | pv -r > /dev/null

大概是 116Mib/s

复用 Buffer 版本

const std = @import("std");

const BUFFER_CAP = 64 * 1024;

fn fill_up_buffer(buf: *[BUFFER_CAP]u8, output: []const u8) []const u8 {
    if (output.len > buf.len / 2) {
        return output;
    }

    std.mem.copy(u8, buf, output);
    std.mem.copy(u8, buf[output.len..], "\n");
    var buffer_size = output.len + 1;
    while (buffer_size < buf.len / 2) {
        std.mem.copy(u8, buf[buffer_size..], buf[0..buffer_size]);
        buffer_size *= 2;
    }

    return buf;
}

pub fn main() !void {
    const argv = std.os.argv;

    const output: []const u8 = if (argv.len > 1)
        argv.ptr[1][0..std.mem.len(argv.ptr[1])]
    else
        "y";

    var buffer: [BUFFER_CAP]u8 = undefined;
    const body = fill_up_buffer(&buffer, output);
    const stdout = std.io.getStdOut();
    var writer = stdout.writer();
    while ((try writer.write(body)) > 0) {}
}

同样地方式编译执行，速率大概是： 5.4Gib/s

而且，这个版本相比 Rust 的实现，个人感觉，要清晰不少。

[jiacai2050]

第二个版本有几个 corner case，重新写了个

//! `yes` unix command in Zig, optimized for speed
//! Reference to: https://github.com/cgati/yes

const std = @import("std");

const BUFFER_CAP = 64 * 1024;

fn fill_up_buffer(buf: *[BUFFER_CAP]u8, word: []const u8) []const u8 {
    if (word.len > buf.len / 2) {
        return word;
    }

    std.mem.copy(u8, buf, word);
    var buffer_size = word.len;
    while (buffer_size < buf.len / 2) {
        std.mem.copy(u8, buf[buffer_size..], buf[0..buffer_size]);
        buffer_size *= 2;
    }

    return buf[0..buffer_size];
}

/// Reference to: https://doc.rust-lang.org/std/borrow/enum.Cow.html
const Cow = union(enum) {
    owned: std.ArrayList(u8),
    ref: []const u8,

    fn as_slice(self: Cow) []const u8 {
        return switch (self) {
            .owned => |o| o.items,
            .ref => |r| r,
        };
    }

    fn deinit(self: Cow) void {
        return switch (self) {
            .owned => |o| o.deinit(),
            .ref => {},
        };
    }
};

pub fn main() !void {
    const argv = std.os.argv;

    const word = if (argv.len > 1) x: {
        var list = std.ArrayList(u8).init(std.heap.page_allocator);
        try list.appendSlice(argv.ptr[1][0..std.mem.len(argv.ptr[1])]);
        try list.append('\n');
        break :x Cow{ .owned = list };
    } else Cow{ .ref = "y\n" };
    defer word.deinit();

    var buffer: [BUFFER_CAP]u8 = undefined;
    const body = fill_up_buffer(&buffer, word.as_slice());
    const stdout = std.io.getStdOut();
    var writer = stdout.writer();
    while (true) {
        try writer.writeAll(body);
    }
}

[jiacai2050]

• https://liujiacai.net/blog/2022/11/17/yes-in-zig/
  后续会更新到这里