Tree(1) in Zig

tree.zig

//! Tree(1) in Zig
//! https://linux.die.net/man/1/tree

const std = @import("std");
const process = std.process;
const fs = std.fs;
const mem = std.mem;
const testing = std.testing;

const VisitedEntry = struct {
    name: []const u8,
    level: usize,
};

const VisitedList = std.ArrayList(VisitedEntry);

pub fn main() anyerror!void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    const args = try process.argsAlloc(allocator);
    defer process.argsFree(allocator, args);

    const root_dir = if (args.len > 1)
        args[1]
    else
        ".";

    var visited = VisitedList.init(allocator);
    try visited.append(.{ .name = root_dir, .level = 0 });

    var iter_dir =
        try fs.cwd().openIterableDir(root_dir, .{});
    defer iter_dir.close();

    try walk(allocator, &iter_dir, &visited, 1);

    var writer = std.io.bufferedWriter(std.io.getStdOut().writer());
    for (visited.items) |entry| {
        _ = try writer.write(try genPrefix(allocator, entry.level));
        _ = try writer.write(entry.name);
        _ = try writer.write("\n");
    }
    try writer.flush();
}

fn genPrefix(allocator: std.mem.Allocator, level: usize) ![]const u8 {
    const prefix = "|  ";
    const base = "|--";
    return switch (level) {
        0 => return try allocator.dupe(u8, ""),
        1 => return try allocator.dupe(u8, base),
        else => {
            var out = try allocator.alloc(u8, (level - 1) * prefix.len + base.len);
            for (0..level - 1) |i| {
                mem.copy(u8, out[i * prefix.len ..], prefix);
            }
            mem.copy(u8, out[(level - 1) * prefix.len ..], base);

            return out;
        },
    };
}

test "testing genPrefix" {
    const allocator = testing.allocator;
    const testcases = [_]std.meta.Tuple(&[_]type{ usize, []const u8 }){
        .{ 0, "" },
        .{ 1, "|--" },
        .{ 2, "|  |--" },
        .{ 3, "|  |  |--" },
    };
    for (testcases) |case| {
        const actual = try genPrefix(allocator, case.@"0");
        defer allocator.free(actual);
        try testing.expectEqualStrings(case.@"1", actual);
    }
}

fn walk(
    allocator: mem.Allocator,
    iter_dir: *fs.IterableDir,
    visited: *VisitedList,
    level: usize,
) !void {
    var it = iter_dir.iterate();
    while (try it.next()) |entry| {
        const dupe_name = try allocator.dupe(u8, entry.name);
        switch (entry.kind) {
            .Directory => {
                try visited.append(.{ .name = dupe_name, .level = level });
                var sub_iter_dir = try iter_dir.dir.openIterableDir(entry.name, .{});
                defer sub_iter_dir.close();

                try walk(
                    allocator,
                    &sub_iter_dir,
                    visited,
                    level + 1,
                );
            },
            else => {
                try visited.append(.{ .name = dupe_name, .level = level });
            },
        }
    }
}

Example

./zig-out/bin/tree zig-cache/
zig-cache/
|--z
|  |--6c2830de7eff2e6c29a466f3dae29f56
|  |--a2d9c9576a2de120dc2fe54d26e22b56
|--o
|  |--31a508b05079f502e9de10cf725ab900
|  |  |--tree
|  |  |--tree.o
|  |--c3a7328c0b6d9e4ac745e62c729ae197
|  |  |--build.o
|  |  |--build
|  |--d4336c38182f8d4321b881a61531ae3c
|  |  |--builtin.zig
|  |--d7233815bac1be707323362604efcc18
|  |  |--dependencies.zig
|  |--d4ac96c2b77544f5a6d309f1f10d02bd
|  |  |--builtin.zig
|--h
|  |--ef8df946ad597ff441d79907e875840e.txt
|  |--acd9b03fd406a6d60e3dbfb18849aace.txt
|  |--timestamp
|--tmp