test-nanoexroinfo.zig

test-nanoexrinfo.zig

const std = @import("std");
const Queue = std.mpmc.MPMCQueue;
const sync = std.sync;
const Allocator = std.heap.ArenaAllocator;
const io = std.io;
const os = std.os;

const MAX_TASKS = 100;

pub fn ThreadPool(allocator: *Allocator, num_threads: usize) anyerror!void {
    const thread_locals = try allocator.allocUnmanaged([]const u8{});
    defer allocator.free(thread_locals);

    const task_queue = Queue(allocator);
    const is_shutdown = @AtomicFlag.init(false);

    for (num_threads) |i| {
        thread_locals.* = undefined;
        std.thread.spawnDetached(threadFunc, .{&is_shutdown, &task_queue, thread_locals});
    }

    pub fn submit(fn: fn() anyerror!void) void {
        try task_queue.push(fn);
    }

    pub fn shutdown() void {
        is_shutdown.store(true, .SeqCst);
    }

    fn threadFunc(arg: []void) anyerror!void {
        var is_shutdown = @as(atomic, arg[0]).load(.SeqCst);
        var task_queue: [*fn() anyerror!void, MAX_TASKS] = @ptrCast([*fn() anyerror!void, MAX_TASKS]void, arg[1]);
        var thread_locals = arg[2].*;

        while (!is_shutdown) : (is_shutdown) {
            var maybe_task: *fn() anyerror!void = null;
            try task_queue.pop(&maybe_task);

            if (maybe_task) {
                defer maybe_task();
            } else {
                // No tasks available, sleep and try again
                std.os.sleep(std.duration.milliseconds(10));
            }
        }
    }
}

pub fn recurse_directory(directory_path: []const u8, thread_pool: *ThreadPool, allocator: *Allocator) anyerror!void {
    const dir = try io.dir.openDir(allocator, directory_path);
    defer dir.close();

    while (dir.hasNext()) : (|entry| entry.close()) {
        var entry = try dir.next();
        if (entry.type == .Dir) {
            // Recursively process the subdirectory in a new task
            thread_pool.submit(fn() anyerror!void {
                try recurse_directory(entry.path, thread_pool, allocator);
            });
        } else if (entry.type == .File && entry.name.endsWith(".exr")) {
            // Print the file name in a new task
            thread_pool.submit(fn() anyerror!void {
                try print_file_name(entry.path);
            });
        }
    }
}

pub fn main() !void {
    const allocator = std.heap.ArenaAllocator.init(std.heap.page_allocator, 1024 * 1024);
    defer allocator.deinit();

    var args = std.process.args();
    if (args.len() != 2) {
        return std.log.err("usage: {} <directory>\n", args[0]);
    }

    const thread_pool = try ThreadPool(&allocator, 4);
    defer thread_pool.shutdown();

    const directory_path = try args[1].dup();
    defer allocator.free(directory_path);

    try recurse_directory(directory_path, &thread_pool, &allocator);

    return Ok.Void;
}