Wasm plugin system (#1555)

Cargo.lock

@@ -494,6 +494,12 @@ dependencies = [
  "syn 2.0.16",
 ]
 
+[[package]]
+name = "downcast-rs"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ea835d29036a4087793836fa931b08837ad5e957da9e23886b29586fb9b6650"
+
 [[package]]
 name = "ecow"
 version = "0.1.1"
@@ -1014,6 +1020,12 @@ dependencies = [
  "hashbrown 0.14.0",
 ]
 
+[[package]]
+name = "indexmap-nostd"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e04e2fd2b8188ea827b32ef11de88377086d690286ab35747ef7f9bf3ccb590"
+
 [[package]]
 name = "inferno"
 version = "0.11.15"
@@ -1066,6 +1078,12 @@ dependencies = [
  "cfg-if",
 ]
 
+[[package]]
+name = "intx"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f6f38a50a899dc47a6d0ed5508e7f601a2e34c3a85303514b5d137f3c10a0c75"
+
 [[package]]
 name = "io-lifetimes"
 version = "1.0.10"
@@ -2590,6 +2608,7 @@ dependencies = [
  "unicode-segmentation",
  "unscanny",
  "usvg",
+ "wasmi",
  "xmlparser",
  "xmlwriter",
  "xmp-writer",
@@ -3037,6 +3056,47 @@ version = "0.2.86"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ed9d5b4305409d1fc9482fee2d7f9bcbf24b3972bf59817ef757e23982242a93"
 
+[[package]]
+name = "wasmi"
+version = "0.30.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e51fb5c61993e71158abf5bb863df2674ca3ec39ed6471c64f07aeaf751d67b4"
+dependencies = [
+ "intx",
+ "smallvec",
+ "spin 0.9.8",
+ "wasmi_arena",
+ "wasmi_core",
+ "wasmparser-nostd",
+]
+
+[[package]]
+name = "wasmi_arena"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "401c1f35e413fac1846d4843745589d9ec678977ab35a384db8ae7830525d468"
+
+[[package]]
+name = "wasmi_core"
+version = "0.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "624e6333e861ef49095d2d678b76ebf30b06bf37effca845be7e5b87c90071b7"
+dependencies = [
+ "downcast-rs",
+ "libm",
+ "num-traits",
+ "paste",
+]
+
+[[package]]
+name = "wasmparser-nostd"
+version = "0.100.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9157cab83003221bfd385833ab587a039f5d6fa7304854042ba358a3b09e0724"
+dependencies = [
+ "indexmap-nostd",
+]
+
 [[package]]
 name = "web-sys"
 version = "0.3.63"

crates/typst-library/src/compute/construct.rs

@@ -3,7 +3,7 @@ use std::str::FromStr;
 
 use time::{Month, PrimitiveDateTime};
 
-use typst::eval::{Bytes, Datetime, Module, Reflect, Regex};
+use typst::eval::{Bytes, Datetime, Module, Plugin, Reflect, Regex};
 
 use crate::prelude::*;
 
@@ -831,6 +831,115 @@ pub fn range(
     Ok(array)
 }
 
+/// Loads a WebAssembly plugin.
+///
+/// This is **advanced functionality** and not be confused with
+/// [Typst packages]($scripting/#packages).
+///
+/// Typst is capable of interfacing with plugins compiled to WebAssembly. Plugin
+/// functions may accept multiple [byte buffers]($type/bytes) as arguments and
+/// return a single byte buffer. They should typically be wrapped in idiomatic
+/// Typst functions that perform the necessary conversions between native Typst
+/// types and bytes.
+///
+/// Plugins run in isolation from your system, which means that printing,
+/// reading files, or anything like that will not be supported for security
+/// reasons. To run as a plugin, a program needs to be compiled to a 32-bit
+/// shared WebAssembly library. Many compilers will use the
+/// [WASI ABI](https://wasi.dev/) by default or as their only option (e.g.
+/// emscripten), which allows printing, reading files, etc. This ABI will not
+/// directly work with Typst. You will either need to compile to a different
+/// target or [stub all functions](https://github.com/astrale-sharp/wasm-minimal-protocol/blob/master/wasi-stub).
+///
+/// ## Example { #example }
+/// ```example
+/// #let myplugin = plugin("hello.wasm")
+/// #let concat(a, b) = str(
+///   myplugin.concatenate(
+///     bytes(a),
+///     bytes(b),
+///   )
+/// )
+///
+/// #concat("hello", "world")
+/// ```
+///
+/// ## Protocol { #protocol }
+/// To be used as a plugin, a WebAssembly module must conform to the following
+/// protocol
+///
+/// ### Exports { #exports }
+/// A plugin module can export functions to make them callable from Typst. To
+/// conform to the protocol, an exported function should:
+///
+/// - Take `n` 32-bit integer arguments `a_1`, `a_2`, ..., `a_n` (interpreted as
+///   lengths, so `usize/size_t` may be preferable), and return one 32-bit
+///   integer.
+///
+/// - The function should first allocate a buffer `buf` of length
+///   `a_1 + a_2 + ... + a_n`, and then call
+///   `wasm_minimal_protocol_write_args_to_buffer(buf.ptr)`.
+///
+/// - The `a_1` first bytes of the buffer now constitute the first argument, the
+///   `a_2` next bytes the second argument, and so on.
+///
+/// - The function can now do its job with the argument and produce an output
+///   buffer. Before returning, it should call
+///   `wasm_minimal_protocol_send_result_to_host` to send its result back to the
+///   host.
+///
+/// - To signal success, the function should return `0`.
+///
+/// - To signal an error, the function should return `1`. The written buffer is
+///   then interpreted as an UTF-8 encoded error message.
+///
+/// ### Imports { #imports }
+/// Plugin modules need to import two functions that are provided by the runtime.
+/// (Types and functions are described using WAT syntax.)
+///
+/// - `(import "typst_env" "wasm_minimal_protocol_write_args_to_buffer" (func (param i32)))`
+///
+///   Writes the arguments for the current function into a plugin-allocated
+///   buffer. When a plugin function is called, it
+///   [receives the lengths](#exported-functions) of its input buffers as
+///   arguments. It should then allocate a buffer whose capacity is at least the
+///   sum of these lengths. It should then call this function with a `ptr` to
+///   the buffer to fill it with the arguments, one after another.
+///
+/// - `(import "typst_env" "wasm_minimal_protocol_send_result_to_host" (func (param i32 i32)))`
+///
+///   Sends the output of the current function to the host (Typst). The first
+///   parameter shall be a pointer to a buffer (`ptr`), while the second is the
+///   length of that buffer (`len`). The memory pointed at by `ptr` can be freed
+///   immediately after this function returns. If the message should be
+///   interpreted as an error message, it should be encoded as UTF-8.
+///
+/// ## Resources { #resources }
+/// For more resources, check out the
+/// [wasm-minimal-protocol repository](https://github.com/astrale-sharp/wasm-minimal-protocol).
+/// It contains:
+///
+/// - A list of example plugin implementations and a test runner for these
+///   examples
+/// - Wrappers to help you write your plugin in Rust (Zig wrapper in
+///   development)
+/// - A stubber for WASI
+///
+/// Display: Plugin
+/// Category: construct
+#[func]
+pub fn plugin(
+    /// Path to a WebAssembly file.
+    path: Spanned<EcoString>,
+    /// The virtual machine.
+    vm: &mut Vm,
+) -> SourceResult<Plugin> {
+    let Spanned { v: path, span } = path;
+    let id = vm.location().join(&path).at(span)?;
+    let data = vm.world().file(id).at(span)?;
+    Plugin::new(data).at(span)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/typst-library/src/compute/mod.rs

@@ -39,4 +39,5 @@ pub(super) fn define(global: &mut Scope) {
     global.define("yaml", yaml_func());
     global.define("xml", xml_func());
     global.define("calc", calc::module());
+    global.define("plugin", plugin_func());
 }

crates/typst/Cargo.toml

@@ -55,6 +55,7 @@ usvg = { version = "0.32", default-features = false, features = ["text"] }
 xmlwriter = "0.1.0"
 xmp-writer = "0.1"
 time = { version = "0.3.20", features = ["std", "formatting"] }
+wasmi = "0.30.0"
 xmlparser = "0.13.5"
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]

crates/typst/src/eval/methods.rs

@@ -2,7 +2,7 @@
 
 use ecow::{eco_format, EcoString};
 
-use super::{Args, IntoValue, Str, Value, Vm};
+use super::{Args, Bytes, IntoValue, Plugin, Str, Value, Vm};
 use crate::diag::{At, Hint, SourceResult};
 use crate::eval::{bail, Datetime};
 use crate::geom::{Align, Axes, Color, Dir, Em, GenAlign};
@@ -280,6 +280,14 @@ pub fn call(
                     "inv" => align2d.map(GenAlign::inv).into_value(),
                     _ => return missing(),
                 }
+            } else if let Some(plugin) = dynamic.downcast::<Plugin>() {
+                if plugin.iter().any(|func_name| func_name == method) {
+                    let bytes = args.all::<Bytes>()?;
+                    args.take().finish()?;
+                    plugin.call(method, bytes).at(span)?.into_value()
+                } else {
+                    return missing();
+                }
             } else {
                 return (vm.items.library_method)(vm, &dynamic, method, args, span);
             }

crates/typst/src/eval/mod.rs

@@ -23,6 +23,7 @@ mod methods;
 mod module;
 mod none;
 pub mod ops;
+mod plugin;
 mod scope;
 mod symbol;
 mod tracer;
@@ -53,6 +54,7 @@ pub use self::library::{set_lang_items, LangItems, Library};
 pub use self::methods::methods_on;
 pub use self::module::Module;
 pub use self::none::NoneValue;
+pub use self::plugin::Plugin;
 pub use self::scope::{Scope, Scopes};
 pub use self::str::{format_str, Regex, Str};
 pub use self::symbol::Symbol;

crates/typst/src/eval/plugin.rs

@@ -0,0 +1,208 @@
+use std::fmt::{self, Debug, Formatter};
+use std::hash::{Hash, Hasher};
+
+use ecow::{eco_format, EcoString};
+use std::sync::{Arc, Mutex};
+use wasmi::{AsContext, AsContextMut, Caller, Engine, Linker, Module};
+
+use super::{cast, Bytes};
+use crate::diag::{bail, StrResult};
+
+/// A plugin loaded from WebAssembly code.
+///
+/// It can run external code conforming to its protocol.
+///
+/// This type is cheap to clone and hash.
+#[derive(Clone)]
+pub struct Plugin(Arc<Repr>);
+
+/// The internal representation of a plugin.
+struct Repr {
+    /// The raw WebAssembly bytes.
+    bytes: Bytes,
+    /// The function defined by the WebAssembly module.
+    functions: Vec<(EcoString, wasmi::Func)>,
+    /// Owns all data associated with the WebAssembly module.
+    store: Mutex<Store>,
+}
+
+/// Owns all data associated with the WebAssembly module.
+type Store = wasmi::Store<StoreData>;
+
+/// The persistent store data used for communication between store and host.
+#[derive(Default)]
+struct StoreData {
+    args: Vec<Bytes>,
+    output: Vec<u8>,
+}
+
+impl Plugin {
+    /// Create a new plugin from raw WebAssembly bytes.
+    #[comemo::memoize]
+    pub fn new(bytes: Bytes) -> StrResult<Self> {
+        let engine = Engine::default();
+        let module = Module::new(&engine, bytes.as_slice())
+            .map_err(|err| format!("failed to load WebAssembly module: {err}"))?;
+
+        let mut linker = Linker::new(&engine);
+        linker
+            .func_wrap(
+                "typst_env",
+                "wasm_minimal_protocol_send_result_to_host",
+                wasm_minimal_protocol_send_result_to_host,
+            )
+            .unwrap();
+        linker
+            .func_wrap(
+                "typst_env",
+                "wasm_minimal_protocol_write_args_to_buffer",
+                wasm_minimal_protocol_write_args_to_buffer,
+            )
+            .unwrap();
+
+        let mut store = Store::new(&engine, StoreData::default());
+        let instance = linker
+            .instantiate(&mut store, &module)
+            .and_then(|pre_instance| pre_instance.start(&mut store))
+            .map_err(|e| eco_format!("{e}"))?;
+
+        // Ensure that the plugin exports its memory.
+        if !matches!(
+            instance.get_export(&store, "memory"),
+            Some(wasmi::Extern::Memory(_))
+        ) {
+            bail!("plugin does not export its memory");
+        }
+
+        // Collect exported functions.
+        let functions = instance
+            .exports(&store)
+            .filter_map(|export| {
+                let name = export.name().into();
+                export.into_func().map(|func| (name, func))
+            })
+            .collect();
+
+        Ok(Plugin(Arc::new(Repr { bytes, functions, store: Mutex::new(store) })))
+    }
+
+    /// Call the plugin function with the given `name`.
+    pub fn call(&self, name: &str, args: Vec<Bytes>) -> StrResult<Bytes> {
+        // Find the function with the given name.
+        let func = self
+            .0
+            .functions
+            .iter()
+            .find(|(v, _)| v == name)
+            .map(|&(_, func)| func)
+            .ok_or_else(|| {
+                eco_format!("plugin does not contain a function called {name}")
+            })?;
+
+        let mut store = self.0.store.lock().unwrap();
+        let ty = func.ty(store.as_context());
+
+        // Check function signature.
+        if ty.params().iter().any(|&v| v != wasmi::core::ValueType::I32) {
+            bail!(
+                "plugin function `{name}` has a parameter that is not a 32-bit integer"
+            );
+        }
+        if ty.results() != [wasmi::core::ValueType::I32] {
+            bail!("plugin function `{name}` does not return exactly one 32-bit integer");
+        }
+
+        // Check inputs.
+        let expected = ty.params().len();
+        let given = args.len();
+        if expected != given {
+            bail!(
+                "plugin function takes {expected} argument{}, but {given} {} given",
+                if expected == 1 { "" } else { "s" },
+                if given == 1 { "was" } else { "were" },
+            );
+        }
+
+        // Collect the lengths of the argument buffers.
+        let lengths = args
+            .iter()
+            .map(|a| wasmi::Value::I32(a.len() as i32))
+            .collect::<Vec<_>>();
+
+        // Store the input data.
+        store.data_mut().args = args;
+
+        // Call the function.
+        let mut code = wasmi::Value::I32(-1);
+        func.call(store.as_context_mut(), &lengths, std::slice::from_mut(&mut code))
+            .map_err(|err| eco_format!("plugin panicked: {err}"))?;
+
+        // Extract the returned data.
+        let output = std::mem::take(&mut store.data_mut().output);
+
+        // Parse the functions return value.
+        match code {
+            wasmi::Value::I32(0) => {}
+            wasmi::Value::I32(1) => match std::str::from_utf8(&output) {
+                Ok(message) => bail!("plugin errored with: {message}"),
+                Err(_) => {
+                    bail!("plugin errored, but did not return a valid error message")
+                }
+            },
+            _ => bail!("plugin did not respect the protocol"),
+        };
+
+        Ok(output.into())
+    }
+
+    /// An iterator over all the function names defined by the plugin.
+    pub fn iter(&self) -> impl Iterator<Item = &EcoString> {
+        self.0.functions.as_slice().iter().map(|(func_name, _)| func_name)
+    }
+}
+
+impl Debug for Plugin {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        f.pad("plugin(..)")
+    }
+}
+
+impl PartialEq for Plugin {
+    fn eq(&self, other: &Self) -> bool {
+        self.0.bytes == other.0.bytes
+    }
+}
+
+impl Hash for Plugin {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.0.bytes.hash(state);
+    }
+}
+
+cast! {
+    type Plugin: "plugin",
+}
+
+/// Write the arguments to the plugin function into the plugin's memory.
+fn wasm_minimal_protocol_write_args_to_buffer(mut caller: Caller<StoreData>, ptr: u32) {
+    let memory = caller.get_export("memory").unwrap().into_memory().unwrap();
+    let arguments = std::mem::take(&mut caller.data_mut().args);
+    let mut offset = ptr as usize;
+    for arg in arguments {
+        memory.write(&mut caller, offset, arg.as_slice()).unwrap();
+        offset += arg.len();
+    }
+}
+
+/// Extracts the output of the plugin function from the plugin's memory.
+fn wasm_minimal_protocol_send_result_to_host(
+    mut caller: Caller<StoreData>,
+    ptr: u32,
+    len: u32,
+) {
+    let memory = caller.get_export("memory").unwrap().into_memory().unwrap();
+    let mut buffer = std::mem::take(&mut caller.data_mut().output);
+    buffer.resize(len as usize, 0);
+    memory.read(&caller, ptr as _, &mut buffer).unwrap();
+    caller.data_mut().output = buffer;
+}

crates/typst/src/ide/complete.rs

@@ -10,7 +10,7 @@ use super::analyze::analyze_labels;
 use super::{analyze_expr, analyze_import, plain_docs_sentence, summarize_font_family};
 use crate::doc::Frame;
 use crate::eval::{
-    fields_on, format_str, methods_on, CastInfo, Func, Library, Scope, Value,
+    fields_on, format_str, methods_on, CastInfo, Func, Library, Plugin, Scope, Value,
 };
 use crate::syntax::{
     ast, is_id_continue, is_id_start, is_ident, LinkedNode, Source, SyntaxKind,
@@ -415,6 +415,18 @@ fn field_access_completions(ctx: &mut CompletionContext, value: &Value) {
                 }
             }
         }
+        Value::Dyn(val) => {
+            if let Some(plugin) = val.downcast::<Plugin>() {
+                for name in plugin.iter() {
+                    ctx.completions.push(Completion {
+                        kind: CompletionKind::Func,
+                        label: name.clone(),
+                        apply: None,
+                        detail: None,
+                    })
+                }
+            }
+        }
         _ => {}
     }
 }

tests/typ/compiler/plugin.typ

@@ -0,0 +1,35 @@
+// Test WebAssembly plugins.
+// Ref: false
+
+---
+#let p = plugin("/files/hello.wasm")
+#test(p.hello(), bytes("Hello from wasm!!!"))
+#test(p.double_it(bytes("hey!")), bytes("hey!.hey!"))
+#test(
+  p.shuffle(bytes("value1"), bytes("value2"), bytes("value3")),
+  bytes("value3-value1-value2"),
+)
+
+---
+#let p = plugin("/files/hello.wasm")
+
+// Error: 2-20 plugin function takes 0 arguments, but 1 was given
+#p.hello(bytes(""))
+
+---
+#let p = plugin("/files/hello.wasm")
+
+// Error: 10-14 unexpected argument
+#p.hello(true)
+
+---
+#let p = plugin("/files/hello.wasm")
+
+// Error: 2-17 plugin errored with: This is an `Err`
+#p.returns_err()
+
+---
+#let p = plugin("/files/hello.wasm")
+
+// Error: 2-16 plugin panicked: wasm `unreachable` instruction executed
+#p.will_panic()