feat: add machine factory support

core/src/main/java/org/extism/sdk/chicory/ChicoryModule.java

@@ -2,10 +2,12 @@
 
 import com.dylibso.chicory.runtime.ByteArrayMemory;
 import com.dylibso.chicory.runtime.Instance;
+import com.dylibso.chicory.runtime.Machine;
 import com.dylibso.chicory.wasm.Parser;
 import com.dylibso.chicory.wasm.WasmModule;
 
 import java.nio.file.Path;
+import java.util.function.Function;
 
 class ChicoryModule {
 
@@ -29,11 +31,11 @@ static WasmModule fromWasm(ManifestWasm m) {
         }
     }
 
-    static Instance.Builder instanceWithOptions(Instance.Builder m, Manifest.Options opts, CachedAotMachineFactory aotMachineFactory) {
+    static Instance.Builder instanceWithOptions(Instance.Builder m, Manifest.Options opts, Function<Instance, Machine> aotMachineFactory) {
         if (opts == null) {
             return m;
         }
-        // This feature is not compatibly with the native-image builder.
+        // This feature is not compatible with the native-image builder.
         if (opts.aot && !IS_NATIVE_IMAGE_AOT) {
             m.withMachineFactory(aotMachineFactory);
         }

core/src/main/java/org/extism/sdk/chicory/DependencyGraph.java

@@ -6,6 +6,7 @@
 import com.dylibso.chicory.runtime.ImportFunction;
 import com.dylibso.chicory.runtime.ImportValues;
 import com.dylibso.chicory.runtime.Instance;
+import com.dylibso.chicory.runtime.Machine;
 import com.dylibso.chicory.runtime.Store;
 import com.dylibso.chicory.runtime.WasmFunctionHandle;
 import com.dylibso.chicory.wasi.WasiExitException;
@@ -26,6 +27,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.Stack;
+import java.util.function.Function;
 
 import static java.util.stream.Collectors.groupingBy;
 
@@ -42,7 +44,7 @@ class DependencyGraph {
 
     private final Store store = new Store();
     private Manifest.Options options;
-    private CachedAotMachineFactory aotMachineFactory;
+    private Function<Instance, Machine> machineFactory;
 
     public DependencyGraph(Logger logger) {
         this.logger = logger;
@@ -53,8 +55,14 @@ public DependencyGraph(Logger logger) {
      */
     public void setOptions(Manifest.Options options) {
         this.options = options;
-        if (options != null && options.aot) {
-            this.aotMachineFactory = new CachedAotMachineFactory();
+        if (options != null) {
+            if (options.aot && options.machineFactory == null) {
+                this.machineFactory = new CachedAotMachineFactory();
+                return;
+            }
+            if (options.machineFactory != null) {
+                this.machineFactory = options.machineFactory;
+            }
         }
     }
 
@@ -98,10 +106,12 @@ private void checkCollision(String moduleName, String symbol) {
      */
     public void registerModule(String name, WasmModule m) {
         checkCollision(name, null);
-        if (aotMachineFactory != null) {
-            aotMachineFactory.compile(m);
+        if (machineFactory != null) {
+            if (machineFactory instanceof CachedAotMachineFactory) {
+                var cachedAotMachineFactory = (CachedAotMachineFactory) machineFactory;
+                cachedAotMachineFactory.compile(m);
+            }
         }
-
         ExportSection exportSection = m.exportSection();
         for (int i = 0; i < exportSection.exportCount(); i++) {
             Export export = exportSection.getExport(i);
@@ -238,7 +248,7 @@ private Instance instantiate(String moduleId, List<HostFunction> moreHostFunctio
 
         Instance instance =
                 ChicoryModule.instanceWithOptions(
-                        Instance.builder(m), this.options, aotMachineFactory)
+                        Instance.builder(m), this.options, machineFactory)
                         .withImportValues(importValues)
                         .withStart(false)
                         .build();

core/src/main/java/org/extism/sdk/chicory/Kernel.java

@@ -3,12 +3,14 @@
 import com.dylibso.chicory.runtime.ExportFunction;
 import com.dylibso.chicory.runtime.HostFunction;
 import com.dylibso.chicory.runtime.Instance;
+import com.dylibso.chicory.runtime.Machine;
 import com.dylibso.chicory.wasm.Parser;
 import com.dylibso.chicory.wasm.WasmModule;
 import com.dylibso.chicory.wasm.types.FunctionType;
 import com.dylibso.chicory.wasm.types.ValType;
 
 import java.util.List;
+import java.util.function.Function;
 
 public class Kernel {
 
@@ -39,7 +41,7 @@ public Kernel() {
         this(null);
     }
 
-    Kernel(CachedAotMachineFactory machineFactory) {
+    Kernel(Function<Instance, Machine> machineFactory) {
         Instance kernel = instance(machineFactory);
         instanceMemory = kernel.memory();
         alloc = kernel.export("alloc");
@@ -64,11 +66,11 @@ public Kernel() {
         memoryBytes = kernel.export("memory_bytes");
     }
 
-    private static Instance instance(CachedAotMachineFactory machineFactory) {
+    private static Instance instance(Function<Instance, Machine> machineFactory) {
         var kernelStream = Kernel.class.getClassLoader().getResourceAsStream("extism-runtime.wasm");
         WasmModule module = Parser.parse(kernelStream);
-        if (machineFactory != null) {
-            machineFactory.compile(module);
+        if (machineFactory != null && machineFactory instanceof CachedAotMachineFactory) {
+            ((CachedAotMachineFactory) machineFactory).compile(module);
         }
         return Instance.builder(module).withMachineFactory(machineFactory).build();
     }

core/src/main/java/org/extism/sdk/chicory/Linker.java

@@ -3,11 +3,13 @@
 import com.dylibso.chicory.log.Logger;
 import com.dylibso.chicory.runtime.HostFunction;
 import com.dylibso.chicory.runtime.Instance;
+import com.dylibso.chicory.runtime.Machine;
 import com.dylibso.chicory.wasi.WasiOptions;
 import com.dylibso.chicory.wasi.WasiPreview1;
 import org.extism.sdk.chicory.http.HttpConfig;
 
 import java.util.Arrays;
+import java.util.function.Function;
 
 
 /**
@@ -37,17 +39,23 @@ Plugin link() {
         ConfigProvider config;
         String[] allowedHosts;
         WasiOptions wasiOptions;
-        CachedAotMachineFactory aotMachineFactory;
+        Function<Instance, Machine> machineFactory = null;
         HttpConfig httpConfig;
-        dg.setOptions(manifest.options);
-        config = manifest.options.config;
-        allowedHosts = manifest.options.allowedHosts;
-        wasiOptions = manifest.options.wasiOptions;
-        aotMachineFactory = manifest.options.aot? new CachedAotMachineFactory() : null;
-        httpConfig = manifest.options.httpConfig;
+        Manifest.Options options = manifest.options;
+        dg.setOptions(options);
+        config = options.config;
+        allowedHosts = options.allowedHosts;
+        wasiOptions = options.wasiOptions;
+        if (options.aot && options.machineFactory == null) {
+            machineFactory = new CachedAotMachineFactory();
+        }
+        if (options.machineFactory != null) {
+            machineFactory = options.machineFactory;
+        }
+        httpConfig = options.httpConfig;
 
         // Register the HostEnv exports.
-        var hostEnv = new HostEnv(new Kernel(aotMachineFactory), config, allowedHosts, httpConfig, logger);
+        var hostEnv = new HostEnv(new Kernel(machineFactory), config, allowedHosts, httpConfig, logger);
         dg.registerFunctions(hostEnv.toHostFunctions());
 
         // Register the WASI host functions.

core/src/main/java/org/extism/sdk/chicory/Manifest.java

@@ -1,12 +1,15 @@
 package org.extism.sdk.chicory;
 
+import com.dylibso.chicory.runtime.Instance;
+import com.dylibso.chicory.runtime.Machine;
 import com.dylibso.chicory.wasi.WasiOptions;
 import org.extism.sdk.chicory.http.HttpConfig;
 
 import com.dylibso.chicory.wasm.types.MemoryLimits;
 import java.util.EnumSet;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 
 public class Manifest {
 
@@ -16,6 +19,7 @@ public enum Validation {
 
     public static class Options {
         boolean aot = false;
+        Function<Instance, Machine> machineFactory;
         EnumSet<Validation> validationFlags = EnumSet.noneOf(Validation.class);
         ConfigProvider config = ConfigProvider.empty();
         WasiOptions wasiOptions = WasiOptions.builder().build();
@@ -32,6 +36,11 @@ public Options withAoT(boolean enabled) {
             return this;
         }
 
+        public Options withMachineFactory(Function<Instance, Machine> machineFactory) {
+            this.machineFactory = machineFactory;
+            return this;
+        }
+
         public Options withConfig(Map<String, String> config) {
             return withConfigProvider(ConfigProvider.ofMap(config));
         }