feat(wasi): #180 ADR-015 S4b — env_count/arg_count + multi-import indexing (#290)

Second S4 slice: env + argv counts via WASI preview1 sizes_get
imports. Also a small principled refactor: ctx.wasi_func_indices
records each WASI import's func idx, so multi-import-on-use
indexing is no longer fragile-hardcoded (S4a's `clock_func_idx = 1`
becomes a `List.assoc` lookup; future WASI builtins just append).

- `lib/wasi_runtime.ml`: `create_environ_sizes_get_import` +
  `create_args_sizes_get_import` + shared
  `gen_count_via_sizes_get` emitter (alloc 8B scratch, call
  sizes_get(count_ptr, bufsize_ptr), drop errno, return *count).
- `lib/codegen.ml`:
  • `context.wasi_func_indices : (string * int) list` (additive
    field; default []).
  • Module-assembly refactored: canonical optional WASI imports
    list (clock_time_get, environ_sizes_get, args_sizes_get) added
    on-demand via the existing Effect_sites pre-scan, indices
    assigned by position (fd_write=0; others 1..). The
    `clock_now_ms` builtin now looks up its idx by name (no more
    hardcoded 1) — fully backward-compatible.
  • `env_count(u: Unit) -> Int / Time` and `arg_count(u: Unit) -> Int
    / Time` builtin special-cases. Unit arg = the zero-param-fn
    collapse wart workaround (per existing memory).
- `lib/typecheck.ml` + `lib/resolve.ml`: register the two builtins.

Tests:
- `tests/codegen/env_count.{affine,mjs}` — host stubs
  environ_sizes_get with count 7; asserts 7.
- `tests/codegen/arg_count.{affine,mjs}` — same shape, count 3.
- `tests/codegen/wasi_combo.{affine,mjs}` — **the critical
  regression**: a single unit using ALL three optional WASI imports
  (clock + env + args). Each lookup in `ctx.wasi_func_indices` must
  return the right idx; stubbed values chosen so the sum (100+20+3
  = 123) uniquely identifies any indexing mistake. PASSED.
- `clock_now_ms` (S4a) still passes through the refactor — proves
  the indexing change is lossless.

String accessors (env_at/arg_at) need byte-level wasm IR ops
(I32Load8U/I32Store8 absent in lib/wasm.ml today) — tracked
follow-up as a small wasm-IR extension PR before/with S5.

Gate: `dune test --force` 295/295; full
`tools/run_codegen_wasm_tests.sh` PASSED incl. all 4 new tests
(print, #199 closure ABI, #234 effect-table, #225 CPS, S4a clock,
S4b env/arg/combo, all green). Zero regression.

Refs #180 — S4 done (counts); next: byte-ops slice for strings,
then S5 wasi:filesystem (unblocks INT-06), then S6 flip default +
WIT export lifting. Not Closes.

Author: hyperpolymath

docs/ECOSYSTEM.adoc

@@ -213,15 +213,16 @@ valid WASI-0.2 component via the fetch-pinned preview1→preview2
 `affinescript.ownership` section is proven to SURVIVE the wrap;
 `tests/componentize/smoke.sh` gates it (SKIP-safe without the
 toolchain). Codegen UNCHANGED — core preview1 stays the default
-(reversible). S4a (clock) DONE: `clock_now_ms(clock_id)` builtin
-lowers to a `wasi_snapshot_preview1.clock_time_get` import (added
-on-demand via Effect_sites pre-scan; idx 1 after fd_write at 0;
-zero impact on units that don't use it). Component-path bridges to
-`wasi:clocks` via the reactor adapter; the component is
-structurally valid + ownership-preserving (real-host invocation
-deferred to S6: requires WIT export-lifting / wasi:cli command
-shape). Next: S4b env+argv (same preview1-import pattern), then
-S5 filesystem, S6 flip default.** WIT world of
+(reversible). S4a (clock) + S4b (env_count, arg_count) DONE:
+builtins lower to on-demand `wasi_snapshot_preview1.*` imports
+(Effect_sites pre-scan, canonical-order indexing through
+`ctx.wasi_func_indices`; zero impact on units that don't use them;
+verified with a multi-import combo regression). Component path
+bridges to `wasi:clocks`/`wasi:cli`. Real-host main-invoke deferred
+to S6 (WIT export-lifting / wasi:cli/run command shape). String
+accessors (env_at/arg_at) gated on a byte-level wasm-IR extension
+(I32Load8U/I32Store8 absent today) — tracked as the next slice
+before/with S5 filesystem.** WIT world of
 record: `wit/affinescript.wit`
 |INT-04 |Publish compiler + runtime to JSR (then npm) |#181 |runtime
 packaging READY (affine-js + affinescript-tea JSR dry-run green;

docs/TECH-DEBT.adoc

@@ -185,14 +185,15 @@ follow-up
 multi-ns import object, ownership accessor); 14 unit tests via pinned
 `deno task test` + `tests/modules/loader-bridge/` e2e on real
 compiler-emitted xmod wasm (closes INT-01↔INT-02). Unblocks INT-05/08/11
-|INT-03 |WASI preview2 / host I/O |S1→S4a |#180 ADR-015 (full
+|INT-03 |WASI preview2 / host I/O |S1→S4b |#180 ADR-015 (full
 Component-Model re-target, S1..S6); S2 toolchain #251 closed;
-S3 componentize on-ramp done; **S4a (clock) DONE —
-`clock_now_ms(clock_id)` builtin lowers to a preview1
-`clock_time_get` import (Effect_sites pre-scan, on-demand → zero
-regression on non-clock units); component path bridges to
-wasi:clocks. Real-host main-invoke deferred to S6 (WIT export
-lifting). Next S4b
+S3 componentize done; **S4a (clock) + S4b (env_count, arg_count)
+DONE — on-demand preview1 imports via Effect_sites pre-scan,
+canonical-order indexing through `ctx.wasi_func_indices`; combo
+regression proves no collision. String accessors (env_at/arg_at)
+gated on byte-level wasm IR (I32Load8U/I32Store8 absent today) —
+tracked next slice. Real-host main-invoke = S6 (WIT export
+lifting). Next S5
 (native clocks/env/argv)**
 |INT-04 |Publish to JSR/npm |S2 |#181 packaging READY (dry-run green,
 manual workflow); JSR publish authorised + dispatched (owner go

lib/codegen.ml

@@ -56,6 +56,14 @@ type context = {
   (** Collected ownership annotations: (func_index, param_kinds, return_kind).
       Emitted as the [affinescript.ownership] Wasm custom section for typed-wasm
       Level 7/10 verification. Kind encoding: 0=Unrestricted, 1=Linear, 2=SharedBorrow, 3=ExclBorrow. *)
+  wasi_func_indices : (string * int) list;
+  (** ADR-015 S4 (#180): WASI preview1 import name → wasm func index.
+      Populated at module-assembly time from the optional-imports
+      pre-scan. `fd_write` is always present at 0; other entries
+      (`clock_time_get`, `environ_sizes_get`, `args_sizes_get`, …) are
+      added on-demand in a canonical order, with indices computed by
+      position. WASI builtin special-cases look up their own import
+      index by name from this map. *)
 }
 
 (** Code generation error *)
@@ -100,6 +108,7 @@ let create_context () : context = {
   next_string_offset = 2048;  (* Start strings after heap at offset 2048 *)
   datas = [];
   ownership_annots = [];
+  wasi_func_indices = [];
 }
 
 (** Extract ownership kind from a parameter declaration.
@@ -833,26 +842,58 @@ let rec gen_expr (ctx : context) (expr : expr) : (context * instr list) result =
         Ok (ctx_with_heap, print_code)
 
       | ExprVar id when id.name = "clock_now_ms" && List.length args = 1 ->
-        (* ADR-015 S4a (#180): clock_now_ms(clock_id) — i32 monotonic /
-           realtime milliseconds. Lowers to a `wasi_snapshot_preview1.
-           clock_time_get` call (import idx 1, hardcoded to match the
-           module-assembly ordering). Under the S3 component path the
-           reactor adapter bridges this to wasi:clocks on a real
-           preview2 host. Effect row is `Time` (tracking-only). *)
+        (* ADR-015 S4a (#180): clock_now_ms(clock_id) -> Int ms. Lowers
+           to a `wasi_snapshot_preview1.clock_time_get` call; the import
+           is added on-demand at module assembly (S4b refactor), and the
+           index lives in [ctx.wasi_func_indices] (no hardcoded idx).
+           Under the S3 component path the reactor adapter bridges this
+           to wasi:clocks. Effect row is `Time` (tracking-only). *)
         let* (ctx_with_arg, arg_code) = gen_expr ctx (List.hd args) in
         let (ctx_with_arg2, clock_arg_local) =
           alloc_local ctx_with_arg "__clock_id" in
         let (ctx_with_scratch, scratch_local) =
           alloc_local ctx_with_arg2 "__clock_scratch" in
         let (ctx_with_heap, heap_idx) = ensure_heap_ptr ctx_with_scratch in
-        let clock_func_idx = 1 in
+        let clock_func_idx =
+          try List.assoc "clock_time_get" ctx.wasi_func_indices
+          with Not_found -> 1 (* defensive; pre-scan guarantees presence *)
+        in
         let code =
           arg_code @ [LocalSet clock_arg_local] @
           Wasi_runtime.gen_clock_now_ms
             heap_idx clock_arg_local scratch_local clock_func_idx
         in
         Ok (ctx_with_heap, code)
 
+      | ExprVar id when (id.name = "env_count" || id.name = "arg_count")
+                        && List.length args = 1 ->
+        (* ADR-015 S4b (#180): env_count(u: Unit) / arg_count(u: Unit)
+           — i32 count returns. Lower to the matching
+           `wasi_snapshot_preview1.{environ,args}_sizes_get` import
+           (added on-demand at module assembly; idx looked up in
+           [ctx.wasi_func_indices]). The Unit arg satisfies the
+           zero-param-fn collapse wart; it is evaluated but its value
+           is unused. String accessors (env_at/arg_at) need byte-level
+           wasm IR ops (currently absent) and are a tracked follow-up. *)
+        let wasi_name =
+          if id.name = "env_count" then "environ_sizes_get"
+          else "args_sizes_get"
+        in
+        let sizes_func_idx =
+          try List.assoc wasi_name ctx.wasi_func_indices
+          with Not_found -> 1
+        in
+        let* (ctx_with_arg, arg_code) = gen_expr ctx (List.hd args) in
+        let (ctx_with_scratch, scratch_local) =
+          alloc_local ctx_with_arg ("__" ^ id.name ^ "_scratch") in
+        let (ctx_with_heap, heap_idx) = ensure_heap_ptr ctx_with_scratch in
+        let code =
+          arg_code @ [Drop] @
+          Wasi_runtime.gen_count_via_sizes_get
+            heap_idx scratch_local sizes_func_idx
+        in
+        Ok (ctx_with_heap, code)
+
       | ExprVar id when List.mem_assoc id.name ctx.variant_tags ->
         (* Enum constructor called as a function: Circle(5), Rect({x:1,y:2}), etc.
            Layout: [tag: i32][field1: i32][field2: i32]...
@@ -2488,37 +2529,50 @@ let generate_module ?loader (prog : program) : wasm_module result =
   let fd_write_type_idx = 0 in  (* Will be first type *)
   let fd_write_import_fixed = { fd_write_import with i_desc = ImportFunc fd_write_type_idx } in
 
-  (* ADR-015 S4a (#180): register WASI `clock_time_get` only when the
-     unit actually calls `clock_now_ms` — adding it unconditionally
-     would force every host (incl. every test harness) to stub it,
-     breaking the principle "import what you use". Pre-scan via the
-     shared Effect_sites traversal. When emitted, the clock takes
-     import idx 1 (deterministically after fd_write at 0), which the
-     `clock_now_ms` builtin hardcodes. *)
-  let needs_clock =
+  (* ADR-015 S4 (#180): register WASI preview1 imports ON DEMAND so
+     non-using units stay byte-identical to pre-S4 (the "import what
+     you use" principle — adding unconditionally would force every
+     host stub them and break every test harness). Pre-scan with the
+     shared Effect_sites traversal: each builtin name maps 1:1 to its
+     WASI import; appended after fd_write (idx 0) in a CANONICAL ORDER
+     so the index assignment is deterministic across compilations.
+     Each builtin's gen_expr case looks up its index by name in
+     `ctx.wasi_func_indices` (no hardcoded indices — clean
+     multi-import indexing). *)
+  let uses (builtin : string) : bool =
     Effect_sites.fold_calls
       (fun acc _ord call ->
         acc ||
         match call with
-        | ExprApp (ExprVar id, _) -> id.name = "clock_now_ms"
+        | ExprApp (ExprVar id, _) -> id.name = builtin
         | _ -> false)
       false prog
   in
+  let optional_wasi =
+    (* (guest_builtin_name, wasi_import_name, factory) — canonical order. *)
+    [ ("clock_now_ms", "clock_time_get",     Wasi_runtime.create_clock_time_get_import);
+      ("env_count",    "environ_sizes_get",  Wasi_runtime.create_environ_sizes_get_import);
+      ("arg_count",    "args_sizes_get",     Wasi_runtime.create_args_sizes_get_import);
+    ]
+    |> List.filter_map
+         (fun (b, w, f) -> if uses b then Some (w, f ()) else None)
+    |> List.mapi (fun i (w, (imp, ty)) -> (i + 1, w, imp, ty))
+  in
+  let opt_types = List.map (fun (_, _, _, ty) -> ty) optional_wasi in
+  let opt_imports =
+    List.map (fun (idx, _, imp, _) -> { imp with i_desc = ImportFunc idx })
+      optional_wasi
+  in
+  let wasi_indices =
+    ("fd_write", fd_write_type_idx) ::
+    List.map (fun (idx, name, _, _) -> (name, idx)) optional_wasi
+  in
   let ctx_with_wasi =
-    if needs_clock then
-      let (clock_import, clock_type) = Wasi_runtime.create_clock_time_get_import () in
-      let clock_type_idx = 1 in
-      let clock_import_fixed = { clock_import with i_desc = ImportFunc clock_type_idx } in
-      {
-        ctx with
-        types = fd_write_type :: clock_type :: ctx.types;
-        imports = fd_write_import_fixed :: clock_import_fixed :: ctx.imports;
-      }
-    else
       {
         ctx with
-        types = fd_write_type :: ctx.types;
-        imports = fd_write_import_fixed :: ctx.imports;
+        types = fd_write_type :: opt_types @ ctx.types;
+        imports = fd_write_import_fixed :: opt_imports @ ctx.imports;
+        wasi_func_indices = wasi_indices;
       }
   in
 

lib/resolve.ml

@@ -55,6 +55,8 @@ let seed_builtins (symbols : Symbol.t) : unit =
   def "print"; def "println"; def "eprint"; def "eprintln";
   (* WASI time (ADR-015 S4a, #180) *)
   def "clock_now_ms";
+  (* WASI env / argv counts (ADR-015 S4b, #180) *)
+  def "env_count"; def "arg_count";
   (* String / char builtins *)
   def "len"; def "slice"; def "string_get"; def "string_sub"; def "string_find";
   def "char_to_int"; def "int_to_char"; def "show";

lib/typecheck.ml

@@ -1322,6 +1322,13 @@ let register_builtins (ctx : context) : unit =
      `wasi:clocks`. *)
   bind_var ctx "clock_now_ms"
     (TArrow (ty_int, QOmega, ty_int, ESingleton "Time"));
+  (* ADR-015 S4b (#180): WASI environment / argv COUNTS. The Unit arg
+     satisfies the zero-param-fn collapse wart (`fn()->T` lowers to
+     bare `T`; callable zero-arg builtins take `Unit -> R`). String
+     accessors (env_at/arg_at) need byte-level wasm IR ops — tracked
+     follow-up. Effect row `Time` (reserved). *)
+  bind_var ctx "env_count" (TArrow (ty_unit, QOmega, ty_int, ESingleton "Time"));
+  bind_var ctx "arg_count" (TArrow (ty_unit, QOmega, ty_int, ESingleton "Time"));
   bind_var ctx "eprint" (TArrow (ty_string, QOmega, ty_unit, ESingleton "IO"));
   bind_var ctx "eprintln" (TArrow (ty_string, QOmega, ty_unit, ESingleton "IO"));
   bind_var ctx "read_line"

lib/wasi_runtime.ml

@@ -354,3 +354,62 @@ let gen_print_str (heap_ptr_global : int) (str_ptr_local : int) (fd_write_idx :
     Drop;
     I32Const 0l;
   ]
+
+(** Create the WASI `environ_sizes_get` import (ADR-015 S4b, #180).
+    Signature: `(envc_out: i32, envbuf_size_out: i32) -> errno: i32`.
+    Writes the env-var count and the total byte size of the
+    null-terminated `KEY=VAL\0…` buffer the next call would need.
+    String accessor (`env_at`) is gated on byte-level wasm IR ops,
+    deferred to a follow-up slice. *)
+let create_environ_sizes_get_import () : import * func_type =
+  let func_type = {
+    ft_params = [I32; I32];  (* envc_out_ptr, envbuf_size_out_ptr *)
+    ft_results = [I32];      (* errno *)
+  } in
+  let import = {
+    i_module = "wasi_snapshot_preview1";
+    i_name = "environ_sizes_get";
+    i_desc = ImportFunc 0;
+  } in
+  (import, func_type)
+
+(** Create the WASI `args_sizes_get` import (ADR-015 S4b, #180).
+    Signature: `(argc_out: i32, argv_buf_size_out: i32) -> errno: i32`. *)
+let create_args_sizes_get_import () : import * func_type =
+  let func_type = {
+    ft_params = [I32; I32];
+    ft_results = [I32];
+  } in
+  let import = {
+    i_module = "wasi_snapshot_preview1";
+    i_name = "args_sizes_get";
+    i_desc = ImportFunc 0;
+  } in
+  (import, func_type)
+
+(** Emit `env_count`/`arg_count`: call the appropriate `*_sizes_get`
+    import (which writes count + buf_size into two i32 scratch slots),
+    drop errno, return the count as i32. Uniform helper for the two
+    builtins — they differ only in the import index. *)
+let gen_count_via_sizes_get
+    (heap_ptr_global : int)
+    (scratch_local : int)
+    (sizes_func_idx : int)
+    : instr list =
+  [
+    (* scratch = heap; heap += 8 (two i32 slots: count, buf_size) *)
+    GlobalGet heap_ptr_global;
+    I32Const 8l; I32Add;
+    GlobalSet heap_ptr_global;
+    GlobalGet heap_ptr_global;
+    I32Const 8l; I32Sub;
+    LocalSet scratch_local;
+    (* sizes_get(count_ptr, bufsize_ptr); drop errno *)
+    LocalGet scratch_local;                       (* count_ptr *)
+    LocalGet scratch_local; I32Const 4l; I32Add;  (* bufsize_ptr *)
+    Call sizes_func_idx;
+    Drop;
+    (* return *count_ptr *)
+    LocalGet scratch_local;
+    I32Load (2, 0);
+  ]

tests/codegen/arg_count.affine

@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: PMPL-1.0-or-later
+// ADR-015 S4b (#180): arg_count() smoke. Lowers to
+// `wasi_snapshot_preview1.args_sizes_get`. Same shape as env_count;
+// proves the multi-WASI-import-on-use indexing handles two optional
+// imports simultaneously without collision.
+pub fn main() -> Int / { Time } {
+  arg_count(())
+}

tests/codegen/env_count.affine

@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: PMPL-1.0-or-later
+// ADR-015 S4b (#180): env_count() smoke. Lowers to
+// `wasi_snapshot_preview1.environ_sizes_get`, returns the count
+// the host wrote into *count_ptr. Unit arg = zero-param-fn collapse
+// workaround (`fn() -> T` lowers to bare `T`; builtins use `Unit -> R`).
+pub fn main() -> Int / { Time } {
+  env_count(())
+}

tests/codegen/test_arg_count.mjs

@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: PMPL-1.0-or-later
+// ADR-015 S4b (#180) — arg_count via WASI preview1 args_sizes_get.
+import assert from 'node:assert/strict';
+import { readFile } from 'node:fs/promises';
+
+const buf = await readFile('./tests/codegen/arg_count.wasm');
+let inst = null;
+let observed = null;
+
+const imports = {
+  wasi_snapshot_preview1: {
+    fd_write: () => 0,
+    args_sizes_get: (argc_ptr, argv_buf_ptr) => {
+      observed = { argc_ptr, argv_buf_ptr };
+      const dv = new DataView(inst.exports.memory.buffer);
+      dv.setUint32(argc_ptr, 3, true);
+      dv.setUint32(argv_buf_ptr, 32, true);
+      return 0;
+    },
+  },
+};
+
+inst = (await WebAssembly.instantiate(buf, imports)).instance;
+const result = inst.exports.main();
+
+assert.ok(observed, 'guest called args_sizes_get');
+assert.equal(result, 3, 'arg_count returns the count');
+console.log('test_arg_count.mjs OK');

tests/codegen/test_env_count.mjs

@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: PMPL-1.0-or-later
+// ADR-015 S4b (#180) — env_count via WASI preview1 environ_sizes_get.
+// Host stubs the import with a known count + buf_size; asserts the
+// guest returns that count.
+import assert from 'node:assert/strict';
+import { readFile } from 'node:fs/promises';
+
+const buf = await readFile('./tests/codegen/env_count.wasm');
+let inst = null;
+let observed = null;
+
+const imports = {
+  wasi_snapshot_preview1: {
+    fd_write: () => 0,
+    environ_sizes_get: (envc_ptr, envbuf_ptr) => {
+      observed = { envc_ptr, envbuf_ptr };
+      const dv = new DataView(inst.exports.memory.buffer);
+      dv.setUint32(envc_ptr, 7, true);        // 7 env vars
+      dv.setUint32(envbuf_ptr, 256, true);    // unused by env_count
+      return 0;
+    },
+  },
+};
+
+inst = (await WebAssembly.instantiate(buf, imports)).instance;
+const result = inst.exports.main();
+
+assert.ok(observed, 'guest called environ_sizes_get');
+assert.equal(result, 7, 'env_count returns the count the host wrote');
+console.log('test_env_count.mjs OK');