authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2020-05-29 15:55:14-04:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2020-05-29 15:55:14-04:00
log65c3833ec2793747961f532510feb4fbb52fd6a6
treeafdafb4d4ce49a0beeca19ea97c79bb96083b3f9
parente9f344dcd43dd2a725ab728b962d7627d6462d02
parent6e347e6180077b2b4243b6dd2cd1ef7a811e7881
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #5388 from kubkon/wasi-args-iter

Add ArgIteratorWasi and integrate it with ArgIterator

1 files changed, 130 insertions(+), 39 deletions(-)

lib/std/process.zig+130-39
......@@ -185,6 +185,79 @@ pub const ArgIteratorPosix = struct {
185185 }
186186};
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
188261pub const ArgIteratorWindows = struct {
189262 index: usize,
190263 cmd_line: [*]const u8,
......@@ -335,14 +408,29 @@ pub const ArgIteratorWindows = struct {
335408};
336409
337410pub const ArgIterator = struct {
338 const InnerType = if (builtin.os.tag == .windows) ArgIteratorWindows else ArgIteratorPosix;
411 const InnerType = switch (builtin.os.tag) {
412 .windows => ArgIteratorWindows,
413 .wasi => ArgIteratorWasi,
414 else => ArgIteratorPosix,
415 };
339416
340417 inner: InnerType,
341418
419 /// Initialize the args iterator.
342420 pub fn init() ArgIterator {
343421 if (builtin.os.tag == .wasi) {
344 // TODO: Figure out a compatible interface accomodating WASI
345 @compileError("ArgIterator is not yet supported in WASI. Use argsAlloc and argsFree instead.");
422 @compileError("In WASI, use initWithAllocator instead.");
423 }
424
425 return ArgIterator{ .inner = InnerType.init() };
426 }
427
428 pub const InitError = ArgIteratorWasi.InitError;
429
430 /// You must deinitialize iterator's internal buffers by calling `deinit` when done.
431 pub fn initWithAllocator(allocator: *mem.Allocator) InitError!ArgIterator {
432 if (builtin.os.tag == .wasi) {
433 return ArgIterator{ .inner = try InnerType.init(allocator) };
346434 }
347435
348436 return ArgIterator{ .inner = InnerType.init() };
......@@ -364,49 +452,62 @@ pub const ArgIterator = struct {
364452 return self.inner.next();
365453 }
366454
455 /// If you only are targeting WASI, you can call this and not need an allocator.
456 pub fn nextWasi(self: *ArgIterator) ?[]const u8 {
457 return self.inner.next();
458 }
459
367460 /// Parse past 1 argument without capturing it.
368461 /// Returns `true` if skipped an arg, `false` if we are at the end.
369462 pub fn skip(self: *ArgIterator) bool {
370463 return self.inner.skip();
371464 }
465
466 /// Call this to free the iterator's internal buffer if the iterator
467 /// was created with `initWithAllocator` function.
468 pub fn deinit(self: *ArgIterator) void {
469 // Unless we're targeting WASI, this is a no-op.
470 if (builtin.os.tag == .wasi) {
471 self.inner.deinit();
472 }
473 }
372474};
373475
374476pub fn args() ArgIterator {
375477 return ArgIterator.init();
376478}
377479
378/// Caller must call argsFree on result.
379pub 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);
480/// You must deinitialize iterator's internal buffers by calling `deinit` when done.
481pub fn argsWithAllocator(allocator: *mem.Allocator) ArgIterator.InitError!ArgIterator {
482 return ArgIterator.initWithAllocator(allocator);
483}
391484
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 }
485test "args iterator" {
486 var ga = std.testing.allocator;
487 var it = if (builtin.os.tag == .wasi) try argsWithAllocator(ga) else args();
488 defer it.deinit(); // no-op unless WASI
397489
398 var result_slice = try allocator.alloc([]u8, count);
490 const prog_name = try it.next(ga) orelse unreachable;
491 defer ga.free(prog_name);
399492
400 var i: usize = 0;
401 while (i < count) : (i += 1) {
402 result_slice[i] = mem.spanZ(argv[i]);
403 }
493 const expected_suffix = switch (builtin.os.tag) {
494 .wasi => "test.wasm",
495 .windows => "test.exe",
496 else => "test",
497 };
498 const given_suffix = std.fs.path.basename(prog_name);
404499
405 return result_slice;
406 }
500 testing.expect(mem.eql(u8, expected_suffix, given_suffix));
501 testing.expectEqual(it.next(ga), null);
502 testing.expect(!it.skip());
503}
407504
505/// Caller must call argsFree on result.
506pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 {
408507 // TODO refactor to only make 1 allocation.
409 var it = args();
508 var it = if (builtin.os.tag == .wasi) try argsWithAllocator(allocator) else args();
509 defer it.deinit();
510
410511 var contents = std.ArrayList(u8).init(allocator);
411512 defer contents.deinit();
412513
......@@ -442,16 +543,6 @@ pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 {
442543}
443544
444545pub 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
455546 var total_bytes: usize = 0;
456547 for (args_alloc) |arg| {
457548 total_bytes += @sizeOf([]u8) + arg.len;