feat: add Plugin.Compiled and call_with_host_context

.github/actions/libextism/action.yml

@@ -12,6 +12,6 @@ runs:
     - uses: ./.extism-cli/.github/actions/extism-cli
     - name: Install
       shell: bash
-      run: sudo extism lib install --version git
+      run: sudo extism lib install --version git --github-token ${{ github.token }}
       env:
         GITHUB_TOKEN: ${{ github.token }}

dune-project

@@ -32,7 +32,7 @@
   (extism-manifest (= :version))
   (ppx_inline_test (>= v0.15.0))
   (cmdliner (>= 1.1.1))
-  (uuidm (>= 0.9.8))
+  (uuidm (>= 0.9.9))
   (mdx (and (>= 2.3.0) :with-test))
  )
  (tags

extism.opam

@@ -20,7 +20,7 @@ depends: [
   "extism-manifest" {= version}
   "ppx_inline_test" {>= "v0.15.0"}
   "cmdliner" {>= "1.1.1"}
-  "uuidm" {>= "0.9.8"}
+  "uuidm" {>= "0.9.9"}
   "mdx" {>= "2.3.0" & with-test}
   "odoc" {with-doc}
 ]

src/bindings.ml

@@ -92,10 +92,19 @@ module Extism_val = struct
 end
 
 let plugin = ptr void
+let compiled_plugin = ptr void
 
 let extism_plugin_new_error_free =
   fn "extism_plugin_new_error_free" (ptr char @-> returning void)
 
+let extism_compiled_plugin_new =
+  fn "extism_compiled_plugin_new"
+    (string @-> uint64_t
+    @-> ptr (ptr void)
+    @-> uint64_t @-> bool
+    @-> ptr (ptr char)
+    @-> returning compiled_plugin)
+
 let extism_plugin_new =
   fn "extism_plugin_new"
     (string @-> uint64_t
@@ -104,6 +113,10 @@ let extism_plugin_new =
     @-> ptr (ptr char)
     @-> returning plugin)
 
+let extism_plugin_new_from_compiled =
+  fn "extism_plugin_new_from_compiled"
+    (compiled_plugin @-> ptr (ptr char) @-> returning plugin)
+
 let extism_plugin_config =
   fn "extism_plugin_config" (plugin @-> string @-> uint64_t @-> returning bool)
 
@@ -115,6 +128,19 @@ let extism_plugin_call_s =
   fn "extism_plugin_call"
     (plugin @-> string @-> string @-> uint64_t @-> returning int32_t)
 
+let extism_current_plugin_host_context =
+  fn "extism_current_plugin_host_context" (ptr void @-> returning (ptr void))
+
+let extism_plugin_call_with_host_contet =
+  fn "extism_plugin_call_with_host_context"
+    (plugin @-> string @-> ptr char @-> uint64_t @-> ptr void
+   @-> returning int32_t)
+
+let extism_plugin_call_s_with_host_context =
+  fn "extism_plugin_call_with_host_context"
+    (plugin @-> string @-> string @-> uint64_t @-> ptr void
+   @-> returning int32_t)
+
 let extism_error = fn "extism_plugin_error" (plugin @-> returning string_opt)
 
 let extism_plugin_output_length =
@@ -136,6 +162,9 @@ let extism_log_drain = fn "extism_log_drain" (log_callback @-> returning void)
 let extism_version = fn "extism_version" (void @-> returning string)
 let extism_plugin_free = fn "extism_plugin_free" (plugin @-> returning void)
 
+let extism_compiled_plugin_free =
+  fn "extism_compiled_plugin_free" (compiled_plugin @-> returning void)
+
 let extism_plugin_function_exists =
   fn "extism_plugin_function_exists" (plugin @-> string @-> returning bool)
 

src/extism.ml

@@ -81,3 +81,27 @@ let%test _ =
     Gc.minor ();
     Gc.full_major ();
     !count > 0
+
+let%test _ =
+  let manifest = Manifest.(create [ Wasm.file "test/code.wasm" ]) in
+  let compiled = Plugin.Compiled.of_manifest_exn manifest in
+  let plugin = Plugin.of_compiled_exn compiled in
+  let _ =
+    Plugin.call Type.string Type.string plugin ~name:"count_vowels"
+      "this is a test"
+  in
+  Gc.minor ();
+  Gc.full_major ();
+  true
+
+let%test _ =
+  let manifest = Manifest.(create [ Wasm.file "test/code.wasm" ]) in
+  let compiled = Plugin.Compiled.of_manifest_exn manifest in
+  let plugin = Plugin.of_compiled_exn compiled in
+  let _ =
+    Plugin.call Type.string Type.string plugin ~name:"count_vowels"
+      "this is a test"
+  in
+  Gc.minor ();
+  Gc.full_major ();
+  true

src/extism.mli

@@ -296,6 +296,9 @@ module Host_function : sig
   val output : (module Type.S with type t = 'a) -> ?index:int -> t -> 'a -> unit
   (** Convert a value, allocate it and update the results array at [index] *)
 
+  val host_context : t -> 'a option
+  (** Get configured host context *)
+
   (** Some helpter functions for reading/writing memory *)
   module Memory_handle : sig
     val memory : ?offs:Unsigned.UInt64.t -> t -> Unsigned.uint8 Ctypes.ptr
@@ -421,6 +424,28 @@ module Plugin : sig
     'b
   (** Similar to {!call} but raises an exception using {!Error.unwrap} *)
 
+
+  val call_with_host_context :
+    (module Type.S with type t = 'a) ->
+    (module Type.S with type t = 'b) ->
+    t ->
+    ctx: 'c ->
+    name:string ->
+    'a ->
+    ('b, Error.t) result
+  (** [call_with_host_context input_type output_type t ~ctx ~name input] executes a function with input
+      and output types defined in {!Type} *)
+
+  val call_with_host_context_exn :
+    (module Type.S with type t = 'a) ->
+    (module Type.S with type t = 'b) ->
+    t ->
+    ctx:'c ->
+    name:string ->
+    'a ->
+    'b
+  (** Similar to {!call_with_host_context} but raises an exception using {!Error.unwrap} *)
+
   val free : t -> unit
   (** Free a plugin immediately, this isn't normally required unless there are a
       lot of plugins open at once *)
@@ -476,6 +501,43 @@ module Plugin : sig
     module Init () : S
     (** Initialize a new typed plugin module *)
   end
+
+  module Compiled : sig
+    type t
+
+    val free : t -> unit
+
+    val create :
+      ?wasi:bool ->
+      ?functions:Function.t list ->
+      string ->
+      (t, Error.t) result
+    (** Make a new compiled plugin from raw WebAssembly or JSON encoded manifest *)
+
+    val create_exn :
+      ?wasi:bool ->
+      ?functions:Function.t list ->
+      string ->
+      t
+    (** Make a new compiled plugin from raw WebAssembly or JSON encoded manifest *)
+
+    val of_manifest :
+      ?wasi:bool ->
+      ?functions:Function.t list ->
+      Manifest.t ->
+      (t, Error.t) result
+    (** Make a new compiled plugin from a {!Manifest} *)
+
+    val of_manifest_exn :
+      ?wasi:bool -> ?functions:Function.t list -> Manifest.t -> t
+    (** Make a new compiled plugin from a {!Manifest} *)
+  end
+
+  val of_compiled: ?config:Manifest.config -> Compiled.t -> (t, Error.t) result
+  (** Make a new plugin from an existing {!Compiled} *)
+
+  val of_compiled_exn: ?config:Manifest.config -> Compiled.t -> t
+  (** Make a new plugin from an existing {!Compiled} *)
 end
 
 val set_log_file :

src/host_function.ml

@@ -103,3 +103,7 @@ let input (type a) (module C : Type.S with type t = a) ?index t =
   C.decode bs
 
 let input_exn a ?index t = input a ?index t |> Error.unwrap
+
+let host_context t =
+  let ptr = Bindings.extism_current_plugin_host_context t.pointer in
+  if Ctypes.is_null ptr then None else Some (Ctypes.Root.get ptr)

src/plugin.ml

@@ -98,6 +98,26 @@ let call' f { pointer; _ } ~name input len =
       in
       Ok buf
 
+let call_with_host_context' f { pointer; _ } ~name input len ctx =
+  if Ctypes.is_null pointer then Error.throw (`Msg "Plugin already freed")
+  else
+    let root = Ctypes.Root.create ctx in
+    let rc = f pointer name input len root in
+    let err = Bindings.extism_error pointer in
+    if rc <> 0l || Option.is_some err then
+      match err with
+      | None -> Error (`Msg "extism_plugin_call failed")
+      | Some msg -> Error (`Msg msg)
+    else
+      let out_len = Bindings.extism_plugin_output_length pointer in
+      let ptr = Bindings.extism_plugin_output_data pointer in
+      let buf =
+        Ctypes.bigarray_of_ptr Ctypes.array1
+          (Unsigned.UInt64.to_int out_len)
+          Char ptr
+      in
+      Ok buf
+
 let call_bigstring (t : t) ~name input =
   let len = Unsigned.UInt64.of_int (Bigstringaf.length input) in
   let ptr = Ctypes.bigarray_start Ctypes.array1 input in
@@ -135,6 +155,23 @@ let call (type a b) (module In : Type.S with type t = a)
 
 let call_exn a b t ~name input = call a b t ~name input |> Error.unwrap
 
+let call_with_host_context (type a b) (module In : Type.S with type t = a)
+    (module Out : Type.S with type t = b) t ~ctx ~name (a : a) :
+    (b, Error.t) result =
+  let input = In.encode a in
+  let len = String.length input in
+  match
+    call_with_host_context' Bindings.extism_plugin_call_s_with_host_context t
+      ~name input
+      (Unsigned.UInt64.of_int len)
+      ctx
+  with
+  | Ok x -> Out.decode x
+  | Error e -> Error e
+
+let call_with_host_context_exn a b t ~ctx ~name input =
+  call_with_host_context a b t ~ctx ~name input |> Error.unwrap
+
 let%test "call" =
   let manifest = Manifest.(create [ Wasm.file "test/code.wasm" ]) in
   let plugin = of_manifest manifest |> Error.unwrap in
@@ -164,6 +201,28 @@ let%test "call_functions" =
   Gc.full_major ();
   b
 
+let%test "call_functions_context" =
+  let open Val_type in
+  let hello_world =
+    Function.create "hello_world" ~params:[ I64 ] ~results:[ I64 ]
+      ~user_data:"Hello again!"
+    @@ fun plugin _user_data ->
+    let ctx = Host_function.host_context plugin in
+    let () = print_endline @@ Option.get ctx in
+    Host_function.output Type.json plugin (`Assoc [ ("count", `Int 999) ])
+  in
+  let functions = [ hello_world ] in
+  let manifest = Manifest.(create [ Wasm.file "test/code-functions.wasm" ]) in
+  let plugin = of_manifest manifest ~functions ~wasi:true |> Error.unwrap in
+  let b =
+    call_with_host_context Type.string Type.string plugin ~name:"count_vowels"
+      "this is a test" ~ctx:"host context"
+    |> Error.unwrap = "{\"count\":999}"
+  in
+  Gc.minor ();
+  Gc.full_major ();
+  b
+
 let function_exists { pointer; _ } name =
   if Ctypes.is_null pointer then Error.throw (`Msg "Plugin already freed")
   else Bindings.extism_plugin_function_exists pointer name
@@ -204,7 +263,7 @@ let id { pointer; _ } =
   else
     let id = Bindings.extism_plugin_id pointer in
     let s = Ctypes.string_from_ptr id ~length:16 in
-    Uuidm.unsafe_of_bytes s
+    Uuidm.unsafe_of_binary_string s
 
 let reset { pointer; _ } = Bindings.extism_plugin_reset pointer
 
@@ -276,3 +335,75 @@ let%test "typed" =
   let n = Yojson.Safe.Util.member "count" res |> Yojson.Safe.Util.to_number in
   Printf.printf "count = %f\n" n;
   n = 3.0
+
+module Compiled = struct
+  type t = {
+    free_lock : Mutex.t;
+    mutable pointer : unit Ctypes.ptr;
+    mutable functions : Function.t list;
+  }
+
+  let free t =
+    let () = Mutex.lock t.free_lock in
+    Fun.protect
+      ~finally:(fun () -> Mutex.unlock t.free_lock)
+      (fun () ->
+        if not (Ctypes.is_null t.pointer) then
+          let () = Bindings.extism_compiled_plugin_free t.pointer in
+          t.pointer <- Ctypes.null)
+
+  let create ?(wasi = false) ?(functions = []) wasm =
+    let func_ptrs = List.map (fun x -> x.Function.pointer) functions in
+    let arr = Ctypes.CArray.of_list Ctypes.(ptr void) func_ptrs in
+    let n_funcs = Ctypes.CArray.length arr in
+    let errmsg =
+      Ctypes.(allocate (ptr char) (coerce (ptr void) (ptr char) null))
+    in
+    let pointer =
+      Bindings.extism_compiled_plugin_new wasm
+        (Unsigned.UInt64.of_int (String.length wasm))
+        (Ctypes.CArray.start arr)
+        (Unsigned.UInt64.of_int n_funcs)
+        wasi errmsg
+    in
+    if Ctypes.is_null pointer then
+      let s = get_errmsg (Ctypes.( !@ ) errmsg) in
+      Error (`Msg s)
+    else
+      let t = { pointer; functions; free_lock = Mutex.create () } in
+      let () = Bindings.set_managed pointer t in
+      let () = Gc.finalise_last (fun () -> free t) t in
+      Ok t
+
+  let create_exn ?wasi ?functions wasm =
+    create ?wasi ?functions wasm |> Error.unwrap
+
+  let of_manifest ?wasi ?functions manifest =
+    let data = Manifest.to_json manifest in
+    create ?wasi ?functions data
+
+  let of_manifest_exn ?wasi ?functions manifest =
+    of_manifest ?wasi ?functions manifest |> Error.unwrap
+end
+
+let of_compiled ?config compiled =
+  let errmsg =
+    Ctypes.(allocate (ptr char) (coerce (ptr void) (ptr char) null))
+  in
+  let pointer =
+    Bindings.extism_plugin_new_from_compiled compiled.Compiled.pointer errmsg
+  in
+  if Ctypes.is_null pointer then
+    let s = get_errmsg (Ctypes.( !@ ) errmsg) in
+    Error (`Msg s)
+  else
+    let t =
+      { pointer; functions = compiled.functions; free_lock = Mutex.create () }
+    in
+    let () = Bindings.set_managed pointer t in
+    let () = Gc.finalise_last (fun () -> free t) t in
+    if not (set_config t config) then Error (`Msg "call to set_config failed")
+    else Ok t
+
+let of_compiled_exn ?config compiled =
+  of_compiled ?config compiled |> Error.unwrap