generate: add multi-threading

1 files changed, 60 insertions, 0 deletions

diff --git a/src/scheduler.zig b/src/scheduler.zig
new file mode 100644
index 0000000..b2ab0db
--- /dev/null
+++ b/src/scheduler.zig
@@ -0,0 +1,60 @@
+const std = @import("std");
+
+pub fn Scheduler(function: anytype) type {
+	const FuncType = @TypeOf(function);
+	const func_info = @typeInfo(FuncType).Fn;
+	const R = func_info.return_type orelse std.builtin.Type.Void;
+	const ArgType = std.meta.ArgsTuple(FuncType);
+
+	const wrapper = struct {
+		pub fn call(args: ArgType, result: *R) void {
+			result.* = @call(.auto, function, args);
+		}
+	}.call;
+
+	return struct {
+		const Self = @This();
+
+		allocator: std.mem.Allocator,
+		pool: std.Thread.Pool,
+		wait_group: std.Thread.WaitGroup,
+		results: std.ArrayList(*R),
+
+		pub fn init(self: *Self, allocator: std.mem.Allocator) !void {
+			self.allocator = allocator;
+			self.results = std.ArrayList(*R).init(allocator);
+			self.wait_group.reset();
+
+			try self.pool.init(.{
+				.allocator = allocator,
+			});
+		}
+
+		pub fn push_task(self: *Self, args: ArgType) !void {
+			const result = try self.allocator.create(R);
+			try self.results.append(result);
+			self.pool.spawnWg(
+				&self.wait_group,
+				wrapper,
+				.{ args, result }
+			);
+		}
+
+		pub fn deinit(self: *Self) ![]R {
+			self.pool.waitAndWork(&self.wait_group);
+			self.pool.deinit();
+
+			const results = try self.allocator.alloc(R, self.results.items.len);
+			
+			for (self.results.items, 0..) |result, index| {
+				results[index] = result.*;
+				self.allocator.destroy(result);
+			}
+
+			self.results.deinit();
+			self.* = undefined;
+
+			return results;
+		}
+	};
+}