diff --git a/core/src/main/java/org/extism/sdk/chicory/ChicoryModule.java b/core/src/main/java/org/extism/sdk/chicory/ChicoryModule.java index 57d0596..ab8ed11 100644 --- a/core/src/main/java/org/extism/sdk/chicory/ChicoryModule.java +++ b/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 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); } diff --git a/core/src/main/java/org/extism/sdk/chicory/DependencyGraph.java b/core/src/main/java/org/extism/sdk/chicory/DependencyGraph.java index 82913fc..e3540b6 100644 --- a/core/src/main/java/org/extism/sdk/chicory/DependencyGraph.java +++ b/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 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 moreHostFunctio Instance instance = ChicoryModule.instanceWithOptions( - Instance.builder(m), this.options, aotMachineFactory) + Instance.builder(m), this.options, machineFactory) .withImportValues(importValues) .withStart(false) .build(); diff --git a/core/src/main/java/org/extism/sdk/chicory/Kernel.java b/core/src/main/java/org/extism/sdk/chicory/Kernel.java index 5337e4b..a422981 100644 --- a/core/src/main/java/org/extism/sdk/chicory/Kernel.java +++ b/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 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 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(); } diff --git a/core/src/main/java/org/extism/sdk/chicory/Linker.java b/core/src/main/java/org/extism/sdk/chicory/Linker.java index e367666..6224d1e 100644 --- a/core/src/main/java/org/extism/sdk/chicory/Linker.java +++ b/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 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. diff --git a/core/src/main/java/org/extism/sdk/chicory/Manifest.java b/core/src/main/java/org/extism/sdk/chicory/Manifest.java index 8d1b2e3..4172f46 100644 --- a/core/src/main/java/org/extism/sdk/chicory/Manifest.java +++ b/core/src/main/java/org/extism/sdk/chicory/Manifest.java @@ -1,5 +1,7 @@ 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; @@ -7,6 +9,7 @@ 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 machineFactory; EnumSet 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 machineFactory) { + this.machineFactory = machineFactory; + return this; + } + public Options withConfig(Map config) { return withConfigProvider(ConfigProvider.ofMap(config)); }