From a41d8a7ab45c8d5e0fcda8e7bbde809c449093cf Mon Sep 17 00:00:00 2001 From: Andreas Karlsson Date: Fri, 23 Oct 2026 14:27:28 +0200 Subject: [PATCH 1/2] feat: compileModule() without the bodies of compiled functions --- .../run/endive/testing/CompileModuleTest.java | 110 ++++++++++++++++++ .../compiler/MachineFactoryCompiler.java | 60 +++++++++- .../java/run/endive/runtime/Instance.java | 21 +++- .../internal/CompilerInterpreterMachine.java | 2 +- .../compiled/tail_call_drop_bodies.wat.wasm | Bin 0 -> 110 bytes .../resources/wat/tail_call_drop_bodies.wat | 12 ++ .../main/java/run/endive/wasm/WasmModule.java | 36 ++++++ .../run/endive/wasm/types/CodeSection.java | 31 ++++- 8 files changed, 266 insertions(+), 6 deletions(-) create mode 100644 compiler-tests/src/test/java/run/endive/testing/CompileModuleTest.java create mode 100644 wasm-corpus/src/main/resources/compiled/tail_call_drop_bodies.wat.wasm create mode 100644 wasm-corpus/src/main/resources/wat/tail_call_drop_bodies.wat diff --git a/compiler-tests/src/test/java/run/endive/testing/CompileModuleTest.java b/compiler-tests/src/test/java/run/endive/testing/CompileModuleTest.java new file mode 100644 index 000000000..627f8faec --- /dev/null +++ b/compiler-tests/src/test/java/run/endive/testing/CompileModuleTest.java @@ -0,0 +1,110 @@ +package run.endive.testing; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static run.endive.corpus.WatGenerator.methodTooLarge; + +import java.util.List; +import java.util.Set; +import org.junit.jupiter.api.Test; +import run.endive.compiler.InterpreterFallback; +import run.endive.compiler.MachineFactoryCompiler; +import run.endive.corpus.CorpusResources; +import run.endive.runtime.CompiledModule; +import run.endive.runtime.HostFunction; +import run.endive.runtime.ImportValues; +import run.endive.runtime.Instance; +import run.endive.wabt.Wat2Wasm; +import run.endive.wasm.Parser; +import run.endive.wasm.WasmModule; +import run.endive.wasm.types.CodeSection; +import run.endive.wasm.types.FunctionType; +import run.endive.wasm.types.ValType; + +public class CompileModuleTest { + + // func_1 (function index 2, code index 1) is too large to compile and is interpreted. + private static final WasmModule METHOD_TOO_LARGE = + Parser.parse(Wat2Wasm.parse(methodTooLarge(20_000))); + + private static final HostFunction HOST_FUNC = + new HostFunction( + "funcs", + "host_func", + FunctionType.of(List.of(ValType.I32), List.of(ValType.I32)), + (inst, args) -> new long[] {35}); + + private static CompiledModule compileModule(WasmModule module) { + return MachineFactoryCompiler.builder(module) + .withInterpreterFallback(InterpreterFallback.SILENT) + .compileModule(); + } + + private static Instance instantiate(CompiledModule compiled) { + return Instance.builder(compiled) + .withImportValues(ImportValues.builder().addFunction(HOST_FUNC).build()) + .withStart(false) + .build(); + } + + private static void assertRetained(CodeSection code, Set retained) { + var bodies = code.functionBodies(); + for (int i = 0; i < bodies.length; i++) { + assertEquals(retained.contains(i), bodies[i] != null, "function body " + i); + } + } + + @Test + public void keepsOnlyInterpretedFunctionBodies() { + var compiled = compileModule(METHOD_TOO_LARGE); + + var code = compiled.wasmModule().codeSection(); + assertEquals(3, code.functionBodyCount()); + assertRetained(code, Set.of(1)); + assertThrows(IllegalStateException.class, () -> code.getFunctionBody(0)); + // The module passed to the compiler is left unchanged. + assertRetained(METHOD_TOO_LARGE.codeSection(), Set.of(0, 1, 2)); + + var instance = instantiate(compiled); + // Through the interpreted func_1, calling back into compiled code directly... + assertEquals(35, instance.export("func_2").apply(0)[0]); + // ...and through call_indirect. + assertEquals(35, instance.export("func_2").apply(1)[0]); + } + + @Test + public void keepsFunctionsReachableThroughTailCalls() { + // interpreted tail calls tail_called, which calls compiled + var module = + Parser.parse( + CorpusResources.getResource("compiled/tail_call_drop_bodies.wat.wasm")); + var compiled = + MachineFactoryCompiler.builder(module) + .withInterpreterFallback(InterpreterFallback.FAIL) + .withInterpretedFunctions(Set.of(0)) + .compileModule(); + + assertRetained(compiled.wasmModule().codeSection(), Set.of(0, 1)); + assertEquals(42, Instance.builder(compiled).build().export("interpreted").apply(41)[0]); + } + + @Test + public void compileKeepsTheModule() { + var factory = + MachineFactoryCompiler.builder(METHOD_TOO_LARGE) + .withInterpreterFallback(InterpreterFallback.SILENT) + .compile(); + + assertSame(METHOD_TOO_LARGE, ((CompiledModule) factory).wasmModule()); + } + + @Test + public void instanceRejectsDroppedBody() { + var instance = instantiate(compileModule(METHOD_TOO_LARGE)); + + var exception = assertThrows(IllegalStateException.class, () -> instance.function(1)); + assertTrue(exception.getMessage().contains("dropped")); + } +} diff --git a/compiler/src/main/java/run/endive/compiler/MachineFactoryCompiler.java b/compiler/src/main/java/run/endive/compiler/MachineFactoryCompiler.java index f4a1e543b..5c841fc1d 100644 --- a/compiler/src/main/java/run/endive/compiler/MachineFactoryCompiler.java +++ b/compiler/src/main/java/run/endive/compiler/MachineFactoryCompiler.java @@ -4,7 +4,9 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.UncheckedIOException; +import java.util.ArrayDeque; import java.util.HashMap; +import java.util.HashSet; import java.util.Map; import java.util.Properties; import java.util.Set; @@ -14,10 +16,13 @@ import java.util.jar.JarOutputStream; import run.endive.compiler.internal.ClassLoadingCollector; import run.endive.compiler.internal.MachineFactory; +import run.endive.runtime.CompiledModule; import run.endive.runtime.Instance; import run.endive.runtime.Machine; import run.endive.wasm.WasmEngineException; import run.endive.wasm.WasmModule; +import run.endive.wasm.types.ExternalType; +import run.endive.wasm.types.OpCode; /** * Compiles WASM function bodies to JVM byte code that can be used as a machine factory for {@link Instance}'s. @@ -127,6 +132,20 @@ public Builder withCache(Cache cache) { } public Function compile() { + return compileFactory(false); + } + + /** + * Compiles the module and returns it together with a machine factory for instances of it, + * see {@link Instance#builder(CompiledModule)}. The returned module is a copy without the + * bodies of compiled functions, unless the compilation is loaded from the cache. The + * module passed to the builder is left unchanged. + */ + public CompiledModule compileModule() { + return compileFactory(true); + } + + private MachineFactory compileFactory(boolean dropCompiledFunctionBodies) { try { // Can we load the byte codes from the cache? @@ -152,13 +171,52 @@ public Function compile() { cache.putIfAbsent(module.digest(), storeClassLoadingCollector(collector)); } - return new MachineFactory(module, collector.machineFactory()); + var compiledModule = + dropCompiledFunctionBodies + ? withoutCompiledFunctionBodies( + module, result.interpretedFunctions()) + : module; + return new MachineFactory(compiledModule, collector.machineFactory()); } catch (IOException e) { throw new WasmEngineException(e); } } } + private static WasmModule withoutCompiledFunctionBodies( + WasmModule module, Set interpretedFunctions) { + var retained = interpretableFunctions(module, interpretedFunctions); + return retained == null ? module : module.retainingFunctionBodies(retained); + } + + /** + * Returns the functions the interpreter may run: the interpreted functions and, since the + * interpreter runs tail calls itself, every function they reach through tail calls. Returns + * {@code null} if an indirect tail call makes the targets unknown. + */ + private static Set interpretableFunctions( + WasmModule module, Set interpretedFunctions) { + int functionImports = module.importSection().count(ExternalType.FUNCTION); + var result = new HashSet(); + var pending = new ArrayDeque<>(interpretedFunctions); + while (!pending.isEmpty()) { + int funcId = pending.pop(); + if (funcId < functionImports || !result.add(funcId)) { + continue; + } + var body = module.codeSection().getFunctionBody(funcId - functionImports); + for (var ins : body.instructions()) { + if (ins.opcode() == OpCode.RETURN_CALL) { + pending.push((int) ins.operand(0)); + } else if (ins.opcode() == OpCode.RETURN_CALL_INDIRECT + || ins.opcode() == OpCode.RETURN_CALL_REF) { + return null; + } + } + } + return result; + } + private static byte[] storeClassLoadingCollector(ClassLoadingCollector collector) { try { // Create JAR in memory diff --git a/runtime/src/main/java/run/endive/runtime/Instance.java b/runtime/src/main/java/run/endive/runtime/Instance.java index 45004fa6a..ea93cc1cb 100644 --- a/runtime/src/main/java/run/endive/runtime/Instance.java +++ b/runtime/src/main/java/run/endive/runtime/Instance.java @@ -346,10 +346,19 @@ public ExportFunction export(String name) { public FunctionBody function(long idx) { if (idx < 0 || idx >= (functions.length + imports.functionCount())) { throw new InvalidException("unknown function " + idx); - } else if (idx < imports.functionCount()) { + } else if (isImportedFunction(idx)) { return null; } - return functions[(int) idx - imports.functionCount()]; + var body = functions[(int) idx - imports.functionCount()]; + if (body == null) { + throw new IllegalStateException( + "Function body " + idx + " has been dropped from the module"); + } + return body; + } + + public boolean isImportedFunction(long idx) { + return idx < imports.functionCount(); } public int functionCount() { @@ -617,6 +626,14 @@ public static Builder builder(WasmModule module) { return new Builder(module); } + /** + * Creates a builder for instances of a compiled module, using its module and machine factory. + */ + public static Builder builder(CompiledModule compiledModule) { + return new Builder(compiledModule.wasmModule()) + .withMachineFactory(compiledModule.machineFactory()); + } + public static final class Builder { private final WasmModule module; diff --git a/runtime/src/main/java/run/endive/runtime/internal/CompilerInterpreterMachine.java b/runtime/src/main/java/run/endive/runtime/internal/CompilerInterpreterMachine.java index 805d3bac6..643bb2f87 100644 --- a/runtime/src/main/java/run/endive/runtime/internal/CompilerInterpreterMachine.java +++ b/runtime/src/main/java/run/endive/runtime/internal/CompilerInterpreterMachine.java @@ -60,7 +60,7 @@ protected long[] call( protected void CALL(Operands operands) { var instance = instance(); var funcId = (int) operands.get(0); - if (interpretedFuncIds.contains(funcId) || instance.function(funcId) == null) { + if (interpretedFuncIds.contains(funcId) || instance.isImportedFunction(funcId)) { // continue interpreting for interpreted functions or imported functions super.CALL(operands); } else { diff --git a/wasm-corpus/src/main/resources/compiled/tail_call_drop_bodies.wat.wasm b/wasm-corpus/src/main/resources/compiled/tail_call_drop_bodies.wat.wasm new file mode 100644 index 0000000000000000000000000000000000000000..2ecb0e7af6ca00318c7903ecb3083edc518e5642 GIT binary patch literal 110 zcmZQbEY4+QU|?WmV@zPIXRK#tVFvQp`5C!0^GZ^S3W`!oQd1ZhxWt*+7!(+U7`Z`| x024PmkaA?q;%3lc$xF;lWmIKmK-R&?U6PoY6Q7)zlardl#F3nzTaXE40RR?$7lHr) literal 0 HcmV?d00001 diff --git a/wasm-corpus/src/main/resources/wat/tail_call_drop_bodies.wat b/wasm-corpus/src/main/resources/wat/tail_call_drop_bodies.wat new file mode 100644 index 000000000..b6b3f184b --- /dev/null +++ b/wasm-corpus/src/main/resources/wat/tail_call_drop_bodies.wat @@ -0,0 +1,12 @@ +(module + ;; An interpreted function reaching another one through a tail call, which calls a third. + (func $interpreted (export "interpreted") (param i32) (result i32) + local.get 0 + return_call $tail_called) + (func $tail_called (param i32) (result i32) + local.get 0 + call $compiled) + (func $compiled (param i32) (result i32) + local.get 0 + i32.const 1 + i32.add)) diff --git a/wasm/src/main/java/run/endive/wasm/WasmModule.java b/wasm/src/main/java/run/endive/wasm/WasmModule.java index ed732b320..131e9271b 100644 --- a/wasm/src/main/java/run/endive/wasm/WasmModule.java +++ b/wasm/src/main/java/run/endive/wasm/WasmModule.java @@ -8,12 +8,15 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; import run.endive.wasm.types.CodeSection; import run.endive.wasm.types.CustomSection; import run.endive.wasm.types.DataCountSection; import run.endive.wasm.types.DataSection; import run.endive.wasm.types.ElementSection; import run.endive.wasm.types.ExportSection; +import run.endive.wasm.types.ExternalType; import run.endive.wasm.types.FunctionSection; import run.endive.wasm.types.GlobalSection; import run.endive.wasm.types.ImportSection; @@ -102,6 +105,39 @@ public CodeSection codeSection() { return codeSection; } + /** + * Returns a copy of this module that keeps only the bodies of the functions in {@code + * funcIds}, sharing every other section. Accessing a dropped body throws {@link + * IllegalStateException}. + * + * @param funcIds function indices in the module's function index space (imports included) + */ + public WasmModule retainingFunctionBodies(Set funcIds) { + int functionImports = importSection.count(ExternalType.FUNCTION); + var indices = + funcIds.stream() + .filter(funcId -> funcId >= functionImports) + .map(funcId -> funcId - functionImports) + .collect(Collectors.toSet()); + return new WasmModule( + typeSection, + importSection, + functionSection, + tableSection, + memorySection, + globalSection, + exportSection, + startSection, + elementSection, + codeSection.retainingFunctionBodies(indices), + dataSection, + dataCountSection, + tagSection, + customSections, + ignoredSections, + digest); + } + public DataSection dataSection() { return dataSection; } diff --git a/wasm/src/main/java/run/endive/wasm/types/CodeSection.java b/wasm/src/main/java/run/endive/wasm/types/CodeSection.java index 41128694a..c35f842a7 100644 --- a/wasm/src/main/java/run/endive/wasm/types/CodeSection.java +++ b/wasm/src/main/java/run/endive/wasm/types/CodeSection.java @@ -1,19 +1,24 @@ package run.endive.wasm.types; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Objects; +import java.util.Set; public final class CodeSection extends Section { + // A null entry is a body dropped by retainingFunctionBodies. private final List functionBodies; private final boolean requiresDataCount; private CodeSection(List functionBodies, boolean requiresDataCount) { super(SectionId.CODE); - this.functionBodies = List.copyOf(functionBodies); + // Not List.copyOf, which rejects the nulls of dropped bodies. + this.functionBodies = Collections.unmodifiableList(new ArrayList<>(functionBodies)); this.requiresDataCount = requiresDataCount; } + /** Returns all function bodies, with {@code null} for dropped ones. */ public FunctionBody[] functionBodies() { return functionBodies.toArray(new FunctionBody[0]); } @@ -22,8 +27,30 @@ public int functionBodyCount() { return functionBodies.size(); } + /** + * @throws IllegalStateException if the function body has been dropped + */ public FunctionBody getFunctionBody(int idx) { - return functionBodies.get(idx); + var body = functionBodies.get(idx); + if (body == null) { + throw new IllegalStateException("Function body " + idx + " has been dropped"); + } + return body; + } + + /** + * Returns a copy keeping only the function bodies at {@code indices}, see {@link + * run.endive.wasm.WasmModule#retainingFunctionBodies(Set)}. + * + * @param indices the code section indices of the function bodies to keep + */ + public CodeSection retainingFunctionBodies(Set indices) { + var retained = + new ArrayList(Collections.nCopies(functionBodies.size(), null)); + for (int idx : indices) { + retained.set(idx, functionBodies.get(idx)); + } + return new CodeSection(retained, requiresDataCount); } public boolean isRequiresDataCount() { From d484a871c2a56c1ce189a13efec8869a49b7598d Mon Sep 17 00:00:00 2001 From: andreatp Date: Mon, 26 Oct 2026 18:12:26 +0000 Subject: [PATCH 2/2] Share body stripping with the .meta path and route call_ref to compiled code --- .../endive/build/time/compiler/Generator.java | 4 +- .../run/endive/testing/CompileModuleTest.java | 139 ++++++++++++------ .../testing/InterpreterFallbackTest.java | 50 +++++++ .../compiler/MachineFactoryCompiler.java | 94 ++++++------ .../internal/InterpretableFunctions.java | 50 +++++++ docs/docs/execution/runtime-compiler.md | 17 +++ .../run/endive/runtime/CompiledModule.java | 2 +- .../java/run/endive/runtime/Instance.java | 11 +- .../endive/runtime/InterpreterMachine.java | 16 +- .../internal/CompilerInterpreterMachine.java | 2 +- .../resources/compiled/drop_bodies.wat.wasm | Bin 0 -> 175 bytes .../compiled/tail_call_drop_bodies.wat.wasm | Bin 110 -> 0 bytes .../src/main/resources/wat/drop_bodies.wat | 25 ++++ .../resources/wat/tail_call_drop_bodies.wat | 12 -- .../main/java/run/endive/wasm/WasmModule.java | 8 +- .../run/endive/wasm/types/CodeSection.java | 2 +- 16 files changed, 298 insertions(+), 134 deletions(-) create mode 100644 compiler/src/main/java/run/endive/compiler/internal/InterpretableFunctions.java create mode 100644 wasm-corpus/src/main/resources/compiled/drop_bodies.wat.wasm delete mode 100644 wasm-corpus/src/main/resources/compiled/tail_call_drop_bodies.wat.wasm create mode 100644 wasm-corpus/src/main/resources/wat/drop_bodies.wat delete mode 100644 wasm-corpus/src/main/resources/wat/tail_call_drop_bodies.wat diff --git a/build-time-compiler/src/main/java/run/endive/build/time/compiler/Generator.java b/build-time-compiler/src/main/java/run/endive/build/time/compiler/Generator.java index f3f8e27f6..d4b6af23f 100644 --- a/build-time-compiler/src/main/java/run/endive/build/time/compiler/Generator.java +++ b/build-time-compiler/src/main/java/run/endive/build/time/compiler/Generator.java @@ -38,6 +38,7 @@ import run.endive.compiler.MethodPrefixer; import run.endive.compiler.internal.ByteClassCollector; import run.endive.compiler.internal.Compiler; +import run.endive.compiler.internal.InterpretableFunctions; import run.endive.runtime.CompiledModule; import run.endive.runtime.Instance; import run.endive.runtime.Machine; @@ -169,6 +170,7 @@ public void generateModuleInterface(String moduleInterfaceName) throws IOExcepti public void generateMetaWasm(Set interpretedFunctions) throws IOException { byte[] wasmBytes = Files.readAllBytes(config.wasmFile()); var module = Parser.parse(wasmBytes); + var retained = InterpretableFunctions.of(module, interpretedFunctions); var writer = new WasmWriter(); Parser.parseWithoutDecoding( @@ -188,7 +190,7 @@ public void generateMetaWasm(Set interpretedFunctions) throws IOExcepti } for (int i = 0; i < count; i++) { var funcId = importFuncs + i; - if (interpretedFunctions.contains(funcId)) { + if (retained.contains(funcId)) { // Copy over the original function body from the source var bodySize = (int) readVarUInt32(source); diff --git a/compiler-tests/src/test/java/run/endive/testing/CompileModuleTest.java b/compiler-tests/src/test/java/run/endive/testing/CompileModuleTest.java index 627f8faec..83779c70c 100644 --- a/compiler-tests/src/test/java/run/endive/testing/CompileModuleTest.java +++ b/compiler-tests/src/test/java/run/endive/testing/CompileModuleTest.java @@ -1,14 +1,18 @@ package run.endive.testing; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static run.endive.corpus.WatGenerator.methodTooLarge; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Set; import org.junit.jupiter.api.Test; +import run.endive.compiler.Cache; import run.endive.compiler.InterpreterFallback; import run.endive.compiler.MachineFactoryCompiler; import run.endive.corpus.CorpusResources; @@ -25,28 +29,19 @@ public class CompileModuleTest { - // func_1 (function index 2, code index 1) is too large to compile and is interpreted. - private static final WasmModule METHOD_TOO_LARGE = - Parser.parse(Wat2Wasm.parse(methodTooLarge(20_000))); - - private static final HostFunction HOST_FUNC = - new HostFunction( - "funcs", - "host_func", - FunctionType.of(List.of(ValType.I32), List.of(ValType.I32)), - (inst, args) -> new long[] {35}); - - private static CompiledModule compileModule(WasmModule module) { - return MachineFactoryCompiler.builder(module) - .withInterpreterFallback(InterpreterFallback.SILENT) - .compileModule(); + // Each export reaches function 0 differently when it is the only interpreted function. + private static final WasmModule DROP_BODIES = + Parser.parse(CorpusResources.getResource("compiled/drop_bodies.wat.wasm")); + private static final int TAIL_CALL = 3; + private static final int CALL_REF = 4; + private static final int TAIL_CALL_REF = 5; + + private static MachineFactoryCompiler.Builder interpreting(int funcId) { + return MachineFactoryCompiler.builder(DROP_BODIES).withInterpretedFunctions(Set.of(funcId)); } - private static Instance instantiate(CompiledModule compiled) { - return Instance.builder(compiled) - .withImportValues(ImportValues.builder().addFunction(HOST_FUNC).build()) - .withStart(false) - .build(); + private static long call(CompiledModule compiled, String export) { + return Instance.builder(compiled).build().export(export).apply(41)[0]; } private static void assertRetained(CodeSection code, Set retained) { @@ -58,16 +53,31 @@ private static void assertRetained(CodeSection code, Set retained) { @Test public void keepsOnlyInterpretedFunctionBodies() { - var compiled = compileModule(METHOD_TOO_LARGE); + // func_1 (function index 2, code index 1) is too large to compile and is interpreted. + var module = Parser.parse(Wat2Wasm.parse(methodTooLarge(20_000))); + var compiled = + MachineFactoryCompiler.builder(module) + .withInterpreterFallback(InterpreterFallback.SILENT) + .compileModule(); var code = compiled.wasmModule().codeSection(); assertEquals(3, code.functionBodyCount()); assertRetained(code, Set.of(1)); assertThrows(IllegalStateException.class, () -> code.getFunctionBody(0)); // The module passed to the compiler is left unchanged. - assertRetained(METHOD_TOO_LARGE.codeSection(), Set.of(0, 1, 2)); - - var instance = instantiate(compiled); + assertRetained(module.codeSection(), Set.of(0, 1, 2)); + + var hostFunc = + new HostFunction( + "funcs", + "host_func", + FunctionType.of(List.of(ValType.I32), List.of(ValType.I32)), + (inst, args) -> new long[] {35}); + var instance = + Instance.builder(compiled) + .withImportValues(ImportValues.builder().addFunction(hostFunc).build()) + .withStart(false) + .build(); // Through the interpreted func_1, calling back into compiled code directly... assertEquals(35, instance.export("func_2").apply(0)[0]); // ...and through call_indirect. @@ -76,35 +86,80 @@ public void keepsOnlyInterpretedFunctionBodies() { @Test public void keepsFunctionsReachableThroughTailCalls() { - // interpreted tail calls tail_called, which calls compiled - var module = - Parser.parse( - CorpusResources.getResource("compiled/tail_call_drop_bodies.wat.wasm")); - var compiled = - MachineFactoryCompiler.builder(module) - .withInterpreterFallback(InterpreterFallback.FAIL) - .withInterpretedFunctions(Set.of(0)) - .compileModule(); + var compiled = interpreting(TAIL_CALL).compileModule(); + + assertRetained(compiled.wasmModule().codeSection(), Set.of(1, 2, TAIL_CALL)); + assertEquals(42, call(compiled, "tail_call")); + } + + @Test + public void callRefRunsCompiledFunction() { + var compiled = interpreting(CALL_REF).compileModule(); + + assertRetained(compiled.wasmModule().codeSection(), Set.of(CALL_REF)); + assertEquals(42, call(compiled, "call_ref")); + } + + @Test + public void keepsAllBodiesOnIndirectTailCall() { + var compiled = interpreting(TAIL_CALL_REF).compileModule(); - assertRetained(compiled.wasmModule().codeSection(), Set.of(0, 1)); - assertEquals(42, Instance.builder(compiled).build().export("interpreted").apply(41)[0]); + assertRetained(compiled.wasmModule().codeSection(), Set.of(0, 1, 2, 3, 4, 5)); + assertEquals(42, call(compiled, "tail_call_ref")); + } + + @Test + public void dropsCustomSections() { + assertNotNull(DROP_BODIES.nameSection()); + + var compiled = interpreting(TAIL_CALL).compileModule(); + + assertTrue(compiled.wasmModule().customSections().isEmpty()); + } + + @Test + public void dropsBodiesOnCacheHit() { + var cache = new MemoryCache(); + interpreting(TAIL_CALL).withCache(cache).compileModule(); + + var compiled = interpreting(TAIL_CALL).withCache(cache).compileModule(); + + assertEquals(1, cache.hits); + assertRetained(compiled.wasmModule().codeSection(), Set.of(1, 2, TAIL_CALL)); + assertEquals(42, call(compiled, "tail_call")); } @Test public void compileKeepsTheModule() { - var factory = - MachineFactoryCompiler.builder(METHOD_TOO_LARGE) - .withInterpreterFallback(InterpreterFallback.SILENT) - .compile(); + var factory = interpreting(CALL_REF).compile(); - assertSame(METHOD_TOO_LARGE, ((CompiledModule) factory).wasmModule()); + assertSame(DROP_BODIES, ((CompiledModule) factory).wasmModule()); } @Test public void instanceRejectsDroppedBody() { - var instance = instantiate(compileModule(METHOD_TOO_LARGE)); + var instance = Instance.builder(interpreting(CALL_REF).compileModule()).build(); - var exception = assertThrows(IllegalStateException.class, () -> instance.function(1)); + var exception = assertThrows(IllegalStateException.class, () -> instance.function(0)); assertTrue(exception.getMessage().contains("dropped")); } + + private static final class MemoryCache implements Cache { + private final Map entries = new HashMap<>(); + private int hits; + + @Override + public byte[] get(String key) { + var data = entries.get(key); + if (data != null) { + hits++; + } + return data; + } + + @Override + public void putIfAbsent(String key, byte[] data) { + entries.putIfAbsent(key, data); + } + } } diff --git a/compiler-tests/src/test/java/run/endive/testing/InterpreterFallbackTest.java b/compiler-tests/src/test/java/run/endive/testing/InterpreterFallbackTest.java index 00e6edbf3..7c604824e 100644 --- a/compiler-tests/src/test/java/run/endive/testing/InterpreterFallbackTest.java +++ b/compiler-tests/src/test/java/run/endive/testing/InterpreterFallbackTest.java @@ -29,6 +29,7 @@ import run.endive.build.time.compiler.Config; import run.endive.build.time.compiler.Generator; import run.endive.compiler.InterpreterFallback; +import run.endive.corpus.CorpusResources; import run.endive.runtime.ByteBufferMemory; import run.endive.runtime.HostFunction; import run.endive.runtime.ImportMemory; @@ -275,6 +276,55 @@ public void testInterpreterFallbackWithMemoryImport() throws Exception { assertEquals(35, instance.export("func_2").apply(0)[0]); } + @Test + public void metaKeepsFunctionsReachableThroughTailCalls() throws Exception { + assertEquals(42, callFromMeta(3, "tail_call")); + } + + @Test + public void metaCallRefRunsCompiledFunction() throws Exception { + assertEquals(42, callFromMeta(4, "call_ref")); + } + + @Test + public void metaKeepsAllBodiesOnIndirectTailCall() throws Exception { + assertEquals(42, callFromMeta(5, "tail_call_ref")); + } + + // Generates drop_bodies.wat with only funcId interpreted and calls export(41) through the .meta + private long callFromMeta(int funcId, String export) throws Exception { + Path dir = Path.of("target", "test-fixtures-meta"); + Files.createDirectories(dir); + var file = dir.resolve("drop_bodies.wasm"); + try (var in = CorpusResources.getResource("compiled/drop_bodies.wat.wasm")) { + Files.write(file, in.readAllBytes()); + } + + var name = "TestMeta" + funcId; + var config = + Config.builder() + .withWasmFile(file) + .withTargetSourceFolder(dir) + .withTargetClassFolder(dir) + .withTargetWasmFolder(dir) + .withName("run.endive.testing." + name) + .withInterpretedFunctions(Set.of(funcId)) + .build(); + generateAll(new Generator(config)); + + var cl = new URLClassLoader(new URL[] {dir.toUri().toURL()}); + var machineClass = cl.loadClass("run.endive.testing." + name + "Machine"); + WasmModule module; + try (InputStream in = machineClass.getResourceAsStream(name + ".meta")) { + module = Parser.parse(in); + } + var instance = + Instance.builder(module) + .withMachineFactory(createMachineFactory(machineClass)) + .build(); + return instance.export(export).apply(41)[0]; + } + private static HostFunction createHostFunc(List hostStackTrace) { return new HostFunction( "funcs", diff --git a/compiler/src/main/java/run/endive/compiler/MachineFactoryCompiler.java b/compiler/src/main/java/run/endive/compiler/MachineFactoryCompiler.java index 5c841fc1d..6ba3c4cfe 100644 --- a/compiler/src/main/java/run/endive/compiler/MachineFactoryCompiler.java +++ b/compiler/src/main/java/run/endive/compiler/MachineFactoryCompiler.java @@ -1,12 +1,14 @@ package run.endive.compiler; +import static java.util.stream.Collectors.joining; +import static java.util.stream.Collectors.toSet; + import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.UncheckedIOException; -import java.util.ArrayDeque; +import java.util.Arrays; import java.util.HashMap; -import java.util.HashSet; import java.util.Map; import java.util.Properties; import java.util.Set; @@ -15,20 +17,22 @@ import java.util.jar.JarInputStream; import java.util.jar.JarOutputStream; import run.endive.compiler.internal.ClassLoadingCollector; +import run.endive.compiler.internal.CompilerResult; +import run.endive.compiler.internal.InterpretableFunctions; import run.endive.compiler.internal.MachineFactory; import run.endive.runtime.CompiledModule; import run.endive.runtime.Instance; import run.endive.runtime.Machine; import run.endive.wasm.WasmEngineException; import run.endive.wasm.WasmModule; -import run.endive.wasm.types.ExternalType; -import run.endive.wasm.types.OpCode; /** * Compiles WASM function bodies to JVM byte code that can be used as a machine factory for {@link Instance}'s. */ public final class MachineFactoryCompiler { + private static final String INTERPRETED_FUNCTIONS = "interpretedFunctions"; + private MachineFactoryCompiler() { // no instances } @@ -137,15 +141,15 @@ public Function compile() { /** * Compiles the module and returns it together with a machine factory for instances of it, - * see {@link Instance#builder(CompiledModule)}. The returned module is a copy without the - * bodies of compiled functions, unless the compilation is loaded from the cache. The - * module passed to the builder is left unchanged. + * see {@link Instance#builder(CompiledModule)}. The returned module is a copy without + * custom sections that keeps only the function bodies the interpreter may run, see {@link + * WasmModule#stripped(Set)}. The module passed to the builder is left unchanged. */ public CompiledModule compileModule() { return compileFactory(true); } - private MachineFactory compileFactory(boolean dropCompiledFunctionBodies) { + private MachineFactory compileFactory(boolean strip) { try { // Can we load the byte codes from the cache? @@ -153,8 +157,7 @@ private MachineFactory compileFactory(boolean dropCompiledFunctionBodies) { if (useCache) { byte[] cachedData = cache.get(module.digest()); if (cachedData != null) { - var collector = loadClassLoadingCollector(cachedData); - return new MachineFactory(module, collector.machineFactory()); + return machineFactory(loadCompilerResult(cachedData), strip); } } @@ -164,60 +167,32 @@ private MachineFactory compileFactory(boolean dropCompiledFunctionBodies) { .withClassCollectorFactory(ClassLoadingCollector::new) .build() .compile(); - var collector = (ClassLoadingCollector) result.collector(); if (useCache) { // store results in the cache to speed the next time. - cache.putIfAbsent(module.digest(), storeClassLoadingCollector(collector)); + cache.putIfAbsent(module.digest(), storeCompilerResult(result)); } - var compiledModule = - dropCompiledFunctionBodies - ? withoutCompiledFunctionBodies( - module, result.interpretedFunctions()) - : module; - return new MachineFactory(compiledModule, collector.machineFactory()); + return machineFactory(result, strip); } catch (IOException e) { throw new WasmEngineException(e); } } - } - - private static WasmModule withoutCompiledFunctionBodies( - WasmModule module, Set interpretedFunctions) { - var retained = interpretableFunctions(module, interpretedFunctions); - return retained == null ? module : module.retainingFunctionBodies(retained); - } - /** - * Returns the functions the interpreter may run: the interpreted functions and, since the - * interpreter runs tail calls itself, every function they reach through tail calls. Returns - * {@code null} if an indirect tail call makes the targets unknown. - */ - private static Set interpretableFunctions( - WasmModule module, Set interpretedFunctions) { - int functionImports = module.importSection().count(ExternalType.FUNCTION); - var result = new HashSet(); - var pending = new ArrayDeque<>(interpretedFunctions); - while (!pending.isEmpty()) { - int funcId = pending.pop(); - if (funcId < functionImports || !result.add(funcId)) { - continue; - } - var body = module.codeSection().getFunctionBody(funcId - functionImports); - for (var ins : body.instructions()) { - if (ins.opcode() == OpCode.RETURN_CALL) { - pending.push((int) ins.operand(0)); - } else if (ins.opcode() == OpCode.RETURN_CALL_INDIRECT - || ins.opcode() == OpCode.RETURN_CALL_REF) { - return null; - } + private MachineFactory machineFactory(CompilerResult result, boolean strip) { + var factory = ((ClassLoadingCollector) result.collector()).machineFactory(); + var interpreted = result.interpretedFunctions(); + // Cache entries written by older versions don't record the interpreted functions. + if (!strip || interpreted == null) { + return new MachineFactory(module, factory); } + var stripped = module.stripped(InterpretableFunctions.of(module, interpreted)); + return new MachineFactory(stripped, factory); } - return result; } - private static byte[] storeClassLoadingCollector(ClassLoadingCollector collector) { + private static byte[] storeCompilerResult(CompilerResult result) { + var collector = (ClassLoadingCollector) result.collector(); try { // Create JAR in memory ByteArrayOutputStream baos = new ByteArrayOutputStream(); @@ -225,6 +200,11 @@ private static byte[] storeClassLoadingCollector(ClassLoadingCollector collector // Store the properties file var properties = new Properties(); properties.put("mainClass", collector.mainClassName()); + properties.put( + INTERPRETED_FUNCTIONS, + result.interpretedFunctions().stream() + .map(String::valueOf) + .collect(joining(","))); ByteArrayOutputStream propsBaos = new ByteArrayOutputStream(); properties.store(propsBaos, ""); @@ -248,9 +228,9 @@ private static byte[] storeClassLoadingCollector(ClassLoadingCollector collector } } - private static ClassLoadingCollector loadClassLoadingCollector(byte[] jarData) - throws IOException { + private static CompilerResult loadCompilerResult(byte[] jarData) throws IOException { var collector = new ClassLoadingCollector(); + Set interpretedFunctions = null; // It was previously compiled, just load it from JAR. try (JarInputStream jis = new JarInputStream(new ByteArrayInputStream(jarData))) { @@ -264,6 +244,14 @@ private static ClassLoadingCollector loadClassLoadingCollector(byte[] jarData) // Load properties properties.load(jis); mainClass = properties.getProperty("mainClass"); + var interpreted = properties.getProperty(INTERPRETED_FUNCTIONS); + if (interpreted != null) { + interpretedFunctions = + Arrays.stream(interpreted.split(",")) + .filter(id -> !id.isEmpty()) + .map(Integer::valueOf) + .collect(toSet()); + } } else if (entry.getName().endsWith(".class")) { // Load class file String className = entry.getName().replace('/', '.').replace(".class", ""); @@ -289,6 +277,6 @@ private static ClassLoadingCollector loadClassLoadingCollector(byte[] jarData) } } } - return collector; + return new CompilerResult(collector, interpretedFunctions); } } diff --git a/compiler/src/main/java/run/endive/compiler/internal/InterpretableFunctions.java b/compiler/src/main/java/run/endive/compiler/internal/InterpretableFunctions.java new file mode 100644 index 000000000..e9cf8feb9 --- /dev/null +++ b/compiler/src/main/java/run/endive/compiler/internal/InterpretableFunctions.java @@ -0,0 +1,50 @@ +package run.endive.compiler.internal; + +import static java.util.stream.Collectors.toSet; + +import java.util.ArrayDeque; +import java.util.HashSet; +import java.util.Set; +import java.util.stream.IntStream; +import run.endive.wasm.WasmModule; +import run.endive.wasm.types.ExternalType; +import run.endive.wasm.types.OpCode; + +/** + * Computes the functions whose bodies the interpreter may run in a compiled module. + */ +public final class InterpretableFunctions { + + private InterpretableFunctions() {} + + /** + * Returns the interpreted functions plus every function they reach through tail calls, or + * every defined function if an indirect tail call makes the targets unknown. + * + * @param interpretedFunctions function indices in the module's function index space + */ + public static Set of(WasmModule module, Set interpretedFunctions) { + int functionImports = module.importSection().count(ExternalType.FUNCTION); + var result = new HashSet(); + var pending = new ArrayDeque<>(interpretedFunctions); + while (!pending.isEmpty()) { + int funcId = pending.pop(); + if (funcId < functionImports || !result.add(funcId)) { + continue; + } + var body = module.codeSection().getFunctionBody(funcId - functionImports); + for (var ins : body.instructions()) { + if (ins.opcode() == OpCode.RETURN_CALL) { + pending.push((int) ins.operand(0)); + } else if (ins.opcode() == OpCode.RETURN_CALL_INDIRECT + || ins.opcode() == OpCode.RETURN_CALL_REF) { + int functions = module.functionSection().functionCount(); + return IntStream.range(functionImports, functionImports + functions) + .boxed() + .collect(toSet()); + } + } + } + return result; + } +} diff --git a/docs/docs/execution/runtime-compiler.md b/docs/docs/execution/runtime-compiler.md index 7e6da5f89..a27b8eef8 100644 --- a/docs/docs/execution/runtime-compiler.md +++ b/docs/docs/execution/runtime-compiler.md @@ -118,6 +118,23 @@ var instance = Instance.builder(module). Typically, you can obtain the list of the functions by running the compiler once with `InterpreterFallback.WARN` +### Releasing Function Bodies + +Once compiled, the parsed function bodies are only needed by functions that run in the interpreter. +`compileModule()` returns the machine factory together with a copy of the module that keeps only +those bodies and no custom sections, so the rest can be garbage collected once the original module +is no longer referenced: + +```java +import run.endive.compiler.MachineFactoryCompiler; +import run.endive.runtime.Instance; +import run.endive.wasm.Parser; +import java.io.File; + +var compiled = MachineFactoryCompiler.builder(Parser.parse(new File("your.wasm"))).compileModule(); +var instance = Instance.builder(compiled).build(); +``` + ### Method Names By default, the compiler names compiled methods `func_0`, `func_1`, etc. A `MethodPrefixer` lets you diff --git a/runtime/src/main/java/run/endive/runtime/CompiledModule.java b/runtime/src/main/java/run/endive/runtime/CompiledModule.java index fadf33f6a..7485fe920 100644 --- a/runtime/src/main/java/run/endive/runtime/CompiledModule.java +++ b/runtime/src/main/java/run/endive/runtime/CompiledModule.java @@ -4,7 +4,7 @@ import run.endive.wasm.WasmModule; /** - * This interface is implemented by build time compiled wasm modules. + * A wasm module together with a machine factory for its instances. */ public interface CompiledModule { WasmModule wasmModule(); diff --git a/runtime/src/main/java/run/endive/runtime/Instance.java b/runtime/src/main/java/run/endive/runtime/Instance.java index ea93cc1cb..7f4204bc0 100644 --- a/runtime/src/main/java/run/endive/runtime/Instance.java +++ b/runtime/src/main/java/run/endive/runtime/Instance.java @@ -343,10 +343,15 @@ public ExportFunction export(String name) { return this.fluentExports.function(name); } + /** + * Returns the body of the function, or {@code null} for an imported function. + * + * @throws IllegalStateException if the body has been dropped from the module + */ public FunctionBody function(long idx) { if (idx < 0 || idx >= (functions.length + imports.functionCount())) { throw new InvalidException("unknown function " + idx); - } else if (isImportedFunction(idx)) { + } else if (idx < imports.functionCount()) { return null; } var body = functions[(int) idx - imports.functionCount()]; @@ -357,10 +362,6 @@ public FunctionBody function(long idx) { return body; } - public boolean isImportedFunction(long idx) { - return idx < imports.functionCount(); - } - public int functionCount() { return imports.functionCount() + functions.length; } diff --git a/runtime/src/main/java/run/endive/runtime/InterpreterMachine.java b/runtime/src/main/java/run/endive/runtime/InterpreterMachine.java index 62d19fc77..53dc6d31d 100644 --- a/runtime/src/main/java/run/endive/runtime/InterpreterMachine.java +++ b/runtime/src/main/java/run/endive/runtime/InterpreterMachine.java @@ -2127,20 +2127,8 @@ private void CALL_REF() { if (funcId == REF_NULL_VALUE) { throw new TrapException("Trapped on call_ref on null function reference"); } - var typeId = instance.functionType(funcId); - var type = instance.type(typeId); - // given a list of param types, let's pop those params off the stack - // and pass as args to the function call - var extracted = extractArgsAndRefsForParams(stack, type.params(), instance); - call( - stack, - instance, - callStack, - funcId, - (long[]) extracted[0], - (Object[]) extracted[1], - type, - false); + // same as a call to funcId + CALL(index -> funcId); } private static void F64_NEG(MStack stack) { diff --git a/runtime/src/main/java/run/endive/runtime/internal/CompilerInterpreterMachine.java b/runtime/src/main/java/run/endive/runtime/internal/CompilerInterpreterMachine.java index 643bb2f87..5d6735f17 100644 --- a/runtime/src/main/java/run/endive/runtime/internal/CompilerInterpreterMachine.java +++ b/runtime/src/main/java/run/endive/runtime/internal/CompilerInterpreterMachine.java @@ -60,7 +60,7 @@ protected long[] call( protected void CALL(Operands operands) { var instance = instance(); var funcId = (int) operands.get(0); - if (interpretedFuncIds.contains(funcId) || instance.isImportedFunction(funcId)) { + if (interpretedFuncIds.contains(funcId) || funcId < instance.imports().functionCount()) { // continue interpreting for interpreted functions or imported functions super.CALL(operands); } else { diff --git a/wasm-corpus/src/main/resources/compiled/drop_bodies.wat.wasm b/wasm-corpus/src/main/resources/compiled/drop_bodies.wat.wasm new file mode 100644 index 0000000000000000000000000000000000000000..76f09a2a00dfcf6783f0f57a029af928d19dc977 GIT binary patch literal 175 zcmX|(F%H5o00ZqNP0~h42*J+4zz9#^B_(QuRFnWJU6}euUPK6}Ue4#iaL-5pR6KLy zB(zGFqW7X2T_2j(g#cl`Xa;|QT+DO!vMMP+)CZ-p#fev|5|_wW@rS~iG`_LMzieP} els6YWcLJup$44K0M?1g!uBfvNKDd{Z1ji4tk0Iay literal 0 HcmV?d00001 diff --git a/wasm-corpus/src/main/resources/compiled/tail_call_drop_bodies.wat.wasm b/wasm-corpus/src/main/resources/compiled/tail_call_drop_bodies.wat.wasm deleted file mode 100644 index 2ecb0e7af6ca00318c7903ecb3083edc518e5642..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 110 zcmZQbEY4+QU|?WmV@zPIXRK#tVFvQp`5C!0^GZ^S3W`!oQd1ZhxWt*+7!(+U7`Z`| x024PmkaA?q;%3lc$xF;lWmIKmK-R&?U6PoY6Q7)zlardl#F3nzTaXE40RR?$7lHr) diff --git a/wasm-corpus/src/main/resources/wat/drop_bodies.wat b/wasm-corpus/src/main/resources/wat/drop_bodies.wat new file mode 100644 index 000000000..98ad3420d --- /dev/null +++ b/wasm-corpus/src/main/resources/wat/drop_bodies.wat @@ -0,0 +1,25 @@ +(module + ;; Each export reaches $compiled in a different way when it is the only interpreted function. + (type $t (func (param i32) (result i32))) + (func $compiled (type $t) + local.get 0 + i32.const 1 + i32.add) + (func $tail_called (type $t) + local.get 0 + call $compiled) + (func $tail_relay (type $t) + local.get 0 + return_call $tail_called) + (func (export "tail_call") (type $t) + local.get 0 + return_call $tail_relay) + (func (export "call_ref") (type $t) + local.get 0 + ref.func $compiled + call_ref $t) + (func (export "tail_call_ref") (type $t) + local.get 0 + ref.func $compiled + return_call_ref $t) + (elem declare func $compiled)) diff --git a/wasm-corpus/src/main/resources/wat/tail_call_drop_bodies.wat b/wasm-corpus/src/main/resources/wat/tail_call_drop_bodies.wat deleted file mode 100644 index b6b3f184b..000000000 --- a/wasm-corpus/src/main/resources/wat/tail_call_drop_bodies.wat +++ /dev/null @@ -1,12 +0,0 @@ -(module - ;; An interpreted function reaching another one through a tail call, which calls a third. - (func $interpreted (export "interpreted") (param i32) (result i32) - local.get 0 - return_call $tail_called) - (func $tail_called (param i32) (result i32) - local.get 0 - call $compiled) - (func $compiled (param i32) (result i32) - local.get 0 - i32.const 1 - i32.add)) diff --git a/wasm/src/main/java/run/endive/wasm/WasmModule.java b/wasm/src/main/java/run/endive/wasm/WasmModule.java index 131e9271b..43e146f43 100644 --- a/wasm/src/main/java/run/endive/wasm/WasmModule.java +++ b/wasm/src/main/java/run/endive/wasm/WasmModule.java @@ -106,13 +106,13 @@ public CodeSection codeSection() { } /** - * Returns a copy of this module that keeps only the bodies of the functions in {@code - * funcIds}, sharing every other section. Accessing a dropped body throws {@link + * Returns a copy of this module without custom sections and without the bodies of the + * functions not in {@code funcIds}. Accessing a dropped body throws {@link * IllegalStateException}. * * @param funcIds function indices in the module's function index space (imports included) */ - public WasmModule retainingFunctionBodies(Set funcIds) { + public WasmModule stripped(Set funcIds) { int functionImports = importSection.count(ExternalType.FUNCTION); var indices = funcIds.stream() @@ -133,7 +133,7 @@ public WasmModule retainingFunctionBodies(Set funcIds) { dataSection, dataCountSection, tagSection, - customSections, + Map.of(), ignoredSections, digest); } diff --git a/wasm/src/main/java/run/endive/wasm/types/CodeSection.java b/wasm/src/main/java/run/endive/wasm/types/CodeSection.java index c35f842a7..4d2ace2d1 100644 --- a/wasm/src/main/java/run/endive/wasm/types/CodeSection.java +++ b/wasm/src/main/java/run/endive/wasm/types/CodeSection.java @@ -40,7 +40,7 @@ public FunctionBody getFunctionBody(int idx) { /** * Returns a copy keeping only the function bodies at {@code indices}, see {@link - * run.endive.wasm.WasmModule#retainingFunctionBodies(Set)}. + * run.endive.wasm.WasmModule#stripped(Set)}. * * @param indices the code section indices of the function bodies to keep */