add ability to dynamically link wasm modules (#596)

* add wasm-link test modules and setup

* (wip) pass modules to link from elixir to rust

* add ability to dynamically link wasm modules

* add support to dynamically link compiled module

* update linking api to use map-of-map instead of list-of-map

* update param name from links to linked_module to avoid confusion with linker

* update dynamic-linking docs with additional notes

* update api to prevent over exposure of rust resources

* rename remaining links to linked_modules

* add unit tests for dynamic module linking

* add dynamic module linking to the changelog

* update error messages returned from rust for clarity

* refactor dynamic module linking code to improve readability and maintainability (#1)

* update dynamic module linking docs

Co-authored-by: Philipp Tessenow <[email protected]>

* add missing new line at the end of the file

Co-authored-by: Philipp Tessenow <[email protected]>

* add support for dynamically linked module dependencies (#2)

* fix nested dynamic link issue

* add test for compiled dynamically linked module deps

---------

Co-authored-by: Sonny Scroggin <[email protected]>
Co-authored-by: Philipp Tessenow <[email protected]>

.gitignore

@@ -25,6 +25,8 @@ wasmex-*.tar
 # Cargo things in the Rust part of this package
 priv/native/libwasmex.so
 test/wasm_source/target/*
+test/wasm_link_test/target/*
+test/wasm_link_dep_test/target/*
 
 .mix_tasks
 **/.DS_Store

CHANGELOG.md

@@ -28,6 +28,7 @@ please check your fuel consumption values.
 - official support for Elixir 1.15 and 1.16
 - fuel-related API got rewritten, because the underlying Wasm library (wasmtime) changed their API and we want to be consistent. Added `Store.get_fuel/1` and `Store.set_fuel/2` which is a much simpler API than before.
 - read and write a global’s value with `Instance.get_global_value/3` and `Instance.set_global_value/4` ([#540](https://github.com/tessi/wasmex/pull/540))
+- ability to dynamically link wasm modules ([#596](https://github.com/tessi/wasmex/pull/596))
 
 ### Removed
 

lib/wasmex.ex

@@ -92,6 +92,55 @@ defmodule Wasmex do
   - `:f32` a 32 bit float
   - `:f64` a 64 bit float
 
+  ### Linking multiple Wasm modules
+
+  Wasm module `links` may be given as an additional option.
+  Links is a map of module names to Wasm modules.
+
+      iex> calculator_wasm = File.read!(TestHelper.wasm_link_test_file_path())
+      iex> utils_wasm = File.read!(TestHelper.wasm_test_file_path())
+      iex> links = %{utils: %{bytes: utils_wasm}}
+      iex> {:ok, pid} = Wasmex.start_link(%{bytes: calculator_wasm, links: links})
+      iex> Wasmex.call_function(pid, "sum_range", [1, 5])
+      {:ok, [15]}
+
+  It is also possible to link an already compiled module.
+  This improves performance if the same module is used many times by compiling it only once.
+
+      iex> calculator_wasm = File.read!(TestHelper.wasm_link_test_file_path())
+      iex> utils_wasm = File.read!(TestHelper.wasm_test_file_path())
+      iex> {:ok, store} = Wasmex.Store.new()
+      iex> {:ok, utils_module} = Wasmex.Module.compile(store, utils_wasm)
+      iex> links = %{utils: %{module: utils_module}}
+      iex> {:ok, pid} = Wasmex.start_link(%{bytes: calculator_wasm, links: links, store: store})
+      iex> Wasmex.call_function(pid, "sum_range", [1, 5])
+      {:ok, [15]}
+
+  **Important:** Make sure to use the same store for the linked modules and the main module.
+
+  When linking multiple Wasm modules, it is important to handle their dependencies properly.
+  This can be achieved by providing a map of module names to their respective Wasm modules in the `links` option.
+
+  For example, if we have a main module that depends on a calculator module, and the calculator module depends on a utils module, we can link them as follows:
+
+      iex> main_wasm = File.read!(TestHelper.wasm_link_dep_test_file_path())
+      iex> calculator_wasm = File.read!(TestHelper.wasm_link_test_file_path())
+      iex> utils_wasm = File.read!(TestHelper.wasm_test_file_path())
+      iex> links = %{
+      ...>   calculator: %{
+      ...>     bytes: calculator_wasm,
+      ...>     links: %{
+      ...>       utils: %{bytes: utils_wasm}
+      ...>     }
+      ...>   }
+      ...> }
+      iex> {:ok, pid} = Wasmex.start_link(%{bytes: main_wasm, links: links})
+
+  In this example, the `links` map specifies that the `calculator` module depends on the `utils` module.
+  The `links` map is a nested map, where each module name is associated with a map that contains the Wasm module bytes and its dependencies.
+
+  The `links` map can also be used to link an already compiled module, as shown in the previous examples.
+
   ### WASI
 
   Optionally, modules can be run with WebAssembly System Interface (WASI) support.
@@ -159,6 +208,9 @@ defmodule Wasmex do
   def start_link(%{} = opts) when not is_map_key(opts, :imports),
     do: start_link(Map.merge(opts, %{imports: %{}}))
 
+  def start_link(%{} = opts) when not is_map_key(opts, :links),
+    do: start_link(Map.merge(opts, %{links: %{}}))
+
   def start_link(%{} = opts) when is_map_key(opts, :module) and not is_map_key(opts, :store),
     do: {:error, :must_specify_store_used_to_compile_module}
 
@@ -182,15 +234,41 @@ defmodule Wasmex do
     end
   end
 
-  def start_link(%{store: store, module: module, imports: imports} = opts)
-      when is_map(imports) and not is_map_key(opts, :bytes) do
+  def start_link(%{links: links, store: store} = opts)
+      when is_map(links) and not is_map_key(opts, :compiled_links) do
+    compiled_links =
+      links
+      |> flatten_links()
+      |> Enum.reverse()
+      |> Enum.uniq_by(&elem(&1, 0))
+      |> Enum.map(&build_compiled_links(&1, store))
+
+    opts
+    |> Map.delete(:links)
+    |> Map.put(:compiled_links, compiled_links)
+    |> start_link()
+  end
+
+  def start_link(%{store: store, module: module, imports: imports, compiled_links: links} = opts)
+      when is_map(imports) and is_list(links) and not is_map_key(opts, :bytes) do
     GenServer.start_link(__MODULE__, %{
       store: store,
       module: module,
+      links: links,
       imports: stringify_keys(imports)
     })
   end
 
+  defp flatten_links(links) do
+    Enum.flat_map(links, fn {name, opts} ->
+      if Map.has_key?(opts, :links) do
+        [{name, Map.drop(opts, [:links])} | flatten_links(opts.links)]
+      else
+        [{name, opts}]
+      end
+    end)
+  end
+
   defp build_store(opts) do
     if Map.has_key?(opts, :wasi) do
       Wasmex.Store.new_wasi(stringify_keys(opts[:wasi]))
@@ -199,6 +277,17 @@ defmodule Wasmex do
     end
   end
 
+  defp build_compiled_links({name, %{bytes: bytes} = opts}, store)
+       when not is_map_key(opts, :module) do
+    with {:ok, module} <- Wasmex.Module.compile(store, bytes) do
+      %{name: stringify(name), module: module}
+    end
+  end
+
+  defp build_compiled_links({name, %{module: module}}, _store) do
+    %{name: stringify(name), module: module}
+  end
+
   @doc ~S"""
   Returns whether a function export with the given `name` exists in the Wasm instance.
 
@@ -353,6 +442,10 @@ defmodule Wasmex do
     for {key, val} <- map, into: %{}, do: {stringify(key), stringify_keys(val)}
   end
 
+  defp stringify_keys(list) when is_list(list) do
+    for val <- list, into: [], do: stringify_keys(val)
+  end
+
   defp stringify_keys(value), do: value
 
   defp stringify(s) when is_binary(s), do: s
@@ -361,8 +454,9 @@ defmodule Wasmex do
   # Server
 
   @impl true
-  def init(%{store: store, module: module, imports: imports} = state) when is_map(imports) do
-    case Wasmex.Instance.new(store, module, imports) do
+  def init(%{store: store, module: module, imports: imports, links: links} = state)
+      when is_map(imports) and is_list(links) do
+    case Wasmex.Instance.new(store, module, imports, links) do
       {:ok, instance} -> {:ok, Map.merge(state, %{instance: instance})}
       {:error, reason} -> {:error, reason}
     end

lib/wasmex/instance.ex

@@ -35,6 +35,8 @@ defmodule Wasmex.Instance do
   The `import` parameter is a nested map of Wasm namespaces.
   Each namespace consists of a name and a map of function names to function signatures.
 
+  The `links` parameter is a list of name-module pairs that are dynamically linked to the instance.
+
   Function signatures are a tuple of the form `{:fn, arg_types, return_types, callback}`.
   Where `arg_types` and `return_types` are lists of `:i32`, `:i64`, `:f32`, `:f64`.
 
@@ -61,17 +63,28 @@ defmodule Wasmex.Instance do
       ...>       "imported_void" => {:fn, [], [], fn _context -> nil end}
       ...>     }
       ...> }
-      iex> {:ok, %Wasmex.Instance{}} = Wasmex.Instance.new(store, module, imports)
+      ...> links = []
+      iex> {:ok, %Wasmex.Instance{}} = Wasmex.Instance.new(store, module, imports, links)
   """
-  @spec new(Wasmex.StoreOrCaller.t(), Wasmex.Module.t(), %{
-          optional(binary()) => (... -> any())
-        }) ::
+  @spec new(
+          Wasmex.StoreOrCaller.t(),
+          Wasmex.Module.t(),
+          %{optional(binary()) => (... -> any())},
+          [%{optional(binary()) => Wasmex.Module.t()}] | []
+        ) ::
           {:ok, __MODULE__.t()} | {:error, binary()}
-  def new(store_or_caller, module, imports) when is_map(imports) do
+  def new(store_or_caller, module, imports, links \\ [])
+      when is_map(imports) and is_list(links) do
     %Wasmex.StoreOrCaller{resource: store_or_caller_resource} = store_or_caller
     %Wasmex.Module{resource: module_resource} = module
 
-    case Wasmex.Native.instance_new(store_or_caller_resource, module_resource, imports) do
+    links =
+      links
+      |> Enum.map(fn %{name: name, module: module} ->
+        %{name: name, module_resource: module.resource}
+      end)
+
+    case Wasmex.Native.instance_new(store_or_caller_resource, module_resource, imports, links) do
       {:error, err} -> {:error, err}
       resource -> {:ok, __wrap_resource__(resource)}
     end

lib/wasmex/native.ex

@@ -32,7 +32,7 @@ defmodule Wasmex.Native do
   def module_serialize(_module_resource), do: error()
   def module_unsafe_deserialize(_binary, _engine_resource), do: error()
 
-  def instance_new(_store_or_caller_resource, _module_resource, _imports), do: error()
+  def instance_new(_store_or_caller_resource, _module_resource, _imports, _links), do: error()
 
   def instance_function_export_exists(
         _store_or_caller_resource,

native/wasmex/src/environment.rs

@@ -8,9 +8,9 @@ use wasmtime::{Caller, FuncType, Linker, Val, ValType};
 use wiggle::anyhow::{self, anyhow};
 
 use crate::{
-    atoms::{self},
+    atoms,
     caller::{remove_caller, set_caller},
-    instance::{map_wasm_values_to_vals, WasmValue},
+    instance::{map_wasm_values_to_vals, LinkedModule, WasmValue},
     memory::MemoryResource,
     store::{StoreData, StoreOrCaller, StoreOrCallerResource},
 };
@@ -25,6 +25,33 @@ pub struct CallbackToken {
     pub return_values: Mutex<Option<(bool, Vec<WasmValue>)>>,
 }
 
+pub fn link_modules(
+    linker: &mut Linker<StoreData>,
+    store: &mut StoreOrCaller,
+    linked_modules: Vec<LinkedModule>,
+) -> Result<(), Error> {
+    for linked_module in linked_modules {
+        let module_name = linked_module.name;
+        let module = linked_module.module_resource.inner.lock().map_err(|e| {
+            rustler::Error::Term(Box::new(format!(
+                "Could not unlock linked module resource as the mutex was poisoned: {e}"
+            )))
+        })?;
+
+        let instance = linker.instantiate(&mut *store, &module).map_err(|e| {
+            rustler::Error::Term(Box::new(format!(
+                "Could not instantiate linked module: {e}"
+            )))
+        })?;
+
+        linker
+            .instance(&mut *store, &module_name, instance)
+            .map_err(|err| Error::Term(Box::new(err.to_string())))?;
+    }
+
+    Ok(())
+}
+
 pub fn link_imports(linker: &mut Linker<StoreData>, imports: MapIterator) -> Result<(), Error> {
     for (namespace_name, namespace_definition) in imports {
         let namespace_name = namespace_name.decode::<String>()?;

native/wasmex/src/instance.rs

@@ -2,9 +2,8 @@ use rustler::{
     dynamic::TermType,
     env::{OwnedEnv, SavedTerm},
     resource::ResourceArc,
-    types::tuple::make_tuple,
-    types::ListIterator,
-    Encoder, Env as RustlerEnv, Error, MapIterator, NifResult, Term,
+    types::{tuple::make_tuple, ListIterator},
+    Encoder, Env as RustlerEnv, Error, MapIterator, NifMap, NifResult, Term,
 };
 use std::ops::Deref;
 use std::sync::Mutex;
@@ -14,13 +13,19 @@ use wasmtime::{Instance, Linker, Module, Val, ValType};
 
 use crate::{
     atoms,
-    environment::{link_imports, CallbackTokenResource},
+    environment::{link_imports, link_modules, CallbackTokenResource},
     functions,
     module::ModuleResource,
     printable_term_type::PrintableTermType,
     store::{StoreData, StoreOrCaller, StoreOrCallerResource},
 };
 
+#[derive(NifMap)]
+pub struct LinkedModule {
+    pub name: String,
+    pub module_resource: ResourceArc<ModuleResource>,
+}
+
 pub struct InstanceResource {
     pub inner: Mutex<Instance>,
 }
@@ -37,6 +42,7 @@ pub fn new(
     store_or_caller_resource: ResourceArc<StoreOrCallerResource>,
     module_resource: ResourceArc<ModuleResource>,
     imports: MapIterator,
+    linked_modules: Vec<LinkedModule>,
 ) -> Result<ResourceArc<InstanceResource>, rustler::Error> {
     let module = module_resource.inner.lock().map_err(|e| {
         rustler::Error::Term(Box::new(format!(
@@ -50,7 +56,7 @@ pub fn new(
             )))
         })?);
 
-    let instance = link_and_create_instance(store_or_caller, &module, imports)?;
+    let instance = link_and_create_instance(store_or_caller, &module, imports, linked_modules)?;
     let resource = ResourceArc::new(InstanceResource {
         inner: Mutex::new(instance),
     });
@@ -61,13 +67,15 @@ fn link_and_create_instance(
     store_or_caller: &mut StoreOrCaller,
     module: &Module,
     imports: MapIterator,
+    linked_modules: Vec<LinkedModule>,
 ) -> Result<Instance, Error> {
     let mut linker = Linker::new(store_or_caller.engine());
     if let Some(_wasi_ctx) = &store_or_caller.data().wasi {
         linker.allow_shadowing(true);
         wasmtime_wasi::add_to_linker(&mut linker, |s: &mut StoreData| s.wasi.as_mut().unwrap())
             .map_err(|err| Error::Term(Box::new(err.to_string())))?;
     }
+    link_modules(&mut linker, store_or_caller, linked_modules)?;
     link_imports(&mut linker, imports)?;
     linker
         .instantiate(store_or_caller, module)