競プロでよく見かける標準入力からの読み込みの支援方法を考えてみた。

read_input.zig

const std = @import("std");

pub fn main() !void {
    const count = try input(struct {x: usize}, 12);

    for (0..count.x) |_| {
        // タプルの分解構文は、0.12.x以降で可能
        // 0.11.xの場合、そのままタプルで受ける必要あり
        const x, const y = try input(struct { [16:0]u8, usize }, 1024);
        std.debug.print("x = {s}, y = {}\n", .{x, y});
    }
}

pub fn input(comptime Result: type, comptime buffer_size: usize) !Result {
    std.debug.assert(@typeInfo(Result) == .Struct);

    const stdin = std.io.getStdIn();
    var buffer: [buffer_size]u8 = undefined;

    var s = try stdin.reader().readUntilDelimiter(buffer, '\n');
    var iter = std.mem.splitAny(u8, s, " \r\n");

    return readInput(Result, &iter);
}

fn readInput(comptime Result: type, iter: *std.mem.SplitIterator(u8, .any)) !Result {
    const fields = std.meta.fields(Result);
    var values: Result = undefined;

    inline for (fields) |field| {
        pick_value: while (iter.next()) |v| {    
            const type_info = @typeInfo(field.type);
            if ((type_info == .Array) and (type_info.Array.child == u8)) {
                // アロケータなしでサイズ不明なスライスを構成できないため、代わりに配列を使う
                // 手を抜くために、センチネルとして0を指定し、NULL終端文字列として構成する。
                // std.mem.sliceToを使えばzig標準のスライスに変更可能
                std.debug.assert(v.len <= type_info.Array.len);

                const sentinel = @as(*u8, @ptrCast(@constCast(type_info.Array.sentinel.?))).*;
                std.debug.assert(sentinel == 0);

                @memset(&@field(values, field.name), 0);
                std.mem.copyForwards(u8, &@field(values, field.name), v);   
                break :pick_value;
            }
            else if (type_info == .Int) {
                @field(values, field.name) = try std.fmt.parseInt(field.type, v, 0);
                break :pick_value;
            }
            else if (type_info == .Float) {
                @field(values, field.name) = try std.fmt.parseFloat(field.type, v);
                break :pick_value;
            }
            else {
                @compileError(std.fmt.comptimePrint("Unsupported type `{s}`", .{@typeName(field.type)}));
            }
        }
    }

    return values;    
}

test "read input to struct" {
    var sample = sample: {
        var sample_buf: [1024]u8 = undefined;
        var stream = std.io.fixedBufferStream(&sample_buf);
        var writer = stream.writer();
        try writer.writeAll("123 456\r\n");
        try writer.writeByte(0);
        try stream.seekTo(0);
        break :sample stream;
    };

    var buffer: [1024]u8 = undefined;    
    var s = try sample.reader().readUntilDelimiter(&buffer, '\n');
    var iter = std.mem.splitAny(u8, s, " \r\n");

    const values = try readInput(struct { x: usize, y: usize }, &iter);

    try std.testing.expectEqual(values.x, 123);
    try std.testing.expectEqual(values.y, 456);
}

test "read input to tuple" {
    var iter = std.mem.splitAny(u8, "123 456\r\n", " \r\n");

    const x, const y = try readInput(struct { usize, usize }, &iter);

    try std.testing.expectEqual(x, 123);
    try std.testing.expectEqual(y, 456);
}

read_input_alloc.zig

/// アロケータを使って良いのであれば、スライスや配列のマッピングもできそう
const std = @import("std");

pub fn main() !void {
    var arena_buffer: [4096]u8 = undefined;
    var arena = std.heap.FixedBufferAllocator.init(&arena_buffer);

    const s, const arr = try inputAlloc(arena.allocator(), struct{ []const u8, []const usize }, 1024);

    std.debug.print("s = {s}\n", .{s});
    std.debug.print("arr = {any}\n", .{arr});
}

pub fn inputAlloc(allocator: std.mem.Allocator, comptime Result: type, comptime buffer_size: usize) !Result {
    std.debug.assert(@typeInfo(Result) == .Struct);

    const stdin = std.io.getStdIn();
    var buffer: [buffer_size]u8 = undefined;

    var s = try stdin.reader().readUntilDelimiter(&buffer, '\n');
    var iter = std.mem.splitAny(u8, s, " \r\n");

    return readInputAlloc(allocator, Result, &iter);
}

fn readInputAlloc(allocator: std.mem.Allocator, comptime Result: type, iter: *std.mem.SplitIterator(u8, .any)) !Result {
    const fields = std.meta.fields(Result);
    var values: Result = undefined;

    inline for (fields) |field| {
        pick_value: while (iter.next()) |v| {  
            if (v.len == 0) continue;

            const type_info = @typeInfo(field.type);
            if ((type_info == .Array) and (type_info.Array.child == u8)) {
                std.debug.assert(v.len <= type_info.Array.len);

                const sentinel = @as(*u8, @ptrCast(@constCast(type_info.Array.sentinel.?))).*;
                std.debug.assert(sentinel == 0);

                @memset(&@field(values, field.name), 0);
                std.mem.copyForwards(u8, &@field(values, field.name), v);   
                break :pick_value;
            }
            else if (field.type == []const u8) {
                @field(values, field.name) = try allocator.dupe(u8, v);
                break :pick_value;
            }
            else if ((type_info == .Pointer) and (type_info.Pointer.size == .Slice)) {
                const size = try std.fmt.parseInt(usize, v, 0);
                var array = try allocator.alloc(type_info.Pointer.child, size);

                const child_type_info = @typeInfo(type_info.Pointer.child);
                for (0..array.len) |i| {
                    while (iter.next()) |vv| {
                        if (vv.len == 0) continue;

                        if (type_info.Pointer.child == []const u8) {
                            array[i] = try allocator.dupe(u8, vv);
                            break;
                        }
                        else if (child_type_info == .Int) {
                            array[i] = try std.fmt.parseInt(type_info.Pointer.child, vv, 0);
                            break;
                        }
                        else if (child_type_info == .Float) {
                            array[i] = try std.fmt.parseFloat(type_info.Pointer.child, vv);
                            break;
                        }
                        else {
                            @compileError(std.fmt.comptimePrint("Unsupported type: {s}", .{@typeName(type_info.Pointer.child)}));
                        }
                    }
                }

                @field(values, field.name) = array;
                break :pick_value;
            }
            else if (type_info == .Int) {
                @field(values, field.name) = try std.fmt.parseInt(field.type, v, 0);
                break :pick_value;
            }
            else if (type_info == .Float) {
                @field(values, field.name) = try std.fmt.parseFloat(field.type, v);
                break :pick_value;
            }
            else {
                @compileError(std.fmt.comptimePrint("Unsupported type `{s}`", .{@typeName(field.type)}));
            }
        }
    }

    return values;    
}

test "read input to struct (string)" {
    var iter = std.mem.splitAny(u8, "abc qwerty\r\n", " \r\n");

    var arena_buffer: [4096]u8 = undefined;
    var arena = std.heap.FixedBufferAllocator.init(&arena_buffer);

    const s1, const s2 = try readInputAlloc(arena.allocator(), struct { []const u8, []const u8 }, &iter);

    try std.testing.expectEqualSlices(u8, s1, "abc");
    try std.testing.expectEqualSlices(u8, s2, "qwerty");
}

test "read input to array (int)" {
    var iter = std.mem.splitAny(u8, "4 7 -3 11 9\r\n", " \r\n");

    var arena_buffer: [4096]u8 = undefined;
    var arena = std.heap.FixedBufferAllocator.init(&arena_buffer);

    const x = try readInputAlloc(arena.allocator(), struct { []const i32 }, &iter);

    try std.testing.expectEqualSlices(i32, x[0], &[_]i32{ 7, -3, 11, 9 });
}

test "read input to array (string)" {
    var iter = std.mem.splitAny(u8, "3 xyz 11 qwerty\r\n", " \r\n");

    var arena_buffer: [4096]u8 = undefined;
    var arena = std.heap.FixedBufferAllocator.init(&arena_buffer);

    const x = try readInputAlloc(arena.allocator(), struct { array: []const []const u8 }, &iter);

    try std.testing.expectEqual(x.array.len, 3);
    try std.testing.expectEqualSlices(u8, x.array[0], "xyz");
    try std.testing.expectEqualSlices(u8, x.array[1], "11");
    try std.testing.expectEqualSlices(u8, x.array[2], "qwerty");
}