pick/pick.zig

const std = @import("std");

fn usage(file: std.Io.File, io: std.Io) !void {
    var buf: [1024]u8 = undefined;
    var w = file.writer(io, &buf);
    try w.interface.writeAll(
        \\usage: pick [--header] FIELD...
        \\
        \\Select tab-separated fields from stdin. FIELD is a 1-based column number,
        \\or a header name when --header is set. Output is TSV/plain records only.
        \\
        \\options:
        \\  --header  resolve FIELD names from the first input record
        \\  --help    show this help on stdout and exit 0
        \\
        \\stdout: selected records only
        \\stderr: usage errors and diagnostics only
        \\exits: 0 on success, nonzero on invalid fields or I/O failure
        \\
        \\No regex, expressions, arithmetic, shell, mutation, files, JSON, color,
        \\prompts, config, network, daemon behavior, or hidden state.
        \\
        \\examples:
        \\  printf 'path\\tsize\\nREADME\\t120\\n' | pick --header path size
        \\  printf 'a\\tb\\n' | pick 2
        \\
    );
    try w.interface.flush();
}

fn die(io: std.Io, comptime fmt: []const u8, args: anytype) noreturn {
    var buf: [512]u8 = undefined;
    var w = std.Io.File.stderr().writer(io, &buf);
    w.interface.print(fmt, args) catch {};
    w.interface.flush() catch {};
    std.process.exit(1);
}

fn fieldAt(line: []const u8, index: usize) []const u8 {
    var it = std.mem.splitScalar(u8, line, '\t');
    var i: usize = 0;
    while (it.next()) |field| : (i += 1) {
        if (i == index) return field;
    }
    return "";
}

fn stripRecordNewlineCr(raw: []const u8) []const u8 {
    if (raw.len > 0 and raw[raw.len - 1] == '\r') return raw[0 .. raw.len - 1];
    return raw;
}

fn headerIndex(header: []const u8, name: []const u8, io: std.Io) usize {
    var it = std.mem.splitScalar(u8, header, '\t');
    var i: usize = 0;
    while (it.next()) |field| : (i += 1) {
        if (std.mem.eql(u8, field, name)) return i;
    }
    die(io, "pick: header not found: {s}\n", .{name});
}

fn writeAll(out: *std.Io.File.Writer, bytes: []const u8) std.Io.File.Writer.Error!void {
    out.interface.writeAll(bytes) catch |err| switch (err) {
        error.WriteFailed => return out.err.?,
    };
}

fn writeByte(out: *std.Io.File.Writer, byte: u8) std.Io.File.Writer.Error!void {
    out.interface.writeByte(byte) catch |err| switch (err) {
        error.WriteFailed => return out.err.?,
    };
}

fn writePicked(out: *std.Io.File.Writer, line: []const u8, cols: []const usize) !void {
    for (cols, 0..) |col, i| {
        if (i > 0) try writeByte(out, '\t');
        try writeAll(out, fieldAt(line, col));
    }
    try writeByte(out, '\n');
}

fn handleLine(
    allocator: std.mem.Allocator,
    io: std.Io,
    has_header: bool,
    names: []const []const u8,
    cols: *std.ArrayList(usize),
    header_done: *bool,
    out: *std.Io.File.Writer,
    raw: []const u8,
) !void {
    const line = stripRecordNewlineCr(raw);
    if (has_header and !header_done.*) {
        for (names) |name| try cols.append(allocator, headerIndex(line, name, io));
        try writePicked(out, line, cols.items);
        header_done.* = true;
    } else {
        try writePicked(out, line, cols.items);
    }
}

fn pickStream(
    allocator: std.mem.Allocator,
    io: std.Io,
    has_header: bool,
    names: []const []const u8,
    cols: *std.ArrayList(usize),
) !void {
    var out_buf: [8192]u8 = undefined;
    var out = std.Io.File.stdout().writer(io, &out_buf);
    var carry = std.ArrayList(u8).empty;
    defer carry.deinit(allocator);

    var header_done = !has_header;
    var buf: [8192]u8 = undefined;
    while (true) {
        const n = std.posix.read(std.posix.STDIN_FILENO, &buf) catch |err| die(io, "pick: stdin: {s}\n", .{@errorName(err)});
        if (n == 0) break;

        var start: usize = 0;
        for (buf[0..n], 0..) |byte, i| {
            if (byte != '\n') continue;
            if (carry.items.len == 0) {
                try handleLine(allocator, io, has_header, names, cols, &header_done, &out, buf[start..i]);
            } else {
                try carry.appendSlice(allocator, buf[start..i]);
                try handleLine(allocator, io, has_header, names, cols, &header_done, &out, carry.items);
                carry.clearRetainingCapacity();
            }
            start = i + 1;
        }
        if (start < n) try carry.appendSlice(allocator, buf[start..n]);
    }

    if (carry.items.len > 0) {
        try handleLine(allocator, io, has_header, names, cols, &header_done, &out, carry.items);
    } else if (has_header and !header_done) {
        try handleLine(allocator, io, has_header, names, cols, &header_done, &out, "");
    }
    try out.flush();
}

pub fn main(init: std.process.Init) !void {
    const allocator = init.gpa;
    var args = try std.process.Args.Iterator.initAllocator(init.minimal.args, allocator);
    defer args.deinit();
    _ = args.skip();

    var has_header = false;
    var names = std.ArrayList([]const u8).empty;
    defer names.deinit(allocator);
    while (args.next()) |arg| {
        if (std.mem.eql(u8, arg, "--help")) {
            try usage(std.Io.File.stdout(), init.io);
            return;
        } else if (std.mem.eql(u8, arg, "--header")) {
            has_header = true;
        } else {
            try names.append(allocator, arg);
        }
    }
    if (names.items.len == 0) {
        try usage(std.Io.File.stderr(), init.io);
        std.process.exit(2);
    }

    var cols = std.ArrayList(usize).empty;
    defer cols.deinit(allocator);
    if (!has_header) {
        for (names.items) |name| {
            const n = std.fmt.parseInt(usize, name, 10) catch die(init.io, "pick: invalid field: {s}\n", .{name});
            if (n == 0) die(init.io, "pick: fields are 1-based\n", .{});
            try cols.append(allocator, n - 1);
        }
    }

    pickStream(allocator, init.io, has_header, names.items, &cols) catch |err| switch (err) {
        error.BrokenPipe => return,
        else => die(init.io, "pick: output: {s}\n", .{@errorName(err)}),
    };
}