From f1a4e1a70f4ecefbae7813d7170f465234c03273 Mon Sep 17 00:00:00 2001 From: Jakub Konka Date: Wed, 20 May 2020 19:42:15 +0200 Subject: [PATCH 1/3] Add ArgIteratorWasi and integrate it with ArgIterator This commit pulls WASI specific implementation of args extraction from the runtime from `process.argsAlloc` and `process.argsFree` into a new iterator struct `process.ArgIteratorWasi`. It also integrates the struct with platform-independent `process.ArgIterator`. --- lib/std/process.zig | 160 ++++++++++++++++++++++++++++++++------------ 1 file changed, 119 insertions(+), 41 deletions(-) diff --git a/lib/std/process.zig b/lib/std/process.zig index 8c1feeffb2..4546ea3def 100644 --- a/lib/std/process.zig +++ b/lib/std/process.zig @@ -185,6 +185,91 @@ pub const ArgIteratorPosix = struct { } }; +pub const ArgIteratorWasi = struct { + allocator: *mem.Allocator, + index: usize, + args: [][]u8, + + pub const InitError = error{OutOfMemory} || os.UnexpectedError; + + /// You must call deinit to free the internal buffer of the + /// iterator after you are done. + pub fn init(allocator: *mem.Allocator) InitError!ArgIteratorWasi { + const fetched_args = try ArgIteratorWasi.internalInit(allocator); + return ArgIteratorWasi{ + .allocator = allocator, + .index = 0, + .args = fetched_args, + }; + } + + fn internalInit(allocator: *mem.Allocator) InitError![][]u8 { + const w = os.wasi; + var count: usize = undefined; + var buf_size: usize = undefined; + + switch (w.args_sizes_get(&count, &buf_size)) { + w.ESUCCESS => {}, + else => |err| return os.unexpectedErrno(err), + } + + var argv = try allocator.alloc([*:0]u8, count); + defer allocator.free(argv); + + var argv_buf = try allocator.alloc(u8, buf_size); + + switch (w.args_get(argv.ptr, argv_buf.ptr)) { + w.ESUCCESS => {}, + else => |err| return os.unexpectedErrno(err), + } + + var result_args = try allocator.alloc([]u8, count); + var i: usize = 0; + while (i < count) : (i += 1) { + result_args[i] = mem.spanZ(argv[i]); + } + + return result_args; + } + + pub fn next(self: *ArgIteratorWasi) ?[]const u8 { + if (self.index == self.args.len) return null; + + const arg = self.args[self.index]; + self.index += 1; + return arg; + } + + pub fn skip(self: *ArgIteratorWasi) bool { + if (self.index == self.args.len) return false; + + self.index += 1; + return true; + } + + /// Call to free the internal buffer of the iterator. + pub fn deinit(self: *ArgIteratorWasi) void { + const last_item = self.args[self.args.len - 1]; + const last_byte_addr = @ptrToInt(last_item.ptr) + last_item.len + 1; // null terminated + const first_item_ptr = self.args[0].ptr; + const len = last_byte_addr - @ptrToInt(first_item_ptr); + self.allocator.free(first_item_ptr[0..len]); + self.allocator.free(self.args); + } +}; + +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, @@ -335,19 +420,37 @@ pub const ArgIteratorWindows = struct { }; pub const ArgIterator = struct { - const InnerType = if (builtin.os.tag == .windows) ArgIteratorWindows else ArgIteratorPosix; + const InnerType = switch (builtin.os.tag) { + .windows => ArgIteratorWindows, + .wasi => ArgIteratorWasi, + else => ArgIteratorPosix, + }; 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) { - // TODO: Figure out a compatible interface accomodating WASI - @compileError("ArgIterator is not yet supported in WASI. Use argsAlloc and argsFree instead."); + const allocator = std.heap.page_allocator; + return ArgIterator.initWasi(allocator) catch @panic("unexpected error occurred when initializing ArgIterator"); } return ArgIterator{ .inner = InnerType.init() }; } + 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) }; + } + pub const NextError = ArgIteratorWindows.NextError; /// You must free the returned memory when done. @@ -364,11 +467,22 @@ pub const ArgIterator = struct { return self.inner.next(); } + /// If you only are targeting WASI, you can call this and not need an allocator. + pub fn nextWasi(self: *ArgIterator) ?[]const u8 { + return self.inner.next(); + } + /// Parse past 1 argument without capturing it. /// Returns `true` if skipped an arg, `false` if we are at the end. pub fn skip(self: *ArgIterator) bool { 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(); + } }; pub fn args() ArgIterator { @@ -377,36 +491,10 @@ pub fn args() ArgIterator { /// Caller must call argsFree on result. pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 { - if (builtin.os.tag == .wasi) { - var count: usize = undefined; - var buf_size: usize = undefined; - - const args_sizes_get_ret = os.wasi.args_sizes_get(&count, &buf_size); - if (args_sizes_get_ret != os.wasi.ESUCCESS) { - return os.unexpectedErrno(args_sizes_get_ret); - } - - var argv = try allocator.alloc([*:0]u8, count); - defer allocator.free(argv); - - var argv_buf = try allocator.alloc(u8, buf_size); - const args_get_ret = os.wasi.args_get(argv.ptr, argv_buf.ptr); - if (args_get_ret != os.wasi.ESUCCESS) { - return os.unexpectedErrno(args_get_ret); - } - - var result_slice = try allocator.alloc([]u8, count); - - var i: usize = 0; - while (i < count) : (i += 1) { - result_slice[i] = mem.spanZ(argv[i]); - } - - return result_slice; - } - // TODO refactor to only make 1 allocation. var it = args(); + defer if (builtin.os.tag == .wasi) it.deinitWasi(); + var contents = std.ArrayList(u8).init(allocator); defer contents.deinit(); @@ -442,16 +530,6 @@ pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 { } pub fn argsFree(allocator: *mem.Allocator, args_alloc: []const []u8) void { - if (builtin.os.tag == .wasi) { - const last_item = args_alloc[args_alloc.len - 1]; - const last_byte_addr = @ptrToInt(last_item.ptr) + last_item.len + 1; // null terminated - const first_item_ptr = args_alloc[0].ptr; - const len = last_byte_addr - @ptrToInt(first_item_ptr); - allocator.free(first_item_ptr[0..len]); - - return allocator.free(args_alloc); - } - var total_bytes: usize = 0; for (args_alloc) |arg| { total_bytes += @sizeOf([]u8) + arg.len; From 6f48842ddb2d8ca8fec639dc11a47753058ae325 Mon Sep 17 00:00:00 2001 From: Jakub Konka Date: Fri, 29 May 2020 08:40:32 +0200 Subject: [PATCH 2/3] Make ArgIterator.init() a compile error in WASI Given that the previous design would require the use of a default allocator to have `ArgIterator.init()` work in WASI, and since in Zig we're trying to avoid default allocators, I've changed the design slightly in that now `init()` is a compile error in WASI, and instead in its message it points to `initWithAllocator(*mem.Allocator)`. The latter by virtue of requiring an allocator as an argument can safely be used in WASI as well as on other OSes (where the allocator argument is simply unused). When using `initWithAllocator` it is then natural to remember to call `deinit()` after being done with the iterator. Also, to make use of this, I've also added `argsWithAllocator` function which is equivalent to `args` minus the requirement of supplying an allocator and being fallible. Finally, I've also modified the WASI only test `process.ArgWasiIterator` to test all OSes. --- lib/std/process.zig | 68 ++++++++++++++++++++++++++------------------- 1 file changed, 39 insertions(+), 29 deletions(-) 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(); From 6e347e6180077b2b4243b6dd2cd1ef7a811e7881 Mon Sep 17 00:00:00 2001 From: Jakub Konka Date: Fri, 29 May 2020 17:12:19 +0200 Subject: [PATCH 3/3] Fix args iterator test --- lib/std/process.zig | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/lib/std/process.zig b/lib/std/process.zig index 744f192f32..a9400e8ce1 100644 --- a/lib/std/process.zig +++ b/lib/std/process.zig @@ -484,25 +484,28 @@ pub fn argsWithAllocator(allocator: *mem.Allocator) ArgIterator.InitError!ArgIte test "args iterator" { var ga = std.testing.allocator; - var it = if (builtin.os.tag == .wasi) argsWithAllocator(ga) else args(); + var it = if (builtin.os.tag == .wasi) try argsWithAllocator(ga) else args(); defer it.deinit(); // no-op unless WASI - testing.expect(it.skip()); - const prog_name = it.next(ga) orelse unreachable; + const prog_name = try it.next(ga) orelse unreachable; defer ga.free(prog_name); - const expected_bin_name = switch (builtin.os.tag) { + const expected_suffix = switch (builtin.os.tag) { .wasi => "test.wasm", .windows => "test.exe", else => "test", }; - testing.expect(mem.eql(u8, expected_bin_name, prog_name)); + const given_suffix = std.fs.path.basename(prog_name); + + testing.expect(mem.eql(u8, expected_suffix, given_suffix)); + testing.expectEqual(it.next(ga), null); + testing.expect(!it.skip()); } /// Caller must call argsFree on result. pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 { // TODO refactor to only make 1 allocation. - var it = if (builtin.os.tag == .wasi) argsWithAllocator(allocator) else args(); + var it = if (builtin.os.tag == .wasi) try argsWithAllocator(allocator) else args(); defer it.deinit(); var contents = std.ArrayList(u8).init(allocator);