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 new file mode 100644 index 000000000..83779c70c --- /dev/null +++ b/compiler-tests/src/test/java/run/endive/testing/CompileModuleTest.java @@ -0,0 +1,165 @@ +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; +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 { + + // 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 long call(CompiledModule compiled, String export) { + return Instance.builder(compiled).build().export(export).apply(41)[0]; + } + + 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() { + // 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(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. + assertEquals(35, instance.export("func_2").apply(1)[0]); + } + + @Test + public void keepsFunctionsReachableThroughTailCalls() { + 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, 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 = interpreting(CALL_REF).compile(); + + assertSame(DROP_BODIES, ((CompiledModule) factory).wasmModule()); + } + + @Test + public void instanceRejectsDroppedBody() { + var instance = Instance.builder(interpreting(CALL_REF).compileModule()).build(); + + 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 f4a1e543b..6ba3c4cfe 100644 --- a/compiler/src/main/java/run/endive/compiler/MachineFactoryCompiler.java +++ b/compiler/src/main/java/run/endive/compiler/MachineFactoryCompiler.java @@ -1,9 +1,13 @@ 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.Arrays; import java.util.HashMap; import java.util.Map; import java.util.Properties; @@ -13,7 +17,10 @@ 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; @@ -24,6 +31,8 @@ */ public final class MachineFactoryCompiler { + private static final String INTERPRETED_FUNCTIONS = "interpretedFunctions"; + private MachineFactoryCompiler() { // no instances } @@ -127,6 +136,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 + * 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 strip) { try { // Can we load the byte codes from the cache? @@ -134,8 +157,7 @@ public Function compile() { 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); } } @@ -145,21 +167,32 @@ public Function compile() { .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)); } - return new MachineFactory(module, collector.machineFactory()); + return machineFactory(result, strip); } catch (IOException e) { throw new WasmEngineException(e); } } + + 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); + } } - 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(); @@ -167,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, ""); @@ -190,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))) { @@ -206,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", ""); @@ -231,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 45004fa6a..7f4204bc0 100644 --- a/runtime/src/main/java/run/endive/runtime/Instance.java +++ b/runtime/src/main/java/run/endive/runtime/Instance.java @@ -343,13 +343,23 @@ 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 (idx < imports.functionCount()) { 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 int functionCount() { @@ -617,6 +627,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/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 805d3bac6..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.function(funcId) == null) { + 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 000000000..76f09a2a0 Binary files /dev/null and b/wasm-corpus/src/main/resources/compiled/drop_bodies.wat.wasm differ 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/src/main/java/run/endive/wasm/WasmModule.java b/wasm/src/main/java/run/endive/wasm/WasmModule.java index ed732b320..43e146f43 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 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 stripped(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, + Map.of(), + 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..4d2ace2d1 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#stripped(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() {