andreas-karlsson · Oct 23, 2026
@@ -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<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 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"));
+    }
+}
@@ -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<Instance, Machine> 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<Instance, Machine> 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<Integer> 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<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

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

@@ -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 {

@@ -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))
@@ -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<Integer> 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;
     }

@@ -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<FunctionBody> functionBodies;
     private final boolean requiresDataCount;
 
     private CodeSection(List<FunctionBody> 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<Integer> indices) {
+        var retained =
+                new ArrayList<FunctionBody>(Collections.nCopies(functionBodies.size(), null));
+        for (int idx : indices) {
+            retained.set(idx, functionBodies.get(idx));
+        }
+        return new CodeSection(retained, requiresDataCount);
     }
 
     public boolean isRequiresDataCount() {