Skip to content

Commit 4fc978c

Browse files
evacchievacchi
authored
feat: add machine factory support (#56)
Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>
1 parent b15b25b commit 4fc978c

5 files changed

Lines changed: 52 additions & 21 deletions

File tree

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,12 @@
22

33
import com.dylibso.chicory.runtime.ByteArrayMemory;
44
import com.dylibso.chicory.runtime.Instance;
5+
import com.dylibso.chicory.runtime.Machine;
56
import com.dylibso.chicory.wasm.Parser;
67
import com.dylibso.chicory.wasm.WasmModule;
78

89
import java.nio.file.Path;
10+
import java.util.function.Function;
911

1012
class ChicoryModule {
1113

@@ -29,11 +31,11 @@ static WasmModule fromWasm(ManifestWasm m) {
2931
}
3032
}
3133

32-
static Instance.Builder instanceWithOptions(Instance.Builder m, Manifest.Options opts, CachedAotMachineFactory aotMachineFactory) {
34+
static Instance.Builder instanceWithOptions(Instance.Builder m, Manifest.Options opts, Function<Instance, Machine> aotMachineFactory) {
3335
if (opts == null) {
3436
return m;
3537
}
36-
// This feature is not compatibly with the native-image builder.
38+
// This feature is not compatible with the native-image builder.
3739
if (opts.aot && !IS_NATIVE_IMAGE_AOT) {
3840
m.withMachineFactory(aotMachineFactory);
3941
}

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

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import com.dylibso.chicory.runtime.ImportFunction;
77
import com.dylibso.chicory.runtime.ImportValues;
88
import com.dylibso.chicory.runtime.Instance;
9+
import com.dylibso.chicory.runtime.Machine;
910
import com.dylibso.chicory.runtime.Store;
1011
import com.dylibso.chicory.runtime.WasmFunctionHandle;
1112
import com.dylibso.chicory.wasi.WasiExitException;
@@ -26,6 +27,7 @@
2627
import java.util.Objects;
2728
import java.util.Set;
2829
import java.util.Stack;
30+
import java.util.function.Function;
2931

3032
import static java.util.stream.Collectors.groupingBy;
3133

@@ -42,7 +44,7 @@ class DependencyGraph {
4244

4345
private final Store store = new Store();
4446
private Manifest.Options options;
45-
private CachedAotMachineFactory aotMachineFactory;
47+
private Function<Instance, Machine> machineFactory;
4648

4749
public DependencyGraph(Logger logger) {
4850
this.logger = logger;
@@ -53,8 +55,14 @@ public DependencyGraph(Logger logger) {
5355
*/
5456
public void setOptions(Manifest.Options options) {
5557
this.options = options;
56-
if (options != null && options.aot) {
57-
this.aotMachineFactory = new CachedAotMachineFactory();
58+
if (options != null) {
59+
if (options.aot && options.machineFactory == null) {
60+
this.machineFactory = new CachedAotMachineFactory();
61+
return;
62+
}
63+
if (options.machineFactory != null) {
64+
this.machineFactory = options.machineFactory;
65+
}
5866
}
5967
}
6068

@@ -98,10 +106,12 @@ private void checkCollision(String moduleName, String symbol) {
98106
*/
99107
public void registerModule(String name, WasmModule m) {
100108
checkCollision(name, null);
101-
if (aotMachineFactory != null) {
102-
aotMachineFactory.compile(m);
109+
if (machineFactory != null) {
110+
if (machineFactory instanceof CachedAotMachineFactory) {
111+
var cachedAotMachineFactory = (CachedAotMachineFactory) machineFactory;
112+
cachedAotMachineFactory.compile(m);
113+
}
103114
}
104-
105115
ExportSection exportSection = m.exportSection();
106116
for (int i = 0; i < exportSection.exportCount(); i++) {
107117
Export export = exportSection.getExport(i);
@@ -238,7 +248,7 @@ private Instance instantiate(String moduleId, List<HostFunction> moreHostFunctio
238248

239249
Instance instance =
240250
ChicoryModule.instanceWithOptions(
241-
Instance.builder(m), this.options, aotMachineFactory)
251+
Instance.builder(m), this.options, machineFactory)
242252
.withImportValues(importValues)
243253
.withStart(false)
244254
.build();

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,14 @@
33
import com.dylibso.chicory.runtime.ExportFunction;
44
import com.dylibso.chicory.runtime.HostFunction;
55
import com.dylibso.chicory.runtime.Instance;
6+
import com.dylibso.chicory.runtime.Machine;
67
import com.dylibso.chicory.wasm.Parser;
78
import com.dylibso.chicory.wasm.WasmModule;
89
import com.dylibso.chicory.wasm.types.FunctionType;
910
import com.dylibso.chicory.wasm.types.ValType;
1011

1112
import java.util.List;
13+
import java.util.function.Function;
1214

1315
public class Kernel {
1416

@@ -39,7 +41,7 @@ public Kernel() {
3941
this(null);
4042
}
4143

42-
Kernel(CachedAotMachineFactory machineFactory) {
44+
Kernel(Function<Instance, Machine> machineFactory) {
4345
Instance kernel = instance(machineFactory);
4446
instanceMemory = kernel.memory();
4547
alloc = kernel.export("alloc");
@@ -64,11 +66,11 @@ public Kernel() {
6466
memoryBytes = kernel.export("memory_bytes");
6567
}
6668

67-
private static Instance instance(CachedAotMachineFactory machineFactory) {
69+
private static Instance instance(Function<Instance, Machine> machineFactory) {
6870
var kernelStream = Kernel.class.getClassLoader().getResourceAsStream("extism-runtime.wasm");
6971
WasmModule module = Parser.parse(kernelStream);
70-
if (machineFactory != null) {
71-
machineFactory.compile(module);
72+
if (machineFactory != null && machineFactory instanceof CachedAotMachineFactory) {
73+
((CachedAotMachineFactory) machineFactory).compile(module);
7274
}
7375
return Instance.builder(module).withMachineFactory(machineFactory).build();
7476
}

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

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@
33
import com.dylibso.chicory.log.Logger;
44
import com.dylibso.chicory.runtime.HostFunction;
55
import com.dylibso.chicory.runtime.Instance;
6+
import com.dylibso.chicory.runtime.Machine;
67
import com.dylibso.chicory.wasi.WasiOptions;
78
import com.dylibso.chicory.wasi.WasiPreview1;
89
import org.extism.sdk.chicory.http.HttpConfig;
910

1011
import java.util.Arrays;
12+
import java.util.function.Function;
1113

1214

1315
/**
@@ -37,17 +39,23 @@ Plugin link() {
3739
ConfigProvider config;
3840
String[] allowedHosts;
3941
WasiOptions wasiOptions;
40-
CachedAotMachineFactory aotMachineFactory;
42+
Function<Instance, Machine> machineFactory = null;
4143
HttpConfig httpConfig;
42-
dg.setOptions(manifest.options);
43-
config = manifest.options.config;
44-
allowedHosts = manifest.options.allowedHosts;
45-
wasiOptions = manifest.options.wasiOptions;
46-
aotMachineFactory = manifest.options.aot? new CachedAotMachineFactory() : null;
47-
httpConfig = manifest.options.httpConfig;
44+
Manifest.Options options = manifest.options;
45+
dg.setOptions(options);
46+
config = options.config;
47+
allowedHosts = options.allowedHosts;
48+
wasiOptions = options.wasiOptions;
49+
if (options.aot && options.machineFactory == null) {
50+
machineFactory = new CachedAotMachineFactory();
51+
}
52+
if (options.machineFactory != null) {
53+
machineFactory = options.machineFactory;
54+
}
55+
httpConfig = options.httpConfig;
4856

4957
// Register the HostEnv exports.
50-
var hostEnv = new HostEnv(new Kernel(aotMachineFactory), config, allowedHosts, httpConfig, logger);
58+
var hostEnv = new HostEnv(new Kernel(machineFactory), config, allowedHosts, httpConfig, logger);
5159
dg.registerFunctions(hostEnv.toHostFunctions());
5260

5361
// Register the WASI host functions.

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
package org.extism.sdk.chicory;
22

3+
import com.dylibso.chicory.runtime.Instance;
4+
import com.dylibso.chicory.runtime.Machine;
35
import com.dylibso.chicory.wasi.WasiOptions;
46
import org.extism.sdk.chicory.http.HttpConfig;
57

68
import com.dylibso.chicory.wasm.types.MemoryLimits;
79
import java.util.EnumSet;
810
import java.util.List;
911
import java.util.Map;
12+
import java.util.function.Function;
1013

1114
public class Manifest {
1215

@@ -16,6 +19,7 @@ public enum Validation {
1619

1720
public static class Options {
1821
boolean aot = false;
22+
Function<Instance, Machine> machineFactory;
1923
EnumSet<Validation> validationFlags = EnumSet.noneOf(Validation.class);
2024
ConfigProvider config = ConfigProvider.empty();
2125
WasiOptions wasiOptions = WasiOptions.builder().build();
@@ -32,6 +36,11 @@ public Options withAoT(boolean enabled) {
3236
return this;
3337
}
3438

39+
public Options withMachineFactory(Function<Instance, Machine> machineFactory) {
40+
this.machineFactory = machineFactory;
41+
return this;
42+
}
43+
3544
public Options withConfig(Map<String, String> config) {
3645
return withConfigProvider(ConfigProvider.ofMap(config));
3746
}

0 commit comments

Comments
 (0)