wasi_runtime.ml

(* SPDX-License-Identifier: MPL-2.0 *)
(* SPDX-FileCopyrightText: 2024-2025 hyperpolymath *)

(** WASI runtime support - I/O bindings for WebAssembly System Interface.

    This module provides helper functions to generate WASM code that calls
    WASI system calls for I/O operations.
*)

open Wasm

(** WASI file descriptor constants *)
let fd_stdout = 1l
let fd_stderr = 2l

(** WASI clock-id constants (preview1 `wasi_snapshot_preview1`):
    0 = REALTIME, 1 = MONOTONIC, 2 = PROCESS_CPUTIME, 3 = THREAD_CPUTIME. *)
let clock_realtime = 0l
let clock_monotonic = 1l

(** Create the WASI `clock_time_get` import (ADR-015 S4a, #180).

    Signature: `(clockid: i32, precision: i64, time_out: i32) -> errno: i32`.
    Writes the timestamp (nanoseconds, i64) to `*time_out`. Through the
    S3 componentize on-ramp the wasmtime preview1->preview2 reactor
    adapter bridges this to `wasi:clocks/{monotonic,wall}-clock@0.2.*`
    on a real host. *)
let create_clock_time_get_import () : import * func_type =
  let func_type = {
    ft_params = [I32; I64; I32];  (* clockid, precision, time_out_ptr *)
    ft_results = [I32];           (* errno *)
  } in
  let import = {
    i_module = "wasi_snapshot_preview1";
    i_name = "clock_time_get";
    i_desc = ImportFunc 0;  (* Will be adjusted when added to module *)
  } in
  (import, func_type)

(** Emit a `clock_now_ms(clock_id)` sequence. The caller has placed
    [clock_arg_local] (the `i32` clock id) and allocated
    [scratch_local] (an `i32` heap-pointer local) and threaded the
    import [clock_func_idx]. Leaves the i32 monotonic/realtime
    millisecond count on the stack (= ns / 1_000_000, wrapped to i32 —
    documented lossy: 32-bit ms wraps after ~24 days, sufficient for
    typical use; the precise i64 reader is a follow-up).

    Layout: 8 bytes scratch (the i64 time_out the WASI host writes). *)
let gen_clock_now_ms
    (heap_ptr_global : int) (clock_arg_local : int)
    (scratch_local : int) (clock_func_idx : int)
    : instr list =
  [
    (* scratch = heap; heap += 8 *)
    GlobalGet heap_ptr_global;
    I32Const 8l; I32Add;
    GlobalSet heap_ptr_global;
    GlobalGet heap_ptr_global;
    I32Const 8l; I32Sub;
    LocalSet scratch_local;
    (* clock_time_get(clock_id, 0 /* precision */, scratch); drop errno *)
    LocalGet clock_arg_local;
    I64Const 0L;
    LocalGet scratch_local;
    Call clock_func_idx;
    Drop;
    (* return (i32) (i64_load(scratch) / 1_000_000) *)
    LocalGet scratch_local;
    I64Load (3, 0);
    I64Const 1_000_000L;
    I64DivU;
    I32WrapI64;
  ]

(** Create WASI fd_write import

    fd_write signature: (fd: i32, iovs: i32, iovs_len: i32, nwritten: i32) -> i32

    Returns (import, func_type_index)
*)
let create_fd_write_import () : import * func_type =
  let func_type = {
    ft_params = [I32; I32; I32; I32];  (* fd, iovs_ptr, iovs_len, nwritten_ptr *)
    ft_results = [I32];                 (* error code *)
  } in
  let import = {
    i_module = "wasi_snapshot_preview1";
    i_name = "fd_write";
    i_desc = ImportFunc 0;  (* Will be adjusted when added to module *)
  } in
  (import, func_type)

(** Generate code to print an integer to stdout

    This function:
    1. Converts int to string in memory (digit by digit)
    2. Creates iovec structure pointing to string
    3. Calls WASI fd_write

    Algorithm:
    - Handle negative sign separately
    - Extract digits using division and modulo
    - Write digits in reverse order (most significant first)

    Memory layout:
    - [0..15]: digit buffer (max 11 chars: "-2147483648")
    - [16..23]: iovec struct (buf_ptr, buf_len)
    - [24..27]: nwritten result

    Returns: code that leaves 0 on stack for success
*)
let gen_print_int (heap_ptr_global : int) (fd_write_idx : int) (num_local : int)
    : instr list =
  (* We need additional locals for the conversion:
     - buf_ptr: pointer to start of allocated buffer
     - write_ptr: current position while writing digits (counts backwards)
     - n: the number being converted (copy of num_local)
     - digit: current digit being extracted
     - is_negative: 1 if negative, 0 if positive
     - len: final string length
  *)

  (* Allocate 28 bytes total *)
  let total_size = 28 in
  let buf_local = num_local + 1 in       (* reuse locals after num_local *)
  let write_ptr_local = num_local + 2 in
  let n_local = num_local + 3 in
  let digit_local = num_local + 4 in
  let is_neg_local = num_local + 5 in
  let len_local = num_local + 6 in

  [
    (* Allocate memory *)
    GlobalGet heap_ptr_global;
    I32Const (Int32.of_int total_size);
    I32Add;
    GlobalSet heap_ptr_global;

    (* buf_ptr = heap_ptr - 28 *)
    GlobalGet heap_ptr_global;
    I32Const (Int32.of_int total_size);
    I32Sub;
    LocalTee buf_local;

    (* Copy number to n_local, check if negative *)
    LocalGet num_local;
    LocalTee n_local;
    I32Const 0l;
    I32LtS;  (* n < 0 ? *)
    LocalTee is_neg_local;

    (* If negative, negate n to make it positive *)
    If (BtEmpty, [
      I32Const 0l;
      LocalGet n_local;
      I32Sub;  (* n = -n *)
      LocalSet n_local;
    ], []);

    (* write_ptr starts at buf + 15 (end of buffer, write backwards) *)
    LocalGet buf_local;
    I32Const 15l;
    I32Add;
    LocalSet write_ptr_local;

    (* Extract digits loop - write from right to left *)
    Block (BtEmpty, [
      Loop (BtEmpty, [
        (* digit = n % 10 *)
        LocalGet n_local;
        I32Const 10l;
        I32RemU;
        LocalTee digit_local;

        (* *write_ptr = '0' + digit *)
        I32Const 48l;  (* ASCII '0' *)
        I32Add;
        LocalGet write_ptr_local;
        I32Store (0, 0);

        (* write_ptr-- *)
        LocalGet write_ptr_local;
        I32Const 1l;
        I32Sub;
        LocalSet write_ptr_local;

        (* n = n / 10 *)
        LocalGet n_local;
        I32Const 10l;
        I32DivU;
        LocalTee n_local;

        (* Continue if n > 0 *)
        I32Const 0l;
        I32GtU;
        BrIf 0;
      ])
    ]);

    (* If negative, add '-' sign *)
    LocalGet is_neg_local;
    If (BtEmpty, [
      I32Const 45l;  (* ASCII '-' *)
      LocalGet write_ptr_local;
      I32Store (0, 0);
      LocalGet write_ptr_local;
      I32Const 1l;
      I32Sub;
      LocalSet write_ptr_local;
    ], []);

    (* Calculate length: (buf + 15) - write_ptr *)
    LocalGet buf_local;
    I32Const 15l;
    I32Add;
    LocalGet write_ptr_local;
    I32Sub;
    LocalSet len_local;

    (* Actual string starts at write_ptr + 1 *)
    LocalGet write_ptr_local;
    I32Const 1l;
    I32Add;
    LocalSet write_ptr_local;

    (* Create iovec structure at buf + 16 *)
    (* iovec.buf_ptr = write_ptr (where string actually starts) *)
    LocalGet buf_local;
    I32Const 16l;
    I32Add;
    LocalGet write_ptr_local;
    I32Store (2, 0);

    (* iovec.buf_len = len *)
    LocalGet buf_local;
    I32Const 16l;
    I32Add;
    LocalGet len_local;
    I32Store (2, 4);

    (* Call fd_write(stdout, iovec_ptr, 1, nwritten_ptr) *)
    I32Const fd_stdout;
    LocalGet buf_local;
    I32Const 16l;
    I32Add;                       (* iovs = buf + 16 *)
    I32Const 1l;                  (* iovs_len = 1 *)
    LocalGet buf_local;
    I32Const 24l;
    I32Add;                       (* nwritten = buf + 24 *)
    Call fd_write_idx;

    (* Drop error code, return success *)
    Drop;
    I32Const 0l;
  ]

(** Generate code to print a newline *)
let gen_println (heap_ptr_global : int) (fd_write_idx : int) (temp_local : int)
    : instr list =
  let newline_byte = 10l in  (* ASCII '\n' *)
  (* Layout — every i32 store sits on a 4-byte boundary:
       offset  0..3:  iovec.buf_ptr (i32) — points at offset 12 below
       offset  4..7:  iovec.buf_len (i32) — always 1
       offset  8..11: nwritten     (i32) — fd_write writes here
       offset 12:     newline byte  ('\n', 1 byte)
       offset 13..15: padding so subsequent allocs stay aligned
     Total 16 bytes. The previous layout stored the newline at temp+0 and
     the iovec at temp+1, which traps under wasmtime's strict alignment
     check. *)
  [
    GlobalGet heap_ptr_global;
    I32Const 16l;
    I32Add;
    GlobalSet heap_ptr_global;

    GlobalGet heap_ptr_global;
    I32Const 16l;
    I32Sub;
    LocalSet temp_local;

    (* iovec.buf_ptr = temp + 12 (where the newline byte lives) *)
    LocalGet temp_local;
    LocalGet temp_local;
    I32Const 12l;
    I32Add;
    I32Store (2, 0);

    (* iovec.buf_len = 1 *)
    LocalGet temp_local;
    I32Const 1l;
    I32Store (2, 4);

    (* Write the newline byte at temp + 12. We use I32Store to lay down a
       full 4 bytes — the upper 3 are slack but harmless because fd_write
       only reads buf_len = 1. *)
    LocalGet temp_local;
    I32Const newline_byte;
    I32Store (2, 12);

    (* Call fd_write(stdout, iovec_ptr=temp, iovs_len=1, nwritten_ptr=temp+8) *)
    I32Const fd_stdout;
    LocalGet temp_local;
    I32Const 1l;
    LocalGet temp_local;
    I32Const 8l;
    I32Add;
    Call fd_write_idx;

    Drop;
    I32Const 0l;
  ]

(** Generate code to print a string (length-prefixed in memory).
    String layout: [len: i32][bytes...]
    Returns: code that leaves 0 on stack for success *)
let gen_print_str (heap_ptr_global : int) (str_ptr_local : int) (fd_write_idx : int) (temp_local : int)
    : instr list =
  [
    (* Allocate 12 bytes for iovec + nwritten *)
    GlobalGet heap_ptr_global;
    I32Const 12l;
    I32Add;
    GlobalSet heap_ptr_global;

    GlobalGet heap_ptr_global;
    I32Const 12l;
    I32Sub;
    LocalSet temp_local;

    (* iovec.buf_ptr = str_ptr + 4 *)
    LocalGet temp_local;
    LocalGet str_ptr_local;
    I32Const 4l;
    I32Add;
    I32Store (2, 0);

    (* iovec.buf_len = *str_ptr (length) *)
    LocalGet temp_local;
    I32Const 4l;
    I32Add;
    LocalGet str_ptr_local;
    I32Load (2, 0);
    I32Store (2, 0);

    (* Call fd_write(stdout, iovec_ptr, 1, nwritten_ptr) *)
    I32Const fd_stdout;
    LocalGet temp_local;           (* iovs *)
    I32Const 1l;                   (* iovs_len *)
    LocalGet temp_local;
    I32Const 8l;
    I32Add;                        (* nwritten *)
    Call 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.
    Paired with `environ_get` (created by
    {!create_environ_get_import}) for the `env_at` string accessor. *)
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);
  ]

(** Create the WASI `environ_get` import (ADR-015 S5, #180).
    Signature: `(environ_ptr_ptr: i32, environ_buf_ptr: i32) -> errno: i32`.
    Fills two regions: a vector of pointers (one per env-var, written
    at `environ_ptr_ptr`) and a contiguous buffer of null-terminated
    `KEY=VAL` strings (written at `environ_buf_ptr`). The sizes that
    must be allocated are reported by `environ_sizes_get`. *)
let create_environ_get_import () : import * func_type =
  let func_type = {
    ft_params = [I32; I32];  (* environ_ptr_ptr, environ_buf_ptr *)
    ft_results = [I32];      (* errno *)
  } in
  let import = {
    i_module = "wasi_snapshot_preview1";
    i_name = "environ_get";
    i_desc = ImportFunc 0;
  } in
  (import, func_type)

(** Create the WASI `args_get` import (ADR-015 S5, #180).
    Signature: `(argv_ptr_ptr: i32, argv_buf_ptr: i32) -> errno: i32`.
    Same shape as `environ_get`. *)
let create_args_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_get";
    i_desc = ImportFunc 0;
  } in
  (import, func_type)

(** Emit `env_at(i)` / `arg_at(i)`: fetch the i-th entry from the WASI
    environ/argv vector and return it as a length-prefixed AffineScript
    string. Sequence:
      1. `*_sizes_get(&count, &bufsize)`
      2. allocate `count*4` bytes for the pointer vector + `bufsize`
         bytes for the string buffer
      3. `*_get(ptrvec, ptrvec + count*4)`
      4. resolve `src = ptrvec[i]`
      5. scan `src` for the null terminator to compute length
      6. allocate `(4 + length)` bytes for the result string,
         store length at +0, byte-copy `src..src+length` to `result+4`
      7. leave the result pointer on the stack

    The byte loops use `I32Load8U`/`I32Store8` (added with the
    byte-level wasm IR extension). The caller has placed the index `i`
    on the stack; this helper consumes it via [LocalSet n_local].

    All locals must be pre-allocated by the caller (8 in total). The
    helper itself does not modify the type or scope context — it only
    emits instructions. *)
let gen_str_at_via_get
    (heap_ptr_global : int)
    (n_local : int)
    (scratch_local : int)
    (count_local : int)
    (bufsize_local : int)
    (ptrvec_local : int)
    (src_local : int)
    (dst_local : int)
    (result_local : int)
    (sizes_func_idx : int)
    (get_func_idx : int)
    : instr list =
  [
    (* Index `i` is on the stack from the caller's arg_code. *)
    LocalSet n_local;

    (* --- Phase 1: sizes_get -> count, bufsize --- *)
    GlobalGet heap_ptr_global;
    I32Const 8l; I32Add;
    GlobalSet heap_ptr_global;
    GlobalGet heap_ptr_global;
    I32Const 8l; I32Sub;
    LocalSet scratch_local;
    LocalGet scratch_local;                       (* count_ptr *)
    LocalGet scratch_local; I32Const 4l; I32Add;  (* bufsize_ptr *)
    Call sizes_func_idx;
    Drop;
    LocalGet scratch_local; I32Load (2, 0); LocalSet count_local;
    LocalGet scratch_local; I32Load (2, 4); LocalSet bufsize_local;

    (* --- Phase 2: allocate ptrvec (count*4) + bytebuf (bufsize) --- *)
    GlobalGet heap_ptr_global;
    LocalSet ptrvec_local;
    GlobalGet heap_ptr_global;
    LocalGet count_local; I32Const 4l; I32Mul;
    LocalGet bufsize_local; I32Add;
    I32Add;
    GlobalSet heap_ptr_global;

    (* --- Phase 3: get(ptrvec, ptrvec + count*4) --- *)
    LocalGet ptrvec_local;
    LocalGet ptrvec_local; LocalGet count_local; I32Const 4l; I32Mul; I32Add;
    Call get_func_idx;
    Drop;

    (* --- Phase 4: src = *(ptrvec + i*4) --- *)
    LocalGet ptrvec_local;
    LocalGet n_local; I32Const 4l; I32Mul; I32Add;
    I32Load (2, 0);
    LocalSet src_local;

    (* --- Phase 5: scan for null terminator. Use scratch as cursor. --- *)
    LocalGet src_local; LocalSet scratch_local;
    Block (BtEmpty, [
      Loop (BtEmpty, [
        LocalGet scratch_local;
        I32Load8U (0, 0);
        I32Eqz; BrIf 1;                           (* exit on 0 byte *)
        LocalGet scratch_local; I32Const 1l; I32Add;
        LocalSet scratch_local;
        Br 0
      ])
    ]);
    (* length = cursor - src (excludes the null terminator).
       Stash it back into count_local, which we are done with. *)
    LocalGet scratch_local; LocalGet src_local; I32Sub;
    LocalSet count_local;

    (* --- Phase 6: allocate (4 + length) for the AS string --- *)
    GlobalGet heap_ptr_global;
    LocalSet result_local;
    GlobalGet heap_ptr_global;
    I32Const 4l; LocalGet count_local; I32Add;
    I32Add;
    GlobalSet heap_ptr_global;

    (* Store length at result+0. *)
    LocalGet result_local;
    LocalGet count_local;
    I32Store (2, 0);

    (* --- Phase 7: byte-copy src..src+length -> result+4 ---
       Reuses scratch as src cursor and count_local as the loop count. *)
    LocalGet src_local; LocalSet scratch_local;
    LocalGet result_local; I32Const 4l; I32Add; LocalSet dst_local;
    Block (BtEmpty, [
      Loop (BtEmpty, [
        LocalGet count_local; I32Eqz; BrIf 1;
        LocalGet dst_local;
        LocalGet scratch_local; I32Load8U (0, 0);
        I32Store8 (0, 0);
        LocalGet scratch_local; I32Const 1l; I32Add; LocalSet scratch_local;
        LocalGet dst_local; I32Const 1l; I32Add; LocalSet dst_local;
        LocalGet count_local; I32Const 1l; I32Sub; LocalSet count_local;
        Br 0
      ])
    ]);

    (* --- Result: leave the string pointer on the stack. --- *)
    LocalGet result_local;
  ]