From a41d8a7ab45c8d5e0fcda8e7bbde809c449093cf Mon Sep 17 00:00:00 2001 From: Andreas Karlsson Date: Fri, 23 Oct 2026 14:27:28 +0200 Subject: [PATCH] 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() {