perf(wasm): retain indexed forge arena across group calls

packages/forge/README.md

@@ -94,8 +94,9 @@ have no `.` or `..` parts, and no two files may share one.
 
 Things to know:
 
-- **Each group runs in an instance of its own.** `generateGroup` sees
-  nothing that `generate` or another group kept in module state.
+- **Groups run in parallel across a pool of instances.** Instances are
+  reused across groups of the same forge, so module state may carry over
+  between calls.
 - **There is no Node.** Plugins have no `fs`, `process`, network or clock.
   Libraries in plain JavaScript bundle in and work.
 - **QuickJS runs the plugin.** `forge-plugin build` compiles it in, for a

packages/forge/plugin/forge.wit

@@ -219,6 +219,7 @@ world plugin {
     /// Generates files for the forge. Called exactly once.
     export generate: func() -> result<list<file>, string>;
     /// Generates files for a group. Called exactly once for each group, in
-    /// parallel with other groups, each in an instance of its own.
+    /// parallel with other groups. Instances are reused across groups of
+    /// the same forge.
     export generate-group: func(group: group-id) -> result<list<file>, string>;
 }

packages/forge/src/index.ts

@@ -18,9 +18,9 @@
  * ```
  *
  * The forge stays in forge's process. Each part is fetched the first time
- * a plugin reads it, and kept until the call ends, so reading a part again
- * is free, and the same part is always the same object: compare groups,
- * methods and models with `===`.
+ * a plugin reads it, and kept in the instance across group calls, so
+ * reading a part again is free, and the same part is always the same
+ * object: compare groups, methods and models with `===`.
  *
  * @module
  */
@@ -44,8 +44,8 @@ export interface Plugin {
   generate?(forge: Forge): Iterable<File>;
   /**
    * Generates files for a group. Called exactly once for each group, in
-   * parallel with other groups, each in an instance of its own. Nothing a
-   * call keeps in module state is seen by the calls for other groups.
+   * parallel with other groups. Instances are reused across groups of the
+   * same forge, so module state may carry over between calls.
    */
   generateGroup?(group: Group): Iterable<File>;
 }
@@ -382,12 +382,13 @@ export interface ModelType {
  * Throw to fail: forge stops and prints the error's message.
  */
 export function definePlugin(plugin: Plugin): Exports {
+  const forge = lazy(() => new ForgeImpl());
   const exports = {
     generate() {
-      return call(() => files(plugin.generate?.(new ForgeImpl()) ?? []));
+      return call(() => files(plugin.generate?.(forge()) ?? []));
     },
     generateGroup(id: host.GroupId) {
-      return call(() => files(plugin.generateGroup?.(new ForgeImpl().group(id)) ?? []));
+      return call(() => files(plugin.generateGroup?.(forge().group(id)) ?? []));
     },
   } satisfies Pick<typeof hostPlugin, 'generate' | 'generateGroup'>;
   return exports;
@@ -497,12 +498,12 @@ function lazy<T>(f: () => T): () => T {
   };
 }
 
-/** Returns what `cache` has under `id`, making it with `make` if it has nothing. */
-function cached<T>(cache: Map<number, T>, id: number, make: () => T): T {
-  let value = cache.get(id);
+/** Returns what `cache` has at index `id`, making it with `make` if it has nothing. */
+function cached<T>(cache: (T | undefined)[], id: number, make: () => T): T {
+  let value = cache[id];
   if (value === undefined) {
     value = make();
-    cache.set(id, value);
+    cache[id] = value;
   }
   return value;
 }
@@ -554,21 +555,24 @@ function upperFirst(word: string): string {
 }
 
 class ForgeImpl implements Forge {
-  readonly version: string;
-  readonly description: string;
-  readonly baseUrl: string;
-  private readonly groupCache = new Map<host.GroupId, GroupImpl>();
-  private readonly methodCache = new Map<host.MethodId, MethodImpl>();
-  private readonly modelCache = new Map<host.ModelId, ModelImpl>();
-  private readonly typeCache = new Map<host.TypeId, Type>();
+  private readonly info = lazy(() => host.info());
+  private readonly groupCache: (GroupImpl | undefined)[] = [];
+  private readonly methodCache: (MethodImpl | undefined)[] = [];
+  private readonly modelCache: (ModelImpl | undefined)[] = [];
+  private readonly typeCache: (Type | undefined)[] = [];
   private readonly allGroups = lazy(() => ids(host.groups()).map((id) => this.group(id)));
   private readonly topModels = lazy(() => ids(host.models()).map((id) => this.model(id)));
 
-  constructor() {
-    const info = host.info();
-    this.version = info.version;
-    this.description = info.description;
-    this.baseUrl = info.baseUrl;
+  get version(): string {
+    return this.info().version;
+  }
+
+  get description(): string {
+    return this.info().description;
+  }
+
+  get baseUrl(): string {
+    return this.info().baseUrl;
   }
 
   get groups(): readonly Group[] {

src/plugin/wasm.rs

@@ -280,7 +280,8 @@ world plugin {
     /// Generates files for the forge. Called exactly once.
     export generate: func() -> result<list<file>, string>;
     /// Generates files for a group. Called exactly once for each group, in
-    /// parallel with other groups, each in an instance of its own.
+    /// parallel with other groups. Instances are reused across groups of
+    /// the same forge.
     export generate-group: func(group: group-id) -> result<list<file>, string>;
 }
 "#
@@ -291,14 +292,42 @@ use bindings::forge::plugin::forge as wit;
 /// A plugin compiled to a WebAssembly component of the `plugin` world in
 /// [`WIT`].
 ///
-/// It compiles the component once, then keeps a pool of instances, so
-/// groups run in parallel without compiling or instantiating it again. A
-/// call takes an instance from the pool, or makes one if none is free, and
-/// returns it unless the plugin trapped.
+/// It compiles the component once, then keeps a pool of instances for the
+/// current [`Forge`], so groups run in parallel without compiling or
+/// instantiating it again, and each instance only ever reads one [`Forge`].
+/// A call takes an instance from the pool, or makes one if none is free,
+/// and returns it unless the plugin trapped.
 pub struct WasmPlugin {
     engine: Engine,
     pre: bindings::PluginPre<State>,
-    pool: Mutex<Vec<Instance>>,
+    pool: Mutex<Pool>,
+}
+
+#[derive(Default)]
+struct Pool {
+    forge: ForgeKey,
+    instances: Vec<Instance>,
+}
+
+#[derive(Clone, Copy, Default, PartialEq, Eq)]
+struct ForgeKey {
+    forge: usize,
+    groups: (usize, usize),
+    methods: (usize, usize),
+    models: (usize, usize),
+    types: (usize, usize),
+}
+
+impl ForgeKey {
+    fn of(forge: &Forge) -> Self {
+        ForgeKey {
+            forge: ptr::from_ref(forge).addr(),
+            groups: (forge.groups.as_ptr().addr(), forge.groups.len()),
+            methods: (forge.methods.as_ptr().addr(), forge.methods.len()),
+            models: (forge.models.as_ptr().addr(), forge.models.len()),
+            types: (forge.types.as_ptr().addr(), forge.types.len()),
+        }
+    }
 }
 
 /// A file a [`WasmPlugin`] generated, with its content already rendered.
@@ -331,7 +360,7 @@ impl WasmPlugin {
         Ok(WasmPlugin {
             engine,
             pre,
-            pool: Mutex::new(Vec::new()),
+            pool: Mutex::new(Pool::default()),
         })
     }
 
@@ -341,7 +370,15 @@ impl WasmPlugin {
         forge: &Forge,
         f: impl FnOnce(&bindings::Plugin, &mut Store<State>) -> wasmtime::Result<Result<Vec<bindings::File>, String>>,
     ) -> Result<std::vec::IntoIter<WasmFile>, Err> {
-        let pooled = self.pool.lock().unwrap_or_else(|err| err.into_inner()).pop();
+        let forge_key = ForgeKey::of(forge);
+        let pooled = {
+            let mut pool = self.pool.lock().unwrap_or_else(|err| err.into_inner());
+            if pool.forge != forge_key {
+                pool.forge = forge_key;
+                pool.instances.clear();
+            }
+            pool.instances.pop()
+        };
         let mut instance = match pooled {
             Some(instance) => instance,
             None => self.instantiate()?,
@@ -353,7 +390,11 @@ impl WasmPlugin {
 
         // A trap can leave the instance broken, so it is dropped instead.
         let files = result.map_err(error)?;
-        self.pool.lock().unwrap_or_else(|err| err.into_inner()).push(instance);
+        let mut pool = self.pool.lock().unwrap_or_else(|err| err.into_inner());
+        if pool.forge == forge_key {
+            pool.instances.push(instance);
+        }
+        drop(pool);
 
         let files = files?.into_iter().map(|file| WasmFile {
             path: file.path,

tests/wasm/ts/Cargo.toml

@@ -21,7 +21,6 @@ crate-type = ["cdylib"]
 dump = []
 
 [dependencies]
-elsa = "1"
 # Reads the JSON text the host sends for a `Json`, keeping objects in order.
 serde_json = { version = "1", features = ["preserve_order"] }
 wit-bindgen = "0.57"