diff --git a/lib/std/process.zig b/lib/std/process.zig index 4546ea3def..744f192f32 100644 --- a/lib/std/process.zig +++ b/lib/std/process.zig @@ -258,18 +258,6 @@ pub const ArgIteratorWasi = struct { } }; -test "process.ArgIteratorWasi" { - if (builtin.os.tag != .wasi) return error.SkipZigTest; - - var ga = std.testing.allocator; - var args_it = try ArgIteratorWasi.init(ga); - defer args_it.deinit(); - - testing.expectEqual(@as(usize, 1), args_it.args.len); - const prog_name = args_it.next() orelse unreachable; - testing.expect(mem.eql(u8, "test.wasm", prog_name)); -} - pub const ArgIteratorWindows = struct { index: usize, cmd_line: [*]const u8, @@ -429,15 +417,9 @@ pub const ArgIterator = struct { inner: InnerType, /// Initialize the args iterator. - /// - /// On WASI, will panic if the default Wasm page allocator runs out of memory - /// or there is an error fetching the args from the runtime. If you want to - /// use custom allocator and handle the errors yourself, call `initWasi()` instead. - /// You also must remember to free the buffer with `deinitWasi()` call. pub fn init() ArgIterator { if (builtin.os.tag == .wasi) { - const allocator = std.heap.page_allocator; - return ArgIterator.initWasi(allocator) catch @panic("unexpected error occurred when initializing ArgIterator"); + @compileError("In WASI, use initWithAllocator instead."); } return ArgIterator{ .inner = InnerType.init() }; @@ -445,10 +427,13 @@ pub const ArgIterator = struct { pub const InitError = ArgIteratorWasi.InitError; - /// If you are targeting WASI, you can call this to manually specify the allocator and - /// handle any errors. - pub fn initWasi(allocator: *mem.Allocator) InitError!ArgIterator { - return ArgIterator{ .inner = try InnerType.init(allocator) }; + /// You must deinitialize iterator's internal buffers by calling `deinit` when done. + pub fn initWithAllocator(allocator: *mem.Allocator) InitError!ArgIterator { + if (builtin.os.tag == .wasi) { + return ArgIterator{ .inner = try InnerType.init(allocator) }; + } + + return ArgIterator{ .inner = InnerType.init() }; } pub const NextError = ArgIteratorWindows.NextError; @@ -478,10 +463,13 @@ pub const ArgIterator = struct { return self.inner.skip(); } - /// If you are targeting WASI, call this to free the iterator's internal buffer - /// after you are done with it. - pub fn deinitWasi(self: *ArgIterator) void { - self.inner.deinit(); + /// Call this to free the iterator's internal buffer if the iterator + /// was created with `initWithAllocator` function. + pub fn deinit(self: *ArgIterator) void { + // Unless we're targeting WASI, this is a no-op. + if (builtin.os.tag == .wasi) { + self.inner.deinit(); + } } }; @@ -489,11 +477,33 @@ pub fn args() ArgIterator { return ArgIterator.init(); } +/// You must deinitialize iterator's internal buffers by calling `deinit` when done. +pub fn argsWithAllocator(allocator: *mem.Allocator) ArgIterator.InitError!ArgIterator { + return ArgIterator.initWithAllocator(allocator); +} + +test "args iterator" { + var ga = std.testing.allocator; + var it = if (builtin.os.tag == .wasi) argsWithAllocator(ga) else args(); + defer it.deinit(); // no-op unless WASI + + testing.expect(it.skip()); + const prog_name = it.next(ga) orelse unreachable; + defer ga.free(prog_name); + + const expected_bin_name = switch (builtin.os.tag) { + .wasi => "test.wasm", + .windows => "test.exe", + else => "test", + }; + testing.expect(mem.eql(u8, expected_bin_name, prog_name)); +} + /// Caller must call argsFree on result. pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 { // TODO refactor to only make 1 allocation. - var it = args(); - defer if (builtin.os.tag == .wasi) it.deinitWasi(); + var it = if (builtin.os.tag == .wasi) argsWithAllocator(allocator) else args(); + defer it.deinit(); var contents = std.ArrayList(u8).init(allocator); defer contents.deinit();