Initial import

pick.zig

@@ -0,0 +1,181 @@
+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)}),
+    };
+}