WASM: implement closure dispatch (call-closure-N, func-idx, element segment)
ober
57c7c4c07f2fd38a76bf261a36c6df78218e7c58
--- a/lib/jerboa/wasm/closure.sls +++ b/lib/jerboa/wasm/closure.sls @@ -471,7 +471,9 @@ to-bool scheme-bool->wasm scheme-ok scheme-err scheme-ok? scheme-err? scheme-unwrap scheme-unwrap-or - scheme-result-value scheme-map-ok)] + scheme-result-value scheme-map-ok + closure-func-idx closure-env-count + call-closure-1 call-closure-2 call-closure-3)] [bound-with-runtime (append runtime-names all-bound)] [free-vars (free-variables `(begin ,@body) bound-with-runtime)] ;; Generate lifted function name @@ -493,12 +495,14 @@ [lifted-def `(define (,lifted-name ,@new-formals) ,@new-body)] ;; The closure allocation expression [n-free (length free-vars)] - ;; Build the closure allocation + env filling + ;; Build the closure allocation + env filling. + ;; Use the lifted function name as a symbolic func-idx placeholder. + ;; wasm-target.sls replaces these symbols with real table indices. [closure-expr (if (= n-free 0) ;; No free variables: still create a closure for uniformity - `(alloc-closure 0 0) ;; func-idx filled in later by wasm-target - `(let ([__clos (alloc-closure 0 ,n-free)]) + `(alloc-closure ,lifted-name 0) + `(let ([__clos (alloc-closure ,lifted-name ,n-free)]) ,@(map (lambda (var) (let ([idx (cdr (assq var env-map))]) `(closure-env-set! __clos ,idx ,var))) --- a/lib/jerboa/wasm/codegen.sls +++ b/lib/jerboa/wasm/codegen.sls @@ -1286,6 +1286,14 @@ [(define-tag) ;; (define-tag type-idx) (wasm-module-add-tag! mod (cadr form))] + [(define-type) + ;; (define-type (param-types) (result-types)) + ;; Pre-registers a type at the next available index. + ;; Used to establish known type indices for call-indirect. + (let* ([ptypes (map scheme->wasm-type (cadr form))] + [rtypes (map scheme->wasm-type (caddr form))] + [type (make-wasm-type ptypes rtypes)]) + (wasm-module-add-type! mod type))] [else (void)]))) forms) --- a/lib/jerboa/wasm/scheme-runtime.sls +++ b/lib/jerboa/wasm/scheme-runtime.sls @@ -30,6 +30,8 @@ runtime-conversion-forms runtime-io-forms runtime-result-forms + runtime-closure-forms + runtime-closure-type-forms runtime-all-forms ) @@ -586,6 +588,70 @@ )) ;; ================================================================ + ;; Closure call dispatch (call_indirect via function table) + ;; + ;; These type indices depend on define-type forms being emitted + ;; BEFORE imports in the forms list. slang->wasm-forms emits: + ;; (define-type (i32 i32) (i32)) → type-idx 0 (arity-1 closure) + ;; (define-type (i32 i32 i32) (i32)) → type-idx 1 (arity-2 closure) + ;; (define-type (i32 i32 i32 i32) (i32)) → type-idx 2 (arity-3 closure) + ;; Then imports follow (indices 3+), then user functions. + ;; + ;; call-indirect form: (call-indirect type-idx env arg... table-index-expr) + ;; The last arg is the table index (func-idx from the closure header). + ;; ================================================================ + + (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))) + + ;; Call a 1-arg closure: (env + 1 user arg) -> result + ;; Type index 0: (i32 i32) -> i32 + (define (call-closure-1 clos arg) + (call-indirect 0 + clos ;; env (i32) + arg ;; user arg (i32) + (closure-func-idx clos))) ;; table index + + ;; Call a 2-arg closure: (env + 2 user args) -> result + ;; Type index 1: (i32 i32 i32) -> i32 + (define (call-closure-2 clos arg1 arg2) + (call-indirect 1 + clos ;; env (i32) + arg1 ;; user arg 1 (i32) + arg2 ;; user arg 2 (i32) + (closure-func-idx clos))) ;; table index + + ;; Call a 3-arg closure: (env + 3 user args) -> result + ;; Type index 2: (i32 i32 i32 i32) -> i32 + (define (call-closure-3 clos arg1 arg2 arg3) + (call-indirect 2 + clos ;; env (i32) + arg1 ;; user arg 1 (i32) + arg2 ;; user arg 2 (i32) + arg3 ;; user arg 3 (i32) + (closure-func-idx clos))) ;; table index + )) + + ;; Pre-type registration forms — must be emitted FIRST in the program, + ;; before any define-import forms, to guarantee stable type indices 0-2. + (define runtime-closure-type-forms + '( + ;; Type 0: arity-1 closure (env i32, arg i32) -> i32 + (define-type (i32 i32) (i32)) + ;; Type 1: arity-2 closure (env i32, arg1 i32, arg2 i32) -> i32 + (define-type (i32 i32 i32) (i32)) + ;; Type 2: arity-3 closure (env i32, arg1 i32, arg2 i32, arg3 i32) -> i32 + (define-type (i32 i32 i32 i32) (i32)) + )) + + ;; ================================================================ ;; Combined: all runtime forms ;; ================================================================ @@ -599,6 +665,7 @@ runtime-equality-forms runtime-conversion-forms runtime-io-forms - runtime-result-forms)) + runtime-result-forms + runtime-closure-forms)) ) ;; end library --- a/lib/std/secure/wasm-target.sls +++ b/lib/std/secure/wasm-target.sls @@ -950,10 +950,21 @@ [static-strings (collect-static-strings lifted)] [string-data-forms (generate-string-data static-strings)] ;; Replace (string-from-static #vu8(...)) with (string-from-static offset) - [lifted (replace-static-strings lifted static-strings)]) + [lifted (replace-static-strings lifted static-strings)] + ;; Assign real table indices to alloc-closure placeholders + [closure-assignment (assign-closure-indices lifted)] + [lifted (car closure-assignment)] + [element-forms (cdr closure-assignment)] + [has-closures (has-closures? lifted)]) ;; Assemble the complete program (append + ;; 0. Closure type pre-registration — MUST come before imports so + ;; type indices 0/1/2 are stable for call-indirect in call-closure-N. + (if has-closures + (closure-type-forms) + '()) + ;; 1. Memory and globals value-memory-forms value-global-forms @@ -973,12 +984,15 @@ (runtime-forms) ;; 5. Function table (for closures via call_indirect) - (if (has-closures? lifted) + (if has-closures '((define-table 64 256)) '()) - ;; 6. User program (lifted + lowered) - lifted))) + ;; 6. User program (lifted + lowered, with real func-idx in closures) + lifted + + ;; 7. Element segment: populate function table with lifted functions + element-forms))) ;; Collect all static string references and assign offsets (define (collect-static-strings forms) @@ -1035,6 +1049,62 @@ [else expr])) (map replace forms)) + ;; Collect lifted function names (symbols starting with "__lifted_") in order. + (define (collect-lifted-names forms) + (let loop ([fs forms] [names '()]) + (if (null? fs) + (reverse names) + (let ([f (car fs)]) + (if (and (pair? f) + (eq? (car f) 'define) + (pair? (cadr f)) + (let ([name (caadr f)]) + (and (symbol? name) + (let ([s (symbol->string name)]) + (and (> (string-length s) 9) + (string=? (substring s 0 9) "__lifted_")))))) + (loop (cdr fs) (cons (caadr f) names)) + (loop (cdr fs) names)))))) + + ;; Replace (alloc-closure sym n) with (alloc-closure table-idx n) + ;; using a mapping of sym → integer table index. + ;; Returns (cons rewritten-forms element-forms). + (define (assign-closure-indices forms) + (let* ([lifted-names (collect-lifted-names forms)] + [name->idx (let loop ([ns lifted-names] [i 0]) + (if (null? ns) '() + (cons (cons (car ns) i) + (loop (cdr ns) (+ i 1)))))]) + (define (rewrite expr) + (cond + [(pair? expr) + (if (and (eq? (car expr) 'alloc-closure) + (symbol? (cadr expr))) + ;; Replace symbolic name with table index + (let ([entry (assq (cadr expr) name->idx)]) + (if entry + `(alloc-closure ,(cdr entry) ,(caddr expr)) + expr)) ;; unknown name — leave as-is + (map rewrite expr))] + [else expr])) + (let ([rewritten (map rewrite forms)] + [element-forms (if (null? lifted-names) + '() + (list `(define-element 0 ,lifted-names)))]) + (cons rewritten element-forms)))) + + ;; Load closure type pre-registration forms from scheme-runtime module. + (define (closure-type-forms) + (let ([rt (with-exception-handler + (lambda (e) '()) + (lambda () + (eval '(begin + (import (jerboa wasm scheme-runtime)) + runtime-closure-type-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]) --- a/tests/test-slang-wasm.ss +++ b/tests/test-slang-wasm.ss @@ -757,6 +757,102 @@ (check (> (bytevector-length wasm) 100) => #t)) ;; ================================================================ +;; Closure Infrastructure (Phase 1A/1B/1C) +;; ================================================================ + +(section "Closure Infrastructure") + +;; runtime-closure-forms contains call-closure-N and closure-func-idx +(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))) + +;; runtime-closure-type-forms contains exactly 3 define-type forms +(check (length runtime-closure-type-forms) => 3) +(for-each + (lambda (form) + (check (car form) => 'define-type)) + runtime-closure-type-forms) + +;; define-type is accepted by compile-program without error +(let ([wasm (compile-program + (append + value-memory-forms + value-global-forms + '((define-type (i32 i32) (i32)) + (define-type (i32 i32 i32) (i32))) + value-tag-forms + value-predicate-forms + value-accessor-forms + value-constructor-forms + gc-all-forms + runtime-all-forms + '((define (test-fn x) x))))]) + (check-pred bytevector? wasm) + (check (> (bytevector-length wasm) 50) => #t)) + +;; lambda-lift now uses symbolic lifted name in alloc-closure (not 0) +(let* ([forms '((define (make-adder x) + (lambda (y) (+ x y))))] + [lifted (lambda-lift forms)]) + ;; Find alloc-closure forms in the lifted output + (define (find-alloc-closure forms) + (let loop ([fs forms] [found '()]) + (if (null? fs) found + (let ([f (car fs)]) + (loop (cdr fs) + (if (and (pair? f) (eq? (car f) 'alloc-closure)) + (cons f found) + (append found (find-alloc-closure + (if (pair? f) f '()))))))))) + (define (collect-alloc-closures expr) + (cond + [(not (pair? expr)) '()] + [(eq? (car expr) 'alloc-closure) (list expr)] + [else (apply append (map collect-alloc-closures expr))])) + (let ([allocs (apply append (map collect-alloc-closures lifted))]) + ;; There should be at least one alloc-closure + (check-pred pair? allocs) + ;; The func-idx field (cadr) should be a symbol (lifted name), not 0 + (for-each + (lambda (ac) + (check-pred symbol? (cadr ac))) + allocs))) + +;; assign-closure-indices replaces symbolic func-idx with integer +;; (test via compile-program: a closure-using program compiles cleanly) +(let ([wasm (compile-program + (append + ;; closure type pre-registration first + 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 + ;; A pre-lifted closure: func-idx 0 (first table slot) + '((define-table 64 256) + (define (__lifted_test env y) + (+ (closure-env-ref env 0) y)) + (define (make-adder x) + (let ([c (alloc-closure 0 1)]) + (closure-env-set! c 0 x) + c)) + (define-element 0 (__lifted_test)))))]) + (check-pred bytevector? wasm) + (check (> (bytevector-length wasm) 100) => #t)) + +;; ================================================================ ;; Summary ;; ================================================================