feat: option to drop compiled function bodies

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<Instance, Machine> machineFactory) {
+        return Instance.builder(module)
+                .withMachineFactory(machineFactory)
+                .withImportValues(ImportValues.builder().addFunction(HOST_FUNC).build())
+                .withStart(false)
+                .build();
+    }
+
+    private static void assertRetained(CodeSection code, Set<Integer> 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"));
+    }
+}

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<Instance, Machine> compile() {
             try {
 
@@ -152,13 +168,59 @@ public Function<Instance, Machine> 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);
             }
         }
     }
 
+    private static void dropCompiledFunctionBodies(
+            WasmModule module, Set<Integer> 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<Integer> interpretableFunctions(
+            WasmModule module, Set<Integer> interpretedFunctions) {
+        int functionImports = module.importSection().count(ExternalType.FUNCTION);
+        var result = new HashSet<Integer>();
+        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

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() {

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 {

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))

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;
     }

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<FunctionBody> functionBodies;
+    // A null entry is a body dropped by dropFunctionBody.
+    private final FunctionBody[] functionBodies;
     private final boolean requiresDataCount;
 
     private CodeSection(List<FunctionBody> 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);
     }
 }