Modular, multi-threaded, transitioning plugins (#5779)

7 files changed, 468 insertions, 184 deletions

diff --git a/crates/typst-library/src/foundations/func.rs b/crates/typst-library/src/foundations/func.rs
index cb3eba16..a05deb1f 100644
--- a/crates/typst-library/src/foundations/func.rs
+++ b/crates/typst-library/src/foundations/func.rs
@@ -9,11 +9,11 @@ use ecow::{eco_format, EcoString};
 use typst_syntax::{ast, Span, SyntaxNode};
 use typst_utils::{singleton, LazyHash, Static};
 
-use crate::diag::{bail, SourceResult, StrResult};
+use crate::diag::{bail, At, SourceResult, StrResult};
 use crate::engine::Engine;
 use crate::foundations::{
-    cast, repr, scope, ty, Args, CastInfo, Content, Context, Element, IntoArgs, Scope,
-    Selector, Type, Value,
+    cast, repr, scope, ty, Args, Bytes, CastInfo, Content, Context, Element, IntoArgs,
+    PluginFunc, Scope, Selector, Type, Value,
 };
 
 /// A mapping from argument values to a return value.
@@ -151,6 +151,8 @@ enum Repr {
     Element(Element),
     /// A user-defined closure.
     Closure(Arc<LazyHash<Closure>>),
+    /// A plugin WebAssembly function.
+    Plugin(Arc<PluginFunc>),
     /// A nested function with pre-applied arguments.
     With(Arc<(Func, Args)>),
 }
@@ -164,6 +166,7 @@ impl Func {
             Repr::Native(native) => Some(native.name),
             Repr::Element(elem) => Some(elem.name()),
             Repr::Closure(closure) => closure.name(),
+            Repr::Plugin(func) => Some(func.name()),
             Repr::With(with) => with.0.name(),
         }
     }
@@ -176,6 +179,7 @@ impl Func {
             Repr::Native(native) => Some(native.title),
             Repr::Element(elem) => Some(elem.title()),
             Repr::Closure(_) => None,
+            Repr::Plugin(_) => None,
             Repr::With(with) => with.0.title(),
         }
     }
@@ -186,6 +190,7 @@ impl Func {
             Repr::Native(native) => Some(native.docs),
             Repr::Element(elem) => Some(elem.docs()),
             Repr::Closure(_) => None,
+            Repr::Plugin(_) => None,
             Repr::With(with) => with.0.docs(),
         }
     }
@@ -204,6 +209,7 @@ impl Func {
             Repr::Native(native) => Some(&native.0.params),
             Repr::Element(elem) => Some(elem.params()),
             Repr::Closure(_) => None,
+            Repr::Plugin(_) => None,
             Repr::With(with) => with.0.params(),
         }
     }
@@ -221,6 +227,7 @@ impl Func {
                 Some(singleton!(CastInfo, CastInfo::Type(Type::of::<Content>())))
             }
             Repr::Closure(_) => None,
+            Repr::Plugin(_) => None,
             Repr::With(with) => with.0.returns(),
         }
     }
@@ -231,6 +238,7 @@ impl Func {
             Repr::Native(native) => native.keywords,
             Repr::Element(elem) => elem.keywords(),
             Repr::Closure(_) => &[],
+            Repr::Plugin(_) => &[],
             Repr::With(with) => with.0.keywords(),
         }
     }
@@ -241,6 +249,7 @@ impl Func {
             Repr::Native(native) => Some(&native.0.scope),
             Repr::Element(elem) => Some(elem.scope()),
             Repr::Closure(_) => None,
+            Repr::Plugin(_) => None,
             Repr::With(with) => with.0.scope(),
         }
     }
@@ -266,6 +275,14 @@ impl Func {
         }
     }
 
+    /// Extract the plugin function, if it is one.
+    pub fn to_plugin(&self) -> Option<&PluginFunc> {
+        match &self.repr {
+            Repr::Plugin(func) => Some(func),
+            _ => None,
+        }
+    }
+
     /// Call the function with the given context and arguments.
     pub fn call<A: IntoArgs>(
         &self,
@@ -307,6 +324,12 @@ impl Func {
                 context,
                 args,
             ),
+            Repr::Plugin(func) => {
+                let inputs = args.all::<Bytes>()?;
+                let output = func.call(inputs).at(args.span)?;
+                args.finish()?;
+                Ok(Value::Bytes(output))
+            }
             Repr::With(with) => {
                 args.items = with.1.items.iter().cloned().chain(args.items).collect();
                 with.0.call(engine, context, args)
@@ -425,12 +448,30 @@ impl From<Repr> for Func {
     }
 }
 
+impl From<&'static NativeFuncData> for Func {
+    fn from(data: &'static NativeFuncData) -> Self {
+        Repr::Native(Static(data)).into()
+    }
+}
+
 impl From<Element> for Func {
     fn from(func: Element) -> Self {
         Repr::Element(func).into()
     }
 }
 
+impl From<Closure> for Func {
+    fn from(closure: Closure) -> Self {
+        Repr::Closure(Arc::new(LazyHash::new(closure))).into()
+    }
+}
+
+impl From<PluginFunc> for Func {
+    fn from(func: PluginFunc) -> Self {
+        Repr::Plugin(Arc::new(func)).into()
+    }
+}
+
 /// A Typst function that is defined by a native Rust type that shadows a
 /// native Rust function.
 pub trait NativeFunc {
@@ -466,12 +507,6 @@ pub struct NativeFuncData {
     pub returns: LazyLock<CastInfo>,
 }
 
-impl From<&'static NativeFuncData> for Func {
-    fn from(data: &'static NativeFuncData) -> Self {
-        Repr::Native(Static(data)).into()
-    }
-}
-
 cast! {
     &'static NativeFuncData,
     self => Func::from(self).into_value(),
@@ -525,12 +560,6 @@ impl Closure {
     }
 }
 
-impl From<Closure> for Func {
-    fn from(closure: Closure) -> Self {
-        Repr::Closure(Arc::new(LazyHash::new(closure))).into()
-    }
-}
-
 cast! {
     Closure,
     self => Value::Func(self.into()),
diff --git a/crates/typst-library/src/foundations/mod.rs b/crates/typst-library/src/foundations/mod.rs
index 2921481b..a790da4f 100644
--- a/crates/typst-library/src/foundations/mod.rs
+++ b/crates/typst-library/src/foundations/mod.rs
@@ -25,7 +25,8 @@ mod int;
 mod label;
 mod module;
 mod none;
-mod plugin;
+#[path = "plugin.rs"]
+mod plugin_;
 mod scope;
 mod selector;
 mod str;
@@ -56,7 +57,7 @@ pub use self::int::*;
 pub use self::label::*;
 pub use self::module::*;
 pub use self::none::*;
-pub use self::plugin::*;
+pub use self::plugin_::*;
 pub use self::repr::Repr;
 pub use self::scope::*;
 pub use self::selector::*;
@@ -114,11 +115,11 @@ pub(super) fn define(global: &mut Scope, inputs: Dict, features: &Features) {
     global.define_type::<Symbol>();
     global.define_type::<Duration>();
     global.define_type::<Version>();
-    global.define_type::<Plugin>();
     global.define_func::<repr::repr>();
     global.define_func::<panic>();
     global.define_func::<assert>();
     global.define_func::<eval>();
+    global.define_func::<plugin>();
     if features.is_enabled(Feature::Html) {
         global.define_func::<target>();
     }
diff --git a/crates/typst-library/src/foundations/module.rs b/crates/typst-library/src/foundations/module.rs
index 2001aca1..3ee59c10 100644
--- a/crates/typst-library/src/foundations/module.rs
+++ b/crates/typst-library/src/foundations/module.rs
@@ -7,14 +7,20 @@ use typst_syntax::FileId;
 use crate::diag::StrResult;
 use crate::foundations::{repr, ty, Content, Scope, Value};
 
-/// An evaluated module, either built-in or resulting from a file.
+/// An module of definitions.
 ///
-/// You can access definitions from the module using
-/// [field access notation]($scripting/#fields) and interact with it using the
-/// [import and include syntaxes]($scripting/#modules). Alternatively, it is
-/// possible to convert a module to a dictionary, and therefore access its
-/// contents dynamically, using the
-/// [dictionary constructor]($dictionary/#constructor).
+/// A module
+/// - be built-in
+/// - stem from a [file import]($scripting/#modules)
+/// - stem from a [package import]($scripting/#packages) (and thus indirectly
+///   its entrypoint file)
+/// - result from a call to the [plugin]($plugin) function
+///
+/// You can access definitions from the module using [field access
+/// notation]($scripting/#fields) and interact with it using the [import and
+/// include syntaxes]($scripting/#modules). Alternatively, it is possible to
+/// convert a module to a dictionary, and therefore access its contents
+/// dynamically, using the [dictionary constructor]($dictionary/#constructor).
 ///
 /// # Example
 /// ```example
diff --git a/crates/typst-library/src/foundations/ops.rs b/crates/typst-library/src/foundations/ops.rs
index 7dbdde8f..6c240844 100644
--- a/crates/typst-library/src/foundations/ops.rs
+++ b/crates/typst-library/src/foundations/ops.rs
@@ -447,7 +447,6 @@ pub fn equal(lhs: &Value, rhs: &Value) -> bool {
         (Args(a), Args(b)) => a == b,
         (Type(a), Type(b)) => a == b,
         (Module(a), Module(b)) => a == b,
-        (Plugin(a), Plugin(b)) => a == b,
         (Datetime(a), Datetime(b)) => a == b,
         (Duration(a), Duration(b)) => a == b,
         (Dyn(a), Dyn(b)) => a == b,
diff --git a/crates/typst-library/src/foundations/plugin.rs b/crates/typst-library/src/foundations/plugin.rs
index d41261ed..cbc0f52d 100644
--- a/crates/typst-library/src/foundations/plugin.rs
+++ b/crates/typst-library/src/foundations/plugin.rs
@@ -4,43 +4,27 @@ use std::sync::{Arc, Mutex};
 
 use ecow::{eco_format, EcoString};
 use typst_syntax::Spanned;
-use wasmi::{AsContext, AsContextMut};
+use wasmi::Memory;
 
 use crate::diag::{bail, At, SourceResult, StrResult};
 use crate::engine::Engine;
-use crate::foundations::{func, repr, scope, ty, Bytes};
+use crate::foundations::{cast, func, scope, Bytes, Func, Module, Scope, Value};
 use crate::loading::{DataSource, Load};
 
-/// A WebAssembly plugin.
+/// Loads a WebAssembly module.
 ///
-/// Typst is capable of interfacing with plugins compiled to WebAssembly. Plugin
-/// functions may accept multiple [byte buffers]($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.
+/// The resulting [module] will contain one Typst [function] for each function
+/// export of the loaded WebAssembly module.
 ///
-/// 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/tree/master/crates/wasi-stub).
+/// Typst WebAssembly plugins need to follow a specific
+/// [protocol]($plugin/#protocol). To run as a plugin, a program needs to be
+/// compiled to a 32-bit shared WebAssembly library. Plugin functions may accept
+/// multiple [byte buffers]($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 and Packages
-/// Plugins are distributed as packages. A package can make use of a plugin
-/// simply by including a WebAssembly file and loading it. Because the
-/// byte-based plugin interface is quite low-level, plugins are typically
-/// exposed through wrapper functions, that also live in the same package.
-///
-/// # Purity
-/// Plugin functions must be pure: Given the same arguments, they must always
-/// return the same value. The reason for this is that Typst functions must be
-/// pure (which is quite fundamental to the language design) and, since Typst
-/// function can call plugin functions, this requirement is inherited. In
-/// particular, if a plugin function is called twice with the same arguments,
-/// Typst might cache the results and call your function only once.
+/// For security reasons, plugins run in isolation from your system. This means
+/// that printing, reading files, or similar things are not supported.
 ///
 /// # Example
 /// ```example
@@ -55,6 +39,50 @@ use crate::loading::{DataSource, Load};
 /// #concat("hello", "world")
 /// ```
 ///
+/// Since the plugin function returns a module, it can be used with import
+/// syntax:
+/// ```typ
+/// #import plugin("hello.wasm"): concatenate
+/// ```
+///
+/// # Purity
+/// Plugin functions **must be pure:** A plugin function call most not have any
+/// observable side effects on future plugin calls and given the same arguments,
+/// it must always return the same value.
+///
+/// The reason for this is that Typst functions must be pure (which is quite
+/// fundamental to the language design) and, since Typst function can call
+/// plugin functions, this requirement is inherited. In particular, if a plugin
+/// function is called twice with the same arguments, Typst might cache the
+/// results and call your function only once. Moreover, Typst may run multiple
+/// instances of your plugin in multiple threads, with no state shared between
+/// them.
+///
+/// Typst does not enforce plugin function purity (for efficiency reasons), but
+/// calling an impure function will lead to unpredictable and irreproducible
+/// results and must be avoided.
+///
+/// That said, mutable operations _can be_ useful for plugins that require
+/// costly runtime initialization. Due to the purity requirement, such
+/// initialization cannot be performed through a normal function call. Instead,
+/// Typst exposes a [plugin transition API]($plugin.transition), which executes
+/// a function call and then creates a derived module with new functions which
+/// will observe the side effects produced by the transition call. The original
+/// plugin remains unaffected.
+///
+/// # Plugins and Packages
+/// Any Typst code can make use of a plugin simply by including a WebAssembly
+/// file and loading it. However, because the byte-based plugin interface is
+/// quite low-level, plugins are typically exposed through a package containing
+/// the plugin and idiomatic wrapper functions.
+///
+/// # WASI
+/// 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/tree/master/crates/wasi-stub).
+///
 /// # Protocol
 /// To be used as a plugin, a WebAssembly module must conform to the following
 /// protocol:
@@ -67,8 +95,8 @@ use crate::loading::{DataSource, Load};
 ///   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
+/// - 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
@@ -85,19 +113,21 @@ use crate::loading::{DataSource, Load};
 ///   then interpreted as an UTF-8 encoded error message.
 ///
 /// ## Imports
-/// Plugin modules need to import two functions that are provided by the runtime.
-/// (Types and functions are described using WAT syntax.)
+/// 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)))`
+/// - `(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](#exports) 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.
+///   buffer. When a plugin function is called, it [receives the
+///   lengths](#exports) 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)))`
+/// - `(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
@@ -106,72 +136,147 @@ use crate::loading::{DataSource, Load};
 ///   interpreted as an error message, it should be encoded as UTF-8.
 ///
 /// # Resources
-/// For more resources, check out the
-/// [wasm-minimal-protocol repository](https://github.com/astrale-sharp/wasm-minimal-protocol).
-/// It contains:
+/// 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
-#[ty(scope, cast)]
-#[derive(Clone)]
-pub struct Plugin(Arc<Repr>);
+#[func(scope)]
+pub fn plugin(
+    engine: &mut Engine,
+    /// A path to a WebAssembly file or raw WebAssembly bytes.
+    ///
+    /// For more details about paths, see the [Paths section]($syntax/#paths).
+    source: Spanned<DataSource>,
+) -> SourceResult<Module> {
+    let data = source.load(engine.world)?;
+    Plugin::module(data).at(source.span)
+}
 
-/// 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>,
+#[scope]
+impl plugin {
+    /// Calls a plugin function that has side effects and returns a new module
+    /// with plugin functions that are guaranteed to have observed the results
+    /// of the mutable call.
+    ///
+    /// Note that calling an impure function through a normal function call
+    /// (without use of the transition API) is forbidden and leads to
+    /// unpredictable behaviour. Read the [section on purity]($plugin/#purity)
+    /// for more details.
+    ///
+    /// In the example below, we load the plugin `hello-mut.wasm` which exports
+    /// two functions: The `get()` function retrieves a global array as a
+    /// string. The `add(value)` function adds a value to the global array.
+    ///
+    /// We call `add` via the transition API. The call `mutated.get()` on the
+    /// derived module will observe the addition. Meanwhile the original module
+    /// remains untouched as demonstrated by the `base.get()` call.
+    ///
+    /// _Note:_ Due to limitations in the internal WebAssembly implementation,
+    /// the transition API can only guarantee to reflect changes in the plugin's
+    /// memory, not in WebAssembly globals. If your plugin relies on changes to
+    /// globals being visible after transition, you might want to avoid use of
+    /// the transition API for now. We hope to lift this limitation in the
+    /// future.
+    ///
+    /// ```typ
+    /// #let base = plugin("hello-mut.wasm")
+    /// #assert.eq(base.get(), "[]")
+    ///
+    /// #let mutated = plugin.transition(base.add, "hello")
+    /// #assert.eq(base.get(), "[]")
+    /// #assert.eq(mutated.get(), "[hello]")
+    /// ```
+    #[func]
+    pub fn transition(
+        /// The plugin function to call.
+        func: PluginFunc,
+        /// The byte buffers to call the function with.
+        #[variadic]
+        arguments: Vec<Bytes>,
+    ) -> StrResult<Module> {
+        func.transition(arguments)
+    }
 }
 
-/// Owns all data associated with the WebAssembly module.
-type Store = wasmi::Store<StoreData>;
+/// A function loaded from a WebAssembly plugin.
+#[derive(Debug, Clone, PartialEq, Hash)]
+pub struct PluginFunc {
+    /// The underlying plugin, shared by this and the other functions.
+    plugin: Arc<Plugin>,
+    /// The name of the plugin function.
+    name: EcoString,
+}
 
-/// If there was an error reading/writing memory, keep the offset + length to
-/// display an error message.
-struct MemoryError {
-    offset: u32,
-    length: u32,
-    write: bool,
+impl PluginFunc {
+    /// The name of the plugin function.
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    /// Call the WebAssembly function with the given arguments.
+    #[comemo::memoize]
+    #[typst_macros::time(name = "call plugin")]
+    pub fn call(&self, args: Vec<Bytes>) -> StrResult<Bytes> {
+        self.plugin.call(&self.name, args)
+    }
+
+    /// Transition a plugin and turn the result into a module.
+    #[comemo::memoize]
+    #[typst_macros::time(name = "transition plugin")]
+    pub fn transition(&self, args: Vec<Bytes>) -> StrResult<Module> {
+        self.plugin.transition(&self.name, args).map(Plugin::into_module)
+    }
 }
-/// The persistent store data used for communication between store and host.
-#[derive(Default)]
-struct StoreData {
-    args: Vec<Bytes>,
-    output: Vec<u8>,
-    memory_error: Option<MemoryError>,
+
+cast! {
+    PluginFunc,
+    self => Value::Func(self.into()),
+    v: Func => v.to_plugin().ok_or("expected plugin function")?.clone(),
 }
 
-#[scope]
-impl Plugin {
-    /// Creates a new plugin from a WebAssembly file.
-    #[func(constructor)]
-    pub fn construct(
-        engine: &mut Engine,
-        /// A path to a WebAssembly file or raw WebAssembly bytes.
-        ///
-        /// For more details about paths, see the [Paths section]($syntax/#paths).
-        source: Spanned<DataSource>,
-    ) -> SourceResult<Plugin> {
-        let data = source.load(engine.world)?;
-        Plugin::new(data).at(source.span)
-    }
+/// A plugin with potentially multiple instances for multi-threaded
+/// execution.
+struct Plugin {
+    /// Shared by all variants of the plugin.
+    base: Arc<PluginBase>,
+    /// A pool of plugin instances.
+    ///
+    /// When multiple plugin calls run concurrently due to multi-threading, we
+    /// create new instances whenever we run out of ones.
+    pool: Mutex<Vec<PluginInstance>>,
+    /// A snapshot that new instances should be restored to.
+    snapshot: Option<Snapshot>,
+    /// A combined hash that incorporates all function names and arguments used
+    /// in transitions of this plugin, such that this plugin has a deterministic
+    /// hash and equality check that can differentiate it from "siblings" (same
+    /// base, different transitions).
+    fingerprint: u128,
 }
 
 impl Plugin {
-    /// Create a new plugin from raw WebAssembly bytes.
+    /// Create a plugin and turn it into a module.
     #[comemo::memoize]
     #[typst_macros::time(name = "load plugin")]
-    pub fn new(bytes: Bytes) -> StrResult<Plugin> {
+    fn module(bytes: Bytes) -> StrResult<Module> {
+        Self::new(bytes).map(Self::into_module)
+    }
+
+    /// Create a new plugin from raw WebAssembly bytes.
+    fn new(bytes: Bytes) -> StrResult<Self> {
         let engine = wasmi::Engine::default();
         let module = wasmi::Module::new(&engine, bytes.as_slice())
             .map_err(|err| format!("failed to load WebAssembly module ({err})"))?;
 
+        // Ensure that the plugin exports its memory.
+        if !matches!(module.get_export("memory"), Some(wasmi::ExternType::Memory(_))) {
+            bail!("plugin does not export its memory");
+        }
+
         let mut linker = wasmi::Linker::new(&engine);
         linker
             .func_wrap(
@@ -188,58 +293,174 @@ impl Plugin {
             )
             .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}"))?;
+        let base = Arc::new(PluginBase { bytes, linker, module });
+        let instance = PluginInstance::new(&base, None)?;
 
-        // 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");
+        Ok(Self {
+            base,
+            snapshot: None,
+            fingerprint: 0,
+            pool: Mutex::new(vec![instance]),
+        })
+    }
+
+    /// Execute a function with access to an instsance.
+    fn call(&self, func: &str, args: Vec<Bytes>) -> StrResult<Bytes> {
+        // Acquire an instance from the pool (potentially creating a new one).
+        let mut instance = self.acquire()?;
+
+        // Execute the call on an instance from the pool. If the call fails, we
+        // return early and _don't_ return the instance to the pool as it might
+        // be irrecoverably damaged.
+        let output = instance.call(func, args)?;
+
+        // Return the instance to the pool.
+        self.pool.lock().unwrap().push(instance);
+
+        Ok(output)
+    }
+
+    /// Call a mutable plugin function, producing a new mutable whose functions
+    /// are guaranteed to be able to observe the mutation.
+    fn transition(&self, func: &str, args: Vec<Bytes>) -> StrResult<Plugin> {
+        // Derive a new transition hash from the old one and the function and arguments.
+        let fingerprint = typst_utils::hash128(&(self.fingerprint, func, &args));
+
+        // Execute the mutable call on an instance.
+        let mut instance = self.acquire()?;
+
+        // Call the function. If the call fails, we return early and _don't_
+        // return the instance to the pool as it might be irrecoverably damaged.
+        instance.call(func, args)?;
+
+        // Snapshot the instance after the mutable call.
+        let snapshot = instance.snapshot();
+
+        // Create a new plugin and move (this is important!) the used instance
+        // into it, so that the old plugin won't observe the mutation. Also
+        // save the snapshot so that instances that are initialized for the
+        // transitioned plugin's pool observe the mutation.
+        Ok(Self {
+            base: self.base.clone(),
+            snapshot: Some(snapshot),
+            fingerprint,
+            pool: Mutex::new(vec![instance]),
+        })
+    }
+
+    /// Acquire an instance from the pool (or create a new one).
+    fn acquire(&self) -> StrResult<PluginInstance> {
+        // Don't use match to ensure that the lock is released before we create
+        // a new instance.
+        if let Some(instance) = self.pool.lock().unwrap().pop() {
+            return Ok(instance);
         }
 
-        // Collect exported functions.
-        let functions = instance
-            .exports(&store)
-            .filter_map(|export| {
-                let name = export.name().into();
-                export.into_func().map(|func| (name, func))
-            })
-            .collect();
+        PluginInstance::new(&self.base, self.snapshot.as_ref())
+    }
 
-        Ok(Plugin(Arc::new(Repr { bytes, functions, store: Mutex::new(store) })))
+    /// Turn a plugin into a Typst module containing plugin functions.
+    fn into_module(self) -> Module {
+        let shared = Arc::new(self);
+
+        // Build a scope from the collected functions.
+        let mut scope = Scope::new();
+        for export in shared.base.module.exports() {
+            if matches!(export.ty(), wasmi::ExternType::Func(_)) {
+                let name = EcoString::from(export.name());
+                let func = PluginFunc { plugin: shared.clone(), name: name.clone() };
+                scope.define(name, Func::from(func));
+            }
+        }
+
+        Module::anonymous(scope)
     }
+}
 
-    /// Call the plugin function with the given `name`.
-    #[comemo::memoize]
-    #[typst_macros::time(name = "call plugin")]
-    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}")
-            })?;
+impl Debug for Plugin {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        f.pad("Plugin(..)")
+    }
+}
 
-        let mut store = self.0.store.lock().unwrap();
-        let ty = func.ty(store.as_context());
+impl PartialEq for Plugin {
+    fn eq(&self, other: &Self) -> bool {
+        self.base.bytes == other.base.bytes && self.fingerprint == other.fingerprint
+    }
+}
+
+impl Hash for Plugin {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.base.bytes.hash(state);
+        self.fingerprint.hash(state);
+    }
+}
+
+/// Shared by all pooled & transitioned variants of the plugin.
+struct PluginBase {
+    /// The raw WebAssembly bytes.
+    bytes: Bytes,
+    /// The compiled WebAssembly module.
+    module: wasmi::Module,
+    /// A linker used to create a `Store` for execution.
+    linker: wasmi::Linker<CallData>,
+}
+
+/// An single plugin instance for single-threaded execution.
+struct PluginInstance {
+    /// The underlying wasmi instance.
+    instance: wasmi::Instance,
+    /// The execution store of this concrete plugin instance.
+    store: wasmi::Store<CallData>,
+}
+
+/// A snapshot of a plugin instance.
+struct Snapshot {
+    /// The number of pages in the main memory.
+    mem_pages: u32,
+    /// The data in the main memory.
+    mem_data: Vec<u8>,
+}
 
-        // Check function signature.
+impl PluginInstance {
+    /// Create a new execution instance of a plugin, potentially restoring
+    /// a snapshot.
+    #[typst_macros::time(name = "create plugin instance")]
+    fn new(base: &PluginBase, snapshot: Option<&Snapshot>) -> StrResult<PluginInstance> {
+        let mut store = wasmi::Store::new(base.linker.engine(), CallData::default());
+        let instance = base
+            .linker
+            .instantiate(&mut store, &base.module)
+            .and_then(|pre_instance| pre_instance.start(&mut store))
+            .map_err(|e| eco_format!("{e}"))?;
+
+        let mut instance = PluginInstance { instance, store };
+        if let Some(snapshot) = snapshot {
+            instance.restore(snapshot);
+        }
+        Ok(instance)
+    }
+
+    /// Call a plugin function with byte arguments.
+    fn call(&mut self, func: &str, args: Vec<Bytes>) -> StrResult<Bytes> {
+        let handle = self
+            .instance
+            .get_export(&self.store, func)
+            .unwrap()
+            .into_func()
+            .unwrap();
+        let ty = handle.ty(&self.store);
+
+        // Check function signature. Do this lazily only when a function is called
+        // because there might be exported functions like `_initialize` that don't
+        // match the schema.
         if ty.params().iter().any(|&v| v != wasmi::core::ValType::I32) {
             bail!(
-                "plugin function `{name}` has a parameter that is not a 32-bit integer"
+                "plugin function `{func}` has a parameter that is not a 32-bit integer"
             );
         }
         if ty.results() != [wasmi::core::ValType::I32] {
-            bail!("plugin function `{name}` does not return exactly one 32-bit integer");
+            bail!("plugin function `{func}` does not return exactly one 32-bit integer");
         }
 
         // Check inputs.
@@ -260,23 +481,26 @@ impl Plugin {
             .collect::<Vec<_>>();
 
         // Store the input data.
-        store.data_mut().args = args;
+        self.store.data_mut().args = args;
 
         // Call the function.
         let mut code = wasmi::Val::I32(-1);
-        func.call(store.as_context_mut(), &lengths, std::slice::from_mut(&mut code))
+        handle
+            .call(&mut self.store, &lengths, std::slice::from_mut(&mut code))
             .map_err(|err| eco_format!("plugin panicked: {err}"))?;
+
         if let Some(MemoryError { offset, length, write }) =
-            store.data_mut().memory_error.take()
+            self.store.data_mut().memory_error.take()
         {
             return Err(eco_format!(
-                "plugin tried to {kind} out of bounds: pointer {offset:#x} is out of bounds for {kind} of length {length}",
+                "plugin tried to {kind} out of bounds: \
+                 pointer {offset:#x} is out of bounds for {kind} of length {length}",
                 kind = if write { "write" } else { "read" }
             ));
         }
 
         // Extract the returned data.
-        let output = std::mem::take(&mut store.data_mut().output);
+        let output = std::mem::take(&mut self.store.data_mut().output);
 
         // Parse the functions return value.
         match code {
@@ -293,39 +517,63 @@ impl Plugin {
         Ok(Bytes::new(output))
     }
 
-    /// 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)
+    /// Creates a snapshot of this instance from which another one can be
+    /// initialized.
+    #[typst_macros::time(name = "save snapshot")]
+    fn snapshot(&self) -> Snapshot {
+        let memory = self.memory();
+        let mem_pages = memory.size(&self.store);
+        let mem_data = memory.data(&self.store).to_vec();
+        Snapshot { mem_pages, mem_data }
     }
-}
 
-impl Debug for Plugin {
-    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
-        f.pad("Plugin(..)")
+    /// Restores the instance to a snapshot.
+    #[typst_macros::time(name = "restore snapshot")]
+    fn restore(&mut self, snapshot: &Snapshot) {
+        let memory = self.memory();
+        let current_size = memory.size(&self.store);
+        if current_size < snapshot.mem_pages {
+            memory
+                .grow(&mut self.store, snapshot.mem_pages - current_size)
+                .unwrap();
+        }
+
+        memory.data_mut(&mut self.store)[..snapshot.mem_data.len()]
+            .copy_from_slice(&snapshot.mem_data);
     }
-}
 
-impl repr::Repr for Plugin {
-    fn repr(&self) -> EcoString {
-        "plugin(..)".into()
+    /// Retrieves a handle to the plugin's main memory.
+    fn memory(&self) -> Memory {
+        self.instance
+            .get_export(&self.store, "memory")
+            .unwrap()
+            .into_memory()
+            .unwrap()
     }
 }
 
-impl PartialEq for Plugin {
-    fn eq(&self, other: &Self) -> bool {
-        self.0.bytes == other.0.bytes
-    }
+/// The persistent store data used for communication between store and host.
+#[derive(Default)]
+struct CallData {
+    /// Arguments for a current call.
+    args: Vec<Bytes>,
+    /// The results of the current call.
+    output: Vec<u8>,
+    /// A memory error that occured during execution of the current call.
+    memory_error: Option<MemoryError>,
 }
 
-impl Hash for Plugin {
-    fn hash<H: Hasher>(&self, state: &mut H) {
-        self.0.bytes.hash(state);
-    }
+/// If there was an error reading/writing memory, keep the offset + length to
+/// display an error message.
+struct MemoryError {
+    offset: u32,
+    length: u32,
+    write: bool,
 }
 
 /// Write the arguments to the plugin function into the plugin's memory.
 fn wasm_minimal_protocol_write_args_to_buffer(
-    mut caller: wasmi::Caller<StoreData>,
+    mut caller: wasmi::Caller<CallData>,
     ptr: u32,
 ) {
     let memory = caller.get_export("memory").unwrap().into_memory().unwrap();
@@ -346,7 +594,7 @@ fn wasm_minimal_protocol_write_args_to_buffer(
 
 /// Extracts the output of the plugin function from the plugin's memory.
 fn wasm_minimal_protocol_send_result_to_host(
-    mut caller: wasmi::Caller<StoreData>,
+    mut caller: wasmi::Caller<CallData>,
     ptr: u32,
     len: u32,
 ) {
diff --git a/crates/typst-library/src/foundations/scope.rs b/crates/typst-library/src/foundations/scope.rs
index 99c9a37e..b7b4a6d9 100644
--- a/crates/typst-library/src/foundations/scope.rs
+++ b/crates/typst-library/src/foundations/scope.rs
@@ -167,6 +167,14 @@ impl Scope {
         Default::default()
     }
 
+    /// Create a new scope with the given capacity.
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            map: IndexMap::with_capacity(capacity),
+            ..Default::default()
+        }
+    }
+
     /// Create a new scope with duplication prevention.
     pub fn deduplicating() -> Self {
         Self { deduplicate: true, ..Default::default() }
diff --git a/crates/typst-library/src/foundations/value.rs b/crates/typst-library/src/foundations/value.rs
index d9902772..4fa380b4 100644
--- a/crates/typst-library/src/foundations/value.rs
+++ b/crates/typst-library/src/foundations/value.rs
@@ -15,8 +15,8 @@ use crate::diag::{HintedStrResult, HintedString, StrResult};
 use crate::foundations::{
     fields, ops, repr, Args, Array, AutoValue, Bytes, CastInfo, Content, Datetime,
     Decimal, Dict, Duration, Fold, FromValue, Func, IntoValue, Label, Module,
-    NativeElement, NativeType, NoneValue, Plugin, Reflect, Repr, Resolve, Scope, Str,
-    Styles, Symbol, SymbolElem, Type, Version,
+    NativeElement, NativeType, NoneValue, Reflect, Repr, Resolve, Scope, Str, Styles,
+    Symbol, SymbolElem, Type, Version,
 };
 use crate::layout::{Abs, Angle, Em, Fr, Length, Ratio, Rel};
 use crate::text::{RawContent, RawElem, TextElem};
@@ -84,8 +84,6 @@ pub enum Value {
     Type(Type),
     /// A module.
     Module(Module),
-    /// A WebAssembly plugin.
-    Plugin(Plugin),
     /// A dynamic value.
     Dyn(Dynamic),
 }
@@ -147,7 +145,6 @@ impl Value {
             Self::Args(_) => Type::of::<Args>(),
             Self::Type(_) => Type::of::<Type>(),
             Self::Module(_) => Type::of::<Module>(),
-            Self::Plugin(_) => Type::of::<Plugin>(),
             Self::Dyn(v) => v.ty(),
         }
     }
@@ -251,7 +248,6 @@ impl Debug for Value {
             Self::Args(v) => Debug::fmt(v, f),
             Self::Type(v) => Debug::fmt(v, f),
             Self::Module(v) => Debug::fmt(v, f),
-            Self::Plugin(v) => Debug::fmt(v, f),
             Self::Dyn(v) => Debug::fmt(v, f),
         }
     }
@@ -289,7 +285,6 @@ impl Repr for Value {
             Self::Args(v) => v.repr(),
             Self::Type(v) => v.repr(),
             Self::Module(v) => v.repr(),
-            Self::Plugin(v) => v.repr(),
             Self::Dyn(v) => v.repr(),
         }
     }
@@ -340,7 +335,6 @@ impl Hash for Value {
             Self::Args(v) => v.hash(state),
             Self::Type(v) => v.hash(state),
             Self::Module(v) => v.hash(state),
-            Self::Plugin(v) => v.hash(state),
             Self::Dyn(v) => v.hash(state),
         }
     }
@@ -661,7 +655,6 @@ primitive! {
 primitive! { Args: "arguments", Args }
 primitive! { Type: "type", Type }
 primitive! { Module: "module", Module }
-primitive! { Plugin: "plugin", Plugin }
 
 impl<T: Reflect> Reflect for Arc<T> {
     fn input() -> CastInfo {