authorgravatar for kubkon@jakubkonka.comJakub Konka <kubkon@jakubkonka.com> 2020-05-20 19:42:15+02:00
committergravatar for kubkon@jakubkonka.comJakub Konka <kubkon@jakubkonka.com> 2020-05-29 10:48:03+02:00
logf1a4e1a70f4ecefbae7813d7170f465234c03273
tree81e831c694484f8737df2ca4f52263ba00300c31
parentdc4fea983d358125a3b2a80dc169cad7649de6a9

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`.

1 files changed, 119 insertions(+), 41 deletions(-)

lib/std/process.zig+119-41
...@@ -185,6 +185,91 @@ pub const ArgIteratorPosix = struct {...@@ -185,6 +185,91 @@ pub const ArgIteratorPosix = struct {
185 }185 }
186};186};
187187
188pub const ArgIteratorWasi = struct {
189 allocator: *mem.Allocator,
190 index: usize,
191 args: [][]u8,
192
193 pub const InitError = error{OutOfMemory} || os.UnexpectedError;
194
195 /// You must call deinit to free the internal buffer of the
196 /// iterator after you are done.
197 pub fn init(allocator: *mem.Allocator) InitError!ArgIteratorWasi {
198 const fetched_args = try ArgIteratorWasi.internalInit(allocator);
199 return ArgIteratorWasi{
200 .allocator = allocator,
201 .index = 0,
202 .args = fetched_args,
203 };
204 }
205
206 fn internalInit(allocator: *mem.Allocator) InitError![][]u8 {
207 const w = os.wasi;
208 var count: usize = undefined;
209 var buf_size: usize = undefined;
210
211 switch (w.args_sizes_get(&count, &buf_size)) {
212 w.ESUCCESS => {},
213 else => |err| return os.unexpectedErrno(err),
214 }
215
216 var argv = try allocator.alloc([*:0]u8, count);
217 defer allocator.free(argv);
218
219 var argv_buf = try allocator.alloc(u8, buf_size);
220
221 switch (w.args_get(argv.ptr, argv_buf.ptr)) {
222 w.ESUCCESS => {},
223 else => |err| return os.unexpectedErrno(err),
224 }
225
226 var result_args = try allocator.alloc([]u8, count);
227 var i: usize = 0;
228 while (i < count) : (i += 1) {
229 result_args[i] = mem.spanZ(argv[i]);
230 }
231
232 return result_args;
233 }
234
235 pub fn next(self: *ArgIteratorWasi) ?[]const u8 {
236 if (self.index == self.args.len) return null;
237
238 const arg = self.args[self.index];
239 self.index += 1;
240 return arg;
241 }
242
243 pub fn skip(self: *ArgIteratorWasi) bool {
244 if (self.index == self.args.len) return false;
245
246 self.index += 1;
247 return true;
248 }
249
250 /// Call to free the internal buffer of the iterator.
251 pub fn deinit(self: *ArgIteratorWasi) void {
252 const last_item = self.args[self.args.len - 1];
253 const last_byte_addr = @ptrToInt(last_item.ptr) + last_item.len + 1; // null terminated
254 const first_item_ptr = self.args[0].ptr;
255 const len = last_byte_addr - @ptrToInt(first_item_ptr);
256 self.allocator.free(first_item_ptr[0..len]);
257 self.allocator.free(self.args);
258 }
259};
260
261test "process.ArgIteratorWasi" {
262 if (builtin.os.tag != .wasi) return error.SkipZigTest;
263
264 var ga = std.testing.allocator;
265 var args_it = try ArgIteratorWasi.init(ga);
266 defer args_it.deinit();
267
268 testing.expectEqual(@as(usize, 1), args_it.args.len);
269 const prog_name = args_it.next() orelse unreachable;
270 testing.expect(mem.eql(u8, "test.wasm", prog_name));
271}
272
188pub const ArgIteratorWindows = struct {273pub const ArgIteratorWindows = struct {
189 index: usize,274 index: usize,
190 cmd_line: [*]const u8,275 cmd_line: [*]const u8,
...@@ -335,19 +420,37 @@ pub const ArgIteratorWindows = struct {...@@ -335,19 +420,37 @@ pub const ArgIteratorWindows = struct {
335};420};
336421
337pub const ArgIterator = struct {422pub const ArgIterator = struct {
338 const InnerType = if (builtin.os.tag == .windows) ArgIteratorWindows else ArgIteratorPosix;423 const InnerType = switch (builtin.os.tag) {
424 .windows => ArgIteratorWindows,
425 .wasi => ArgIteratorWasi,
426 else => ArgIteratorPosix,
427 };
339428
340 inner: InnerType,429 inner: InnerType,
341430
431 /// Initialize the args iterator.
432 ///
433 /// On WASI, will panic if the default Wasm page allocator runs out of memory
434 /// or there is an error fetching the args from the runtime. If you want to
435 /// use custom allocator and handle the errors yourself, call `initWasi()` instead.
436 /// You also must remember to free the buffer with `deinitWasi()` call.
342 pub fn init() ArgIterator {437 pub fn init() ArgIterator {
343 if (builtin.os.tag == .wasi) {438 if (builtin.os.tag == .wasi) {
344 // TODO: Figure out a compatible interface accomodating WASI439 const allocator = std.heap.page_allocator;
345 @compileError("ArgIterator is not yet supported in WASI. Use argsAlloc and argsFree instead.");440 return ArgIterator.initWasi(allocator) catch @panic("unexpected error occurred when initializing ArgIterator");
346 }441 }
347442
348 return ArgIterator{ .inner = InnerType.init() };443 return ArgIterator{ .inner = InnerType.init() };
349 }444 }
350445
446 pub const InitError = ArgIteratorWasi.InitError;
447
448 /// If you are targeting WASI, you can call this to manually specify the allocator and
449 /// handle any errors.
450 pub fn initWasi(allocator: *mem.Allocator) InitError!ArgIterator {
451 return ArgIterator{ .inner = try InnerType.init(allocator) };
452 }
453
351 pub const NextError = ArgIteratorWindows.NextError;454 pub const NextError = ArgIteratorWindows.NextError;
352455
353 /// You must free the returned memory when done.456 /// You must free the returned memory when done.
...@@ -364,11 +467,22 @@ pub const ArgIterator = struct {...@@ -364,11 +467,22 @@ pub const ArgIterator = struct {
364 return self.inner.next();467 return self.inner.next();
365 }468 }
366469
470 /// If you only are targeting WASI, you can call this and not need an allocator.
471 pub fn nextWasi(self: *ArgIterator) ?[]const u8 {
472 return self.inner.next();
473 }
474
367 /// Parse past 1 argument without capturing it.475 /// Parse past 1 argument without capturing it.
368 /// Returns `true` if skipped an arg, `false` if we are at the end.476 /// Returns `true` if skipped an arg, `false` if we are at the end.
369 pub fn skip(self: *ArgIterator) bool {477 pub fn skip(self: *ArgIterator) bool {
370 return self.inner.skip();478 return self.inner.skip();
371 }479 }
480
481 /// If you are targeting WASI, call this to free the iterator's internal buffer
482 /// after you are done with it.
483 pub fn deinitWasi(self: *ArgIterator) void {
484 self.inner.deinit();
485 }
372};486};
373487
374pub fn args() ArgIterator {488pub fn args() ArgIterator {
...@@ -377,36 +491,10 @@ pub fn args() ArgIterator {...@@ -377,36 +491,10 @@ pub fn args() ArgIterator {
377491
378/// Caller must call argsFree on result.492/// Caller must call argsFree on result.
379pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 {493pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 {
380 if (builtin.os.tag == .wasi) {
381 var count: usize = undefined;
382 var buf_size: usize = undefined;
383
384 const args_sizes_get_ret = os.wasi.args_sizes_get(&count, &buf_size);
385 if (args_sizes_get_ret != os.wasi.ESUCCESS) {
386 return os.unexpectedErrno(args_sizes_get_ret);
387 }
388
389 var argv = try allocator.alloc([*:0]u8, count);
390 defer allocator.free(argv);
391
392 var argv_buf = try allocator.alloc(u8, buf_size);
393 const args_get_ret = os.wasi.args_get(argv.ptr, argv_buf.ptr);
394 if (args_get_ret != os.wasi.ESUCCESS) {
395 return os.unexpectedErrno(args_get_ret);
396 }
397
398 var result_slice = try allocator.alloc([]u8, count);
399
400 var i: usize = 0;
401 while (i < count) : (i += 1) {
402 result_slice[i] = mem.spanZ(argv[i]);
403 }
404
405 return result_slice;
406 }
407
408 // TODO refactor to only make 1 allocation.494 // TODO refactor to only make 1 allocation.
409 var it = args();495 var it = args();
496 defer if (builtin.os.tag == .wasi) it.deinitWasi();
497
410 var contents = std.ArrayList(u8).init(allocator);498 var contents = std.ArrayList(u8).init(allocator);
411 defer contents.deinit();499 defer contents.deinit();
412500
...@@ -442,16 +530,6 @@ pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 {...@@ -442,16 +530,6 @@ pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 {
442}530}
443531
444pub fn argsFree(allocator: *mem.Allocator, args_alloc: []const []u8) void {532pub fn argsFree(allocator: *mem.Allocator, args_alloc: []const []u8) void {
445 if (builtin.os.tag == .wasi) {
446 const last_item = args_alloc[args_alloc.len - 1];
447 const last_byte_addr = @ptrToInt(last_item.ptr) + last_item.len + 1; // null terminated
448 const first_item_ptr = args_alloc[0].ptr;
449 const len = last_byte_addr - @ptrToInt(first_item_ptr);
450 allocator.free(first_item_ptr[0..len]);
451
452 return allocator.free(args_alloc);
453 }
454
455 var total_bytes: usize = 0;533 var total_bytes: usize = 0;
456 for (args_alloc) |arg| {534 for (args_alloc) |arg| {
457 total_bytes += @sizeOf([]u8) + arg.len;535 total_bytes += @sizeOf([]u8) + arg.len;