From 0232cc5c97a84d2dd6629ac469ea80db0b91ce78 Mon Sep 17 00:00:00 2001 From: Andreas Karlsson Date: Fri, 23 Oct 2026 04:28:36 +0200 Subject: [PATCH] feat: option to drop compiled function bodies --- .../DropCompiledFunctionBodiesTest.java | 115 ++++++++++++++++++ .../compiler/MachineFactoryCompiler.java | 62 ++++++++++ .../java/run/endive/runtime/Instance.java | 13 +- .../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 | 13 ++ .../run/endive/wasm/types/CodeSection.java | 29 +++-- 8 files changed, 236 insertions(+), 10 deletions(-) create mode 100644 compiler-tests/src/test/java/run/endive/testing/DropCompiledFunctionBodiesTest.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/DropCompiledFunctionBodiesTest.java b/compiler-tests/src/test/java/run/endive/testing/DropCompiledFunctionBodiesTest.java new file mode 100644 index 000000000..242e9de88 --- /dev/null +++ b/compiler-tests/src/test/java/run/endive/testing/DropCompiledFunctionBodiesTest.java @@ -0,0 +1,115 @@ +package run.endive.testing; + +import static org.junit.jupiter.api.Assertions.assertEquals; +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 java.util.function.Function; +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.HostFunction; +import run.endive.runtime.ImportValues; +import run.endive.runtime.Instance; +import run.endive.runtime.Machine; +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 DropCompiledFunctionBodiesTest { + + // func_1 (function index 2, code index 1) is too large to compile and is interpreted. + private static final byte[] METHOD_TOO_LARGE = 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 MachineFactoryCompiler.Builder dropping(WasmModule module) { + return MachineFactoryCompiler.builder(module) + .withInterpreterFallback(InterpreterFallback.SILENT) + .withDropCompiledFunctionBodies(true); + } + + private static Instance instantiate( + WasmModule module, Function machineFactory) { + return Instance.builder(module) + .withMachineFactory(machineFactory) + .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 module = Parser.parse(METHOD_TOO_LARGE); + var factory = dropping(module).compile(); + + var code = module.codeSection(); + assertEquals(3, code.functionBodyCount()); + assertRetained(code, Set.of(1)); + assertThrows(IllegalStateException.class, () -> code.getFunctionBody(0)); + + var instance = instantiate(module, factory); + // 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 keepsAllBodiesByDefault() { + var module = Parser.parse(METHOD_TOO_LARGE); + var factory = + MachineFactoryCompiler.builder(module) + .withInterpreterFallback(InterpreterFallback.SILENT) + .compile(); + + assertRetained(module.codeSection(), Set.of(0, 1, 2)); + assertEquals(35, instantiate(module, factory).export("func_2").apply(0)[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 factory = + MachineFactoryCompiler.builder(module) + .withInterpreterFallback(InterpreterFallback.FAIL) + .withInterpretedFunctions(Set.of(0)) + .withDropCompiledFunctionBodies(true) + .compile(); + + assertRetained(module.codeSection(), Set.of(0, 1)); + var instance = Instance.builder(module).withMachineFactory(factory).build(); + assertEquals(42, instance.export("interpreted").apply(41)[0]); + } + + @Test + public void instanceRejectsDroppedBody() { + var module = Parser.parse(METHOD_TOO_LARGE); + var instance = instantiate(module, dropping(module).compile()); + + 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..df7cb6b25 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; @@ -18,6 +20,8 @@ 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. @@ -86,6 +90,7 @@ public static final class Builder { private final WasmModule module; private final run.endive.compiler.internal.Compiler.Builder compilerBuilder; private Cache cache; + private boolean dropCompiledFunctionBodies; private Builder(WasmModule module) { this.module = module; @@ -126,6 +131,17 @@ public Builder withCache(Cache cache) { return this; } + /** + * Drops the bodies of compiled functions from the module once it has been compiled, + * keeping only those the interpreter may still run (see {@link + * WasmModule#dropFunction(int)}). The module can then only be run with the returned + * machine factory. Nothing is dropped when the compilation is loaded from the cache. + */ + public Builder withDropCompiledFunctionBodies(boolean drop) { + this.dropCompiledFunctionBodies = drop; + return this; + } + public Function compile() { try { @@ -152,6 +168,9 @@ public Function compile() { cache.putIfAbsent(module.digest(), storeClassLoadingCollector(collector)); } + if (dropCompiledFunctionBodies) { + dropCompiledFunctionBodies(module, result.interpretedFunctions()); + } return new MachineFactory(module, collector.machineFactory()); } catch (IOException e) { throw new WasmEngineException(e); @@ -159,6 +178,49 @@ public Function compile() { } } + private static void dropCompiledFunctionBodies( + WasmModule module, Set interpretedFunctions) { + var retained = interpretableFunctions(module, interpretedFunctions); + if (retained == null) { + return; + } + int functionImports = module.importSection().count(ExternalType.FUNCTION); + int functionCount = functionImports + module.codeSection().functionBodyCount(); + for (int funcId = functionImports; funcId < functionCount; funcId++) { + if (!retained.contains(funcId)) { + module.dropFunction(funcId); + } + } + } + + /** + * 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..838badb5e 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() { 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..be47cf042 100644 --- a/wasm/src/main/java/run/endive/wasm/WasmModule.java +++ b/wasm/src/main/java/run/endive/wasm/WasmModule.java @@ -14,6 +14,7 @@ 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 +103,18 @@ public CodeSection codeSection() { return codeSection; } + /** + * Drops the body of a function to free its memory, for example once it has been compiled. + * Accessing a dropped body throws {@link IllegalStateException}, so this module can only be + * run by machines that don't need the dropped bodies. + * + * @param funcIdx the index of a function defined in this module, in the function index space + * (imports included) + */ + public void dropFunction(int funcIdx) { + codeSection.dropFunctionBody(funcIdx - importSection.count(ExternalType.FUNCTION)); + } + 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..d1ec8543b 100644 --- a/wasm/src/main/java/run/endive/wasm/types/CodeSection.java +++ b/wasm/src/main/java/run/endive/wasm/types/CodeSection.java @@ -1,29 +1,44 @@ package run.endive.wasm.types; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.Objects; public final class CodeSection extends Section { - private final List functionBodies; + // A null entry is a body dropped by dropFunctionBody. + private final FunctionBody[] functionBodies; private final boolean requiresDataCount; private CodeSection(List functionBodies, boolean requiresDataCount) { super(SectionId.CODE); - this.functionBodies = List.copyOf(functionBodies); + this.functionBodies = functionBodies.toArray(new FunctionBody[0]); this.requiresDataCount = requiresDataCount; } + /** Returns all function bodies, with {@code null} for dropped ones. */ public FunctionBody[] functionBodies() { - return functionBodies.toArray(new FunctionBody[0]); + return functionBodies.clone(); } public int functionBodyCount() { - return functionBodies.size(); + return functionBodies.length; } + /** + * @throws IllegalStateException if the function body has been dropped + */ public FunctionBody getFunctionBody(int idx) { - return functionBodies.get(idx); + var body = functionBodies[idx]; + if (body == null) { + throw new IllegalStateException("Function body " + idx + " has been dropped"); + } + return body; + } + + /** Drops the function body at {@code idx}, see {@link run.endive.wasm.WasmModule#dropFunction(int)}. */ + public void dropFunctionBody(int idx) { + functionBodies[idx] = null; } public boolean isRequiresDataCount() { @@ -72,11 +87,11 @@ public boolean equals(Object o) { } CodeSection that = (CodeSection) o; return requiresDataCount == that.requiresDataCount - && Objects.equals(functionBodies, that.functionBodies); + && Arrays.equals(functionBodies, that.functionBodies); } @Override public int hashCode() { - return Objects.hash(functionBodies, requiresDataCount); + return Objects.hash(Arrays.hashCode(functionBodies), requiresDataCount); } }