Racket Wasmtime FFI Example

readme.md

This program shows how to interact with WASM using FFI and a shared buffer between the WASM module and Racket.

The WASM function "run" calls the Fibonacci function assuming the input is stored in the first 8 bytes and returns the 64 bit unsigned result to that same 8 bytes. You could set things up to call $fibo directly, but it is somewhat tedious since you then need to worry about defining the ABI of said functions much more carefully.

On another note, this example uses current-cleanup-context to manage cleanup of resources with an arena style cleanup mechanism. Instead of worrying about cleaning up each individual object, we mostly reason about things on broader scales and try to run multiple cleanup steps all at once.

example-wasmtime.rkt

#lang racket
(require ffi/unsafe
         ffi/unsafe/define)
(require syntax/parse/define)
(require pipechain)
(require (for-syntax pipechain))
(require (for-syntax racket/match racket/format))

;; Download wasmtime C lib and set this file to the DLL/library file
(define wasmtime-dll "")

(define current-cleanup-context (make-parameter (box void)))

(define (make-once-thunk thunk)
  (let ([already-called? (box #f)])
    (lambda args
      (unless (unbox already-called?)
        (set-box! already-called? #t) ; Do this immediately in case the thunk does some continuation stuff
        (apply thunk args)))))

(define (add-cleanup-step thunk)
  (let* ([cleanup-context-box (current-cleanup-context)]
         [old-thunk (unbox cleanup-context-box)])
    (set-box!
     cleanup-context-box
     (make-once-thunk
      (lambda ()
        (dynamic-wind
         void
         thunk
         old-thunk))))))

(define (new-cleanup-context thunk)
  (let* ([new-context (box void)]
         [run-cleanup (make-once-thunk (lambda () ((unbox new-context))))])
    (parameterize ([current-cleanup-context new-context])
      (dynamic-wind void
                    thunk
                    run-cleanup))))

(define-syntax-parse-rule (attach-cleanup val-expr:expr cleanup:expr)
  (~> val-expr (~>effect (add-cleanup-step (lambda () cleanup)))))

(define-syntax-parse-rule (attach-ptr-cleanup val-expr:expr cleanup:expr)
  (~> val-expr (~>effect (add-cleanup-step (lambda () (unless (equal? _ #f) cleanup))))))

(define-ffi-definer define-wasmtime (ffi-lib wasmtime-dll))

(define-syntax (define-wasm-type stx)
  (let* ([name-stx (~> stx syntax->list cdr car)]
         [name-sym (syntax->datum name-stx)])
    (match (~a name-sym)
      [(regexp #rx"^_(.*)-pointer$" (list _ cname))
       (define cname-sym (string->symbol cname))
       #`(define #,name-stx (_cpointer/null '#,cname-sym))])))
  
#;(define _wasm_engine_t-pointer (_cpointer 'wasm_engine_t))
(define-wasm-type _wasm_engine_t-pointer)
(define-wasm-type _wasmtime_store_t-pointer)
(define-wasm-type _wasmtime_context_t-pointer)
(define-wasm-type _wasmtime_error_t-pointer)
(define-wasm-type _wasm_valtype_t-pointer)
(define-cstruct _wasm_byte_vec_t ([size _size]
                                  [data _pointer]))
(define-wasm-type _wasmtime_module_t-pointer)
(define-wasm-type _wasm_functype_t-pointer)
(define-cstruct _wasmtime_func_t ([store_id _uint64]
                                  [index _size]))
(define-cstruct _wasmtime_global_t ([store_id _uint64]
                                    [index _size]))
(define-cstruct _wasmtime_table_t ([store_id _uint64]
                                    [index _size]))
(define-cstruct _wasmtime_memory_t ([store_id _uint64]
                                    [index _size]))
(define-wasm-type _wasmtime_caller-pointer)
(define-wasm-type _wasm_trap_t-pointer)

(define _wasmtime_extern_union_t
  (_union _wasmtime_func_t
          _wasmtime_global_t
          _wasmtime_table_t
          _wasmtime_memory_t))
(define-cstruct _wasmtime_extern_t ([kind _uint8]
                                    [of _wasmtime_extern_union_t]))

(define-wasm-type _wasmtime_externref_t-pointer)
(define _wasmtime_valunion_t
  (_union _int32
          _int64
          _float
          _double
          _wasmtime_func_t
          _wasmtime_externref_t-pointer
          (_array _uint8 16))) ;; wasmtime_v128
(define-cstruct _wasmtime_val_t ([kind _uint8]
                                 [of _wasmtime_valunion_t]))

(define-cstruct _wasmtime_instance_t ([store_id _uint64]
                                      [index _size]))

(define-cstruct _wasm_valtype_vec_t ([size _size]
                                     [data _pointer]))

(define-wasmtime wasm_trap_delete (_fun _wasm_trap_t-pointer -> _void))
(define (add-trap-cleanup trap)
  (attach-ptr-cleanup trap (wasm_trap_delete _)))
(define-wasmtime wasm_engine_delete (_fun _wasm_engine_t-pointer -> _void))
(define-wasmtime wasm_engine_new (_fun -> (res : _wasm_engine_t-pointer) -> (attach-ptr-cleanup res (wasm_engine_delete _))))
(define-wasmtime wasmtime_store_delete (_fun _wasmtime_store_t-pointer -> _void))
(define-wasmtime wasmtime_store_context (_fun _wasmtime_store_t-pointer -> _wasmtime_context_t-pointer))
(define-wasmtime wasm_byte_vec_new_uninitialized (_fun (output-bytes : (_ptr o _wasm_byte_vec_t)) _size -> _void -> output-bytes))
(define-wasmtime wasm_byte_vec_delete (_fun (_ptr i _wasm_byte_vec_t) -> _void))
(define-wasmtime wasmtime_error_delete (_fun _wasmtime_error_t-pointer -> _void))
(define-wasmtime wasmtime_error_message (_fun _wasmtime_error_t-pointer (output-bytes : (_ptr o _wasm_byte_vec_t))
                                              -> _void
                                              -> (let ([delete-msg (make-once-thunk (lambda () (wasm_byte_vec_delete output-bytes)))])
                                                   (dynamic-wind
                                                    void
                                                    (lambda () (~> output-bytes wasm_byte_vec_t-data (cast _ _pointer _string)))
                                                    delete-msg))))

(define-syntax-parse-rule (check-wasm-err err-sym err res)
  (let ([err-val err])
    (if err-val ;; need to make sure err-val always gets cleaned up quicly.
        (let ([err-msg (dynamic-wind void (lambda () (wasmtime_error_message err-val)) (lambda () (wasmtime_error_delete err-val)))])
          (error err-sym "wasm failure: ~a" err-msg))
        res)))

(define-wasmtime wasmtime_wat2wasm (_fun _pointer _size (output-bytes : (_ptr o _wasm_byte_vec_t))
                                         -> (err : _wasmtime_error_t-pointer)
                                         -> (check-wasm-err 'wasmtime_wat2wasm err
                                                            (attach-cleanup output-bytes wasm_byte_vec_delete))))

(define-wasmtime wasmtime_module_delete (_fun _wasmtime_module_t-pointer -> _void))
(define-wasmtime wasmtime_module_new (_fun _wasm_engine_t-pointer _pointer _size (ret : (_ptr o _wasmtime_module_t-pointer))
                                           -> (err : _wasmtime_error_t-pointer)
                                           -> (check-wasm-err 'wasmtime_module_new err
                                                              (attach-ptr-cleanup ret (wasmtime_module_delete _)))))
(define-wasmtime wasm_valtype_vec_new_empty (_fun (res : (_ptr o _wasm_valtype_vec_t)) -> _void -> res))
(define-wasmtime wasm_functype_delete (_fun _wasm_functype_t-pointer -> _void))
(define-wasmtime wasm_functype_new (_fun (_ptr i _wasm_valtype_vec_t) (_ptr i _wasm_valtype_vec_t)
                                         -> (res : _wasm_functype_t-pointer)
                                         -> (attach-ptr-cleanup res (wasm_functype_delete _))))

(define _wasmtime_finalizer_func_t (_fun _pointer -> _void))
(define _wasmtime_func_callback_t-pointer (_fun _pointer _wasmtime_caller-pointer _wasmtime_val_t-pointer _size _wasmtime_val_t-pointer _size -> _wasm_trap_t-pointer))

(define-wasmtime wasmtime_func_new (_fun _wasmtime_context_t-pointer _wasm_functype_t-pointer _wasmtime_func_callback_t-pointer _pointer _wasmtime_finalizer_func_t
                                         (ret : (_ptr o _wasmtime_func_t)) -> _void -> ret))

(define-wasmtime wasmtime_store_new (_fun _wasm_engine_t-pointer _pointer _wasmtime_finalizer_func_t
                                          -> (res : _wasmtime_store_t-pointer)
                                          -> (attach-cleanup res (when _ (wasmtime_store_delete _)))))

(define-wasmtime wasmtime_instance_new (_fun _wasmtime_context_t-pointer _wasmtime_module_t-pointer
                                             [externs : (_list i _wasmtime_extern_t)] [_size = (length externs)]
                                             [inst : (_ptr o _wasmtime_instance_t)]
                                             [trap : (_ptr io _wasm_trap_t-pointer)] ;; Very tedious, but it needs to be IO since wasmtime doesn't set this on success.
                                             -> (err : _wasmtime_error_t-pointer)
                                             -> (let ()
                                                  (add-trap-cleanup trap)
                                                  (check-wasm-err 'wasmtime_instance_new err #f)
                                                  (values inst trap)))) ;; instances have no destructor

(define-wasmtime wasmtime_instance_export_get (_fun _wasmtime_context_t-pointer (_ptr i _wasmtime_instance_t) _pointer _size (res : (_ptr o _wasmtime_extern_t)) -> (has-item : _stdbool) -> (values has-item res)))
(define-wasmtime wasmtime_func_call (_fun _wasmtime_context_t-pointer (_ptr i _wasmtime_func_t)
                                          [args : (_list i _wasmtime_val_t)] [_size = (length args)]
                                          (_or-null _pointer) _size
                                          [trap : (_ptr o _wasm_trap_t-pointer)]
                                          -> (err : _wasmtime_error_t-pointer)
                                          -> (let ()
                                               (add-trap-cleanup trap)
                                               (check-wasm-err 'wasmtime_func_call err
                                                               trap))))

(define-wasmtime wasm_trap_message (_fun _wasm_trap_t-pointer (res : (_ptr o _wasm_byte_vec_t)) -> _void
                                                                   -> (let ()
                                                                        (dynamic-wind
                                                                         void
                                                                         (lambda ()
                                                                           (let ([bytes-ptr (wasm_byte_vec_t-data res)]
                                                                                 [bytes-size (wasm_byte_vec_t-size res)])
                                                                             (~> (cblock->list bytes-ptr _byte bytes-size) list->bytes bytes->string/utf-8)))
                                                                         (lambda () (wasm_byte_vec_delete res))))))

(define-wasmtime wasmtime_memory_data (_fun _wasmtime_context_t-pointer (_ptr i _wasmtime_memory_t) -> _pointer))

(define my-module
  '(module
       (func $hello (import "" "hello"))
     (memory (export "memory") 1 1)
     (func $fibo (param $n i64) (result i64)
           (block $end
                  local.get $n
                  i64.const 2
                  i64.le_u
                  (if
                   (then
                    i64.const 1
                    local.set $n
                    br $end))
                  local.get $n
                  i64.const 1
                  i64.sub
                  call $fibo
                  local.get $n
                  i64.const 2
                  i64.sub
                  call $fibo
                  i64.add
                  local.set $n)
           local.get $n)
     (func (export "run")
           (local $result i64)
           i32.const 0
           i64.load
           call $fibo
           local.set $result
           i32.const 0
           local.get $result
           i64.store)
     ))

(define called-callback #f)

(define (hello-callback env caller args nargs results nresults)
  (set! called-callback #t)
  (thread (lambda () ;; Need another thread to do IO since we can't do anything that blocks here.
            (displayln "Calling back...")
            (displayln "> Hello World!")))
  #f)

(define WASMTIME_EXTERN_FUNC 0)
(define WASMTIME_EXTERN_GLOBAL 1)
(define WASMTIME_EXTERN_TABLE 2)
(define WASMTIME_EXTERN_MEMORY 3)

(new-cleanup-context
 (lambda ()
   (let* ([engine (wasm_engine_new)]
          [store (wasmtime_store_new engine #f #f)]
          [store-context (wasmtime_store_context store)]
          [my-module-bytes (~> my-module ~s string->bytes/utf-8)]
          [wasminput (wasmtime_wat2wasm my-module-bytes (bytes-length my-module-bytes))]
          [wasm-module (wasmtime_module_new engine (wasm_byte_vec_t-data wasminput) (wasm_byte_vec_t-size wasminput))]
          [param-types (wasm_valtype_vec_new_empty)] ;; ownership of these is taken by wasm_functype_new.
          [result-types (wasm_valtype_vec_new_empty)]
          [callback-thunk-type (wasm_functype_new param-types result-types)]
          [host-func (wasmtime_func_new store-context callback-thunk-type hello-callback #f #f)]
          [func-name #"run"]
          [memory-name #"memory"]
          [wasm-imp (~> (cast host-func _wasmtime_func_t _wasmtime_extern_union_t) (make-wasmtime_extern_t WASMTIME_EXTERN_FUNC _))])
     (define-values (inst trap) (wasmtime_instance_new store-context wasm-module (list wasm-imp) #f)) ;; TODO: actually handle traps
     #;(displayln (wasm_trap_message trap))
     (define-values (has-export exp) (wasmtime_instance_export_get store-context inst func-name (bytes-length func-name)))
     (define-values (has-mem memory-export) (wasmtime_instance_export_get store-context inst memory-name (bytes-length memory-name)))
     (define memory-export-ptr (wasmtime_memory_data store-context
                                                     (~> memory-export wasmtime_extern_t-of (cast _ _wasmtime_valunion_t _wasmtime_memory_t))))
     (define n 10)
     (ptr-set! memory-export-ptr _uint64 n)
     (wasmtime_func_call store-context (~> exp wasmtime_extern_t-of (cast _ _wasmtime_valunion_t _wasmtime_func_t)) '() #f 0)
     (printf "fibo(~a) = ~a~n" n (ptr-ref memory-export-ptr _uint64)))))

license.md

Copyright 2023 Dyllon Gagnier

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.