WASM: fix wasmi validation errors and add end-to-end e2e tests (Phase 4)

ober

afb4cdc9bd22a588893757725f987e7e536e84fe

diff --git a/lib/jerboa/wasm/codegen.sls b/lib/jerboa/wasm/codegen.sls
index 5ca9d7f..0d409f2 100644
--- a/lib/jerboa/wasm/codegen.sls
+++ b/lib/jerboa/wasm/codegen.sls
@@ -545,9 +545,16 @@
       (compile-expr (car args) ctx)
       (bytevector opcode)))
 
-  ;; Memory load: (type.load addr) or (type.load offset align addr)
-  (define (compile-mem-load args ctx opcode)
-    (let ([align 2] [offset 0] [addr-expr (car args)])
+  ;; Memory load: (type.load addr)
+  ;; align is the log2 alignment hint. Must not exceed natural alignment:
+  ;;   byte ops (load8, store8): align=0  (1-byte natural)
+  ;;   16-bit ops (load16, store16): align=1  (2-byte natural)
+  ;;   32-bit ops (load, store i32/f32): align=2  (4-byte natural)
+  ;;   64-bit ops (load, store i64/f64): align=3  (8-byte natural)
+  (define (compile-mem-load args ctx opcode . align-hint)
+    (let ([align (if (pair? align-hint) (car align-hint) 2)]
+          [offset 0]
+          [addr-expr (car args)])
       (bv-concat
         (compile-expr addr-expr ctx)
         (bytevector opcode)
@@ -555,8 +562,9 @@
         (encode-u32-leb128 offset))))
 
   ;; Memory store: (type.store addr val)
-  (define (compile-mem-store args ctx opcode)
-    (let ([align 2] [offset 0])
+  (define (compile-mem-store args ctx opcode . align-hint)
+    (let ([align (if (pair? align-hint) (car align-hint) 2)]
+          [offset 0])
       (bv-concat
         (compile-expr (car args) ctx)
         (compile-expr (cadr args) ctx)
@@ -638,19 +646,31 @@
             (compile-cond args ctx)]
 
            [(when)
-            (bv-concat
-              (compile-expr (car args) ctx)
-              (bytevector wasm-opcode-if wasm-type-void)
-              (compile-body (cdr args) ctx)
-              (bytevector wasm-opcode-end))]
+            ;; when uses a void block — drop any value left by the last body form
+            (let* ([body-exprs (cdr args)]
+                   [body-bv (compile-body body-exprs ctx)]
+                   [last-expr (and (pair? body-exprs) (car (last-pair body-exprs)))]
+                   [need-drop (and last-expr (not (void-expr? last-expr)))])
+              (bv-concat
+                (compile-expr (car args) ctx)
+                (bytevector wasm-opcode-if wasm-type-void)
+                body-bv
+                (if need-drop (bytevector wasm-opcode-drop) (bytevector))
+                (bytevector wasm-opcode-end)))]
 
            [(unless)
-            (bv-concat
-              (compile-expr (car args) ctx)
-              (bytevector wasm-opcode-i32-eqz)
-              (bytevector wasm-opcode-if wasm-type-void)
-              (compile-body (cdr args) ctx)
-              (bytevector wasm-opcode-end))]
+            ;; unless uses a void block — drop any value left by the last body form
+            (let* ([body-exprs (cdr args)]
+                   [body-bv (compile-body body-exprs ctx)]
+                   [last-expr (and (pair? body-exprs) (car (last-pair body-exprs)))]
+                   [need-drop (and last-expr (not (void-expr? last-expr)))])
+              (bv-concat
+                (compile-expr (car args) ctx)
+                (bytevector wasm-opcode-i32-eqz)
+                (bytevector wasm-opcode-if wasm-type-void)
+                body-bv
+                (if need-drop (bytevector wasm-opcode-drop) (bytevector))
+                (bytevector wasm-opcode-end)))]
 
            [(and)
             (if (null? args)
@@ -733,22 +753,27 @@
            [(while)
             ;; (while test body...) →
             ;; (block (loop (br_if (eqz test) 1) body... (br 0)))
+            ;; while uses a void block — drop any value left by the last body form
             (context-push-block! ctx 'block)
             (context-push-block! ctx 'loop)
-            (let ([test-bv (compile-expr (car args) ctx)]
-                  [body-bv (compile-body (cdr args) ctx)])
+            (let* ([body-exprs (cdr args)]
+                   [body-bv (compile-body body-exprs ctx)]
+                   [last-expr (and (pair? body-exprs) (car (last-pair body-exprs)))]
+                   [need-drop (and last-expr (not (void-expr? last-expr)))])
               (context-pop-block! ctx)
               (context-pop-block! ctx)
               (bv-concat
                 (bytevector wasm-opcode-block wasm-type-void)
                 (bytevector wasm-opcode-loop wasm-type-void)
                 ;; Test: if false, break outer block
-                test-bv
+                (compile-expr (car args) ctx)
                 (bytevector wasm-opcode-i32-eqz)
                 (bytevector wasm-opcode-br-if)
                 (encode-u32-leb128 1)  ; br 1 = exit outer block
                 ;; Body
                 body-bv
+                ;; Drop any value left by the body before looping back
+                (if need-drop (bytevector wasm-opcode-drop) (bytevector))
                 ;; Continue loop
                 (bytevector wasm-opcode-br)
                 (encode-u32-leb128 0)  ; br 0 = continue loop
@@ -933,20 +958,20 @@
            [(i64.trunc_f64_u) (compile-unop args ctx wasm-opcode-i64-trunc-f64-u)]
 
            ;; -- Memory operations --
-           [(i32.load)   (compile-mem-load args ctx wasm-opcode-i32-load)]
-           [(i64.load)   (compile-mem-load args ctx wasm-opcode-i64-load)]
-           [(f32.load)   (compile-mem-load args ctx wasm-opcode-f32-load)]
-           [(f64.load)   (compile-mem-load args ctx wasm-opcode-f64-load)]
-           [(i32.load8_s)  (compile-mem-load args ctx wasm-opcode-i32-load8-s)]
-           [(i32.load8_u)  (compile-mem-load args ctx wasm-opcode-i32-load8-u)]
-           [(i32.load16_s) (compile-mem-load args ctx wasm-opcode-i32-load16-s)]
-           [(i32.load16_u) (compile-mem-load args ctx wasm-opcode-i32-load16-u)]
-           [(i32.store)  (compile-mem-store args ctx wasm-opcode-i32-store)]
-           [(i64.store)  (compile-mem-store args ctx wasm-opcode-i64-store)]
-           [(f32.store)  (compile-mem-store args ctx wasm-opcode-f32-store)]
-           [(f64.store)  (compile-mem-store args ctx wasm-opcode-f64-store)]
-           [(i32.store8)  (compile-mem-store args ctx wasm-opcode-i32-store8)]
-           [(i32.store16) (compile-mem-store args ctx wasm-opcode-i32-store16)]
+           [(i32.load)   (compile-mem-load args ctx wasm-opcode-i32-load 2)]
+           [(i64.load)   (compile-mem-load args ctx wasm-opcode-i64-load 3)]
+           [(f32.load)   (compile-mem-load args ctx wasm-opcode-f32-load 2)]
+           [(f64.load)   (compile-mem-load args ctx wasm-opcode-f64-load 3)]
+           [(i32.load8_s)  (compile-mem-load args ctx wasm-opcode-i32-load8-s 0)]
+           [(i32.load8_u)  (compile-mem-load args ctx wasm-opcode-i32-load8-u 0)]
+           [(i32.load16_s) (compile-mem-load args ctx wasm-opcode-i32-load16-s 1)]
+           [(i32.load16_u) (compile-mem-load args ctx wasm-opcode-i32-load16-u 1)]
+           [(i32.store)  (compile-mem-store args ctx wasm-opcode-i32-store 2)]
+           [(i64.store)  (compile-mem-store args ctx wasm-opcode-i64-store 3)]
+           [(f32.store)  (compile-mem-store args ctx wasm-opcode-f32-store 2)]
+           [(f64.store)  (compile-mem-store args ctx wasm-opcode-f64-store 3)]
+           [(i32.store8)  (compile-mem-store args ctx wasm-opcode-i32-store8 0)]
+           [(i32.store16) (compile-mem-store args ctx wasm-opcode-i32-store16 1)]
 
            [(memory.size)
             (bv-concat (bytevector wasm-opcode-memory-size) (bytevector #x00))]
@@ -1354,9 +1379,23 @@
                        [_ (for-each
                             (lambda (p) (context-add-local! ctx (car p) (cdr p)))
                             params)]
-                       ;; Compile body
+                       ;; Compile body.
+                       ;; If the function returns i32 but the body ends with a void
+                       ;; expression (store, set!, while, etc.), wasmi's strict
+                       ;; validator rejects it: "expected i32 but nothing on stack".
+                       ;; Fix: detect this case and push i32.const 0 as a dummy.
                        [body-bv (compile-body body-forms ctx)]
-                       [full-body (bv-concat body-bv (bytevector wasm-opcode-end))]
+                       [needs-dummy-return
+                        (and (= rtype wasm-type-i32)
+                             (pair? body-forms)
+                             (void-expr? (car (last-pair body-forms))))]
+                       [full-body (bv-concat
+                                    body-bv
+                                    (if needs-dummy-return
+                                      (bv-concat (bytevector wasm-opcode-i32-const)
+                                                 (encode-i32-leb128 0))
+                                      (bytevector))
+                                    (bytevector wasm-opcode-end))]
                        ;; Collect extra locals (beyond params)
                        [all-locals (compile-context-locals ctx)]
                        [let-locals
diff --git a/lib/jerboa/wasm/scheme-runtime.sls b/lib/jerboa/wasm/scheme-runtime.sls
index 1104ba7..3c32e5e 100644
--- a/lib/jerboa/wasm/scheme-runtime.sls
+++ b/lib/jerboa/wasm/scheme-runtime.sls
@@ -603,13 +603,9 @@
 
   (define runtime-closure-forms
     '(
-      ;; Read the func-idx stored in a closure header
-      (define (closure-func-idx clos)
-        (i32.load (+ clos 4)))
-
-      ;; Read the env-count stored in a closure header
-      (define (closure-env-count clos)
-        (i32.load (+ clos 8)))
+      ;; Note: closure-func-idx, closure-env-count, closure-env-ref, and
+      ;; closure-env-set! are already defined in value-accessor-forms and
+      ;; value-constructor-forms (from values.sls).  Do NOT redefine them here.
 
       ;; Call a 1-arg closure: (env + 1 user arg) -> result
       ;; Type index 0: (i32 i32) -> i32
@@ -655,6 +651,14 @@
   ;; Combined: all runtime forms
   ;; ================================================================
 
+  ;; runtime-all-forms: base runtime without closure dispatch.
+  ;; Does NOT include runtime-closure-forms (call-closure-N) or
+  ;; runtime-closure-type-forms because those require a function table
+  ;; (define-table) to be present — simple programs without closures
+  ;; would fail wasmi validation with "unknown table 0".
+  ;;
+  ;; slang->wasm-forms adds runtime-closure-type-forms + runtime-closure-forms
+  ;; automatically when the program uses closures.
   (define runtime-all-forms
     (append runtime-list-forms
             runtime-bytevector-forms
@@ -665,7 +669,6 @@
             runtime-equality-forms
             runtime-conversion-forms
             runtime-io-forms
-            runtime-result-forms
-            runtime-closure-forms))
+            runtime-result-forms))
 
 ) ;; end library
diff --git a/lib/std/secure/wasm-target.sls b/lib/std/secure/wasm-target.sls
index 6e17fc2..6b50042 100644
--- a/lib/std/secure/wasm-target.sls
+++ b/lib/std/secure/wasm-target.sls
@@ -983,6 +983,13 @@
         gc-all-forms
         (runtime-forms)
 
+        ;; 4b. Closure dispatch runtime (only when closures are present).
+        ;;     call-closure-N uses call_indirect which requires a table — omitting
+        ;;     these for closure-free programs avoids "unknown table 0" errors.
+        (if has-closures
+          (closure-runtime-forms)
+          '())
+
         ;; 5. Function table (for closures via call_indirect)
         (if has-closures
           '((define-table 64 256))
@@ -1105,6 +1112,20 @@
               #:handle-all)])
       (if (pair? rt) rt '())))
 
+  ;; Load closure runtime dispatch forms (call-closure-1/2/3) from scheme-runtime.
+  ;; Only included when closures are present — these forms use call_indirect which
+  ;; requires a function table.
+  (define (closure-runtime-forms)
+    (let ([rt (with-exception-handler
+                (lambda (e) '())
+                (lambda ()
+                  (eval '(begin
+                           (import (jerboa wasm scheme-runtime))
+                           runtime-closure-forms)
+                        (environment '(chezscheme) '(jerboa wasm scheme-runtime))))
+              #:handle-all)])
+      (if (pair? rt) rt '())))
+
   ;; Check if any form references closures
   (define (has-closures? forms)
     (let ([found #f])
diff --git a/tests/test-slang-wasm.ss b/tests/test-slang-wasm.ss
index 4a3c02f..e56adca 100644
--- a/tests/test-slang-wasm.ss
+++ b/tests/test-slang-wasm.ss
@@ -762,18 +762,25 @@
 
 (section "Closure Infrastructure")
 
-;; runtime-closure-forms contains call-closure-N and closure-func-idx
+;; runtime-closure-forms contains call-closure-N
+;; (closure-func-idx and closure-env-count are in value-accessor-forms, not here)
 (check-pred pair? runtime-closure-forms)
 (let ([names (map (lambda (f)
                     (and (pair? f) (eq? (car f) 'define) (pair? (cadr f))
                          (caadr f)))
                   runtime-closure-forms)])
-  (check-pred pair? (memq 'closure-func-idx names))
-  (check-pred pair? (memq 'closure-env-count names))
   (check-pred pair? (memq 'call-closure-1 names))
   (check-pred pair? (memq 'call-closure-2 names))
   (check-pred pair? (memq 'call-closure-3 names)))
 
+;; closure-func-idx and closure-env-count live in value-accessor-forms
+(let ([names (map (lambda (f)
+                    (and (pair? f) (eq? (car f) 'define) (pair? (cadr f))
+                         (caadr f)))
+                  value-accessor-forms)])
+  (check-pred pair? (memq 'closure-func-idx names))
+  (check-pred pair? (memq 'closure-env-count names)))
+
 ;; runtime-closure-type-forms contains exactly 3 define-type forms
 (check (length runtime-closure-type-forms) => 3)
 (for-each
diff --git a/tests/test-wasm-sandbox.ss b/tests/test-wasm-sandbox.ss
index 5c7b6ca..0fd5aa5 100644
--- a/tests/test-wasm-sandbox.ss
+++ b/tests/test-wasm-sandbox.ss
@@ -7,6 +7,9 @@
 (import (except (chezscheme) compile-program)
         (jerboa wasm format)
         (jerboa wasm codegen)
+        (jerboa wasm values)
+        (jerboa wasm gc)
+        (jerboa wasm scheme-runtime)
         (std wasm sandbox))
 
 (define pass 0)
@@ -244,6 +247,172 @@
   '(111 222))
 
 ;;; ============================================================
+;;; Section 7: End-to-end Scheme runtime in wasmi (Phase 4)
+;;;
+;;; These tests compile the full tagged-value runtime + user code
+;;; and execute it in wasmi.  This is the critical integration proof:
+;;; Scheme semantics (cons/car/cdr, tagged fixnums, etc.) executed
+;;; inside the Rust sandbox.
+;;; ============================================================
+(printf "~%--- Section 7: End-to-end Scheme runtime in wasmi ---~%")
+
+;; Helper: compile the full runtime + user functions into a WASM binary.
+;; User-forms is a list of (define ...) forms using the runtime API.
+(define (compile-scheme-runtime user-forms)
+  (compile-program
+    (append
+      value-memory-forms
+      value-global-forms
+      value-tag-forms
+      value-predicate-forms
+      value-accessor-forms
+      value-constructor-forms
+      gc-all-forms
+      runtime-all-forms
+      user-forms)))
+
+;; Tagged fixnums: tag-fixnum(5) = 11, untag-fixnum(11) = 5
+(test "tag-fixnum round-trip in wasmi"
+  (let* ([bv (compile-scheme-runtime
+               '((define (roundtrip n)
+                   (untag-fixnum (tag-fixnum n)))))]
+         [mod-h (wasm-sandbox-load bv)]
+         [inst (wasm-sandbox-instantiate mod-h)])
+    (let ([r (wasm-sandbox-call inst "roundtrip" 42)])
+      (wasm-sandbox-free inst)
+      (wasm-sandbox-free-module mod-h)
+      r))
+  42)
+
+;; Scheme cons/car/cdr in wasmi
+(test "cons + car in wasmi"
+  (let* ([bv (compile-scheme-runtime
+               '((define (car-of-cons a b)
+                   (scheme-car (scheme-cons a b)))))]
+         [mod-h (wasm-sandbox-load bv)]
+         [inst (wasm-sandbox-instantiate mod-h)])
+    ;; Pass tagged fixnum 7 = (tag-fixnum 3) = 7
+    ;; Pass tagged fixnum 9 = (tag-fixnum 4) = 9
+    ;; car should return 7
+    (let ([r (wasm-sandbox-call inst "car-of-cons" 7 9)])
+      (wasm-sandbox-free inst)
+      (wasm-sandbox-free-module mod-h)
+      r))
+  7)
+
+(test "cons + cdr in wasmi"
+  (let* ([bv (compile-scheme-runtime
+               '((define (cdr-of-cons a b)
+                   (scheme-cdr (scheme-cons a b)))))]
+         [mod-h (wasm-sandbox-load bv)]
+         [inst (wasm-sandbox-instantiate mod-h)])
+    (let ([r (wasm-sandbox-call inst "cdr-of-cons" 7 9)])
+      (wasm-sandbox-free inst)
+      (wasm-sandbox-free-module mod-h)
+      r))
+  9)
+
+;; is-pair? predicate
+(test "is-pair returns true for cons in wasmi"
+  (let* ([bv (compile-scheme-runtime
+               '((define (test-pair a b)
+                   (is-pair (scheme-cons a b)))))]
+         [mod-h (wasm-sandbox-load bv)]
+         [inst (wasm-sandbox-instantiate mod-h)])
+    (let ([r (wasm-sandbox-call inst "test-pair" 3 5)])
+      (wasm-sandbox-free inst)
+      (wasm-sandbox-free-module mod-h)
+      r))
+  1)  ;; true in WASM = 1
+
+;; Scheme list length
+(test "scheme-length of 3-element list in wasmi"
+  (let* ([bv (compile-scheme-runtime
+               '((define (make-list3 a b c)
+                   (scheme-cons a (scheme-cons b (scheme-cons c 4))))
+                 (define (test-length a b c)
+                   (untag-fixnum (scheme-length (make-list3 a b c))))))]
+         [mod-h (wasm-sandbox-load bv)]
+         [inst (wasm-sandbox-instantiate mod-h)])
+    (let ([r (wasm-sandbox-call inst "test-length" 3 5 7)])
+      (wasm-sandbox-free inst)
+      (wasm-sandbox-free-module mod-h)
+      r))
+  3)
+
+;; Closure allocation and func-idx storage + retrieval
+(test "closure func-idx round-trip in wasmi"
+  (let* ([bv (compile-program
+               (append
+                 runtime-closure-type-forms
+                 value-memory-forms
+                 value-global-forms
+                 value-tag-forms
+                 value-predicate-forms
+                 value-accessor-forms
+                 value-constructor-forms
+                 gc-all-forms
+                 runtime-all-forms
+                 '((define-table 64 256)
+                   ;; Lifted closure: env=clos, arg=y. Returns (env[0] + y)
+                   (define (__lifted_adder env y)
+                     (+ (closure-env-ref env 0) y))
+                   ;; Allocate a closure pointing to table slot 0
+                   (define (make-adder x)
+                     (let ([c (alloc-closure 0 1)])
+                       (closure-env-set! c 0 (tag-fixnum x))
+                       c))
+                   ;; Read back the func-idx from the closure header
+                   (define (adder-func-idx x)
+                     (closure-func-idx (make-adder x)))
+                   (define-element 0 (__lifted_adder)))))]
+         [mod-h (wasm-sandbox-load bv)]
+         [inst (wasm-sandbox-instantiate mod-h)])
+    (let ([r (wasm-sandbox-call inst "adder-func-idx" 10)])
+      (wasm-sandbox-free inst)
+      (wasm-sandbox-free-module mod-h)
+      r))
+  0)  ;; table slot 0
+
+;; call-closure-1: dispatch via function table in wasmi
+(test "call-closure-1 dispatches correctly in wasmi"
+  (let* ([bv (compile-program
+               (append
+                 runtime-closure-type-forms
+                 value-memory-forms
+                 value-global-forms
+                 value-tag-forms
+                 value-predicate-forms
+                 value-accessor-forms
+                 value-constructor-forms
+                 gc-all-forms
+                 runtime-all-forms
+                 runtime-closure-forms  ;; call-closure-N (needs table)
+                 '((define-table 64 256)
+                   ;; Lifted closure: adds env[0] to y (both tagged fixnums)
+                   (define (__lifted_adder env y)
+                     ;; env[0] = tagged fixnum, y = tagged fixnum
+                     ;; fx+ untags, adds, retags
+                     (fx+ (closure-env-ref env 0) y))
+                   (define (make-adder x)
+                     (let ([c (alloc-closure 0 1)])
+                       (closure-env-set! c 0 x)
+                       c))
+                   (define (test-call-closure base delta)
+                     (let ([adder (make-adder base)])
+                       (call-closure-1 adder delta)))
+                   (define-element 0 (__lifted_adder)))))]
+         [mod-h (wasm-sandbox-load bv)]
+         [inst (wasm-sandbox-instantiate mod-h)])
+    ;; Pass tagged fixnums: tag-fixnum(10)=21, tag-fixnum(5)=11
+    ;; Result should be tag-fixnum(15)=31
+    (let ([r (wasm-sandbox-call inst "test-call-closure" 21 11)])
+      (wasm-sandbox-free inst)
+      (wasm-sandbox-free-module mod-h)
+      r))
+  31)  ;; tag-fixnum(15) = 31
+
+;;; ============================================================
 ;;; Summary
 ;;; ============================================================