Move bound module instance from WasmFunctionNode to WasmRootNode.

zakkak · Sep 13, 2023 · 6ecdd1e · 6ecdd1e
1 parent fd7f2f9
commit 6ecdd1e
Show file tree

Hide file tree

Showing 52 changed files with 128 additions and 104 deletions.
diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/WasmInstantiator.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/WasmInstantiator.java
@@ -489,7 +489,7 @@ private CallTarget instantiateCodeEntry(WasmContext context, WasmModule module,
         if (context.language().isMultiContext()) {
             function.setTarget(callTarget);
         } else {
-            functionNode.setBoundModuleInstance(instance);
+            rootNode.setBoundModuleInstance(instance);
         }
         return callTarget;
     }

diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/api/ExecuteInParentContextNode.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/api/ExecuteInParentContextNode.java
@@ -42,6 +42,7 @@
 
 import org.graalvm.wasm.WasmArguments;
 import org.graalvm.wasm.WasmContext;
+import org.graalvm.wasm.WasmInstance;
 import org.graalvm.wasm.WasmLanguage;
 import org.graalvm.wasm.WasmModule;
 import org.graalvm.wasm.exception.Failure;
@@ -68,7 +69,7 @@ public ExecuteInParentContextNode(WasmLanguage language, WasmModule module, Obje
     }
 
     @Override
-    public Object executeWithContext(VirtualFrame frame, WasmContext context) {
+    public Object executeWithContext(VirtualFrame frame, WasmContext context, WasmInstance instance) {
         // Imported executables come from the parent context
         TruffleContext truffleContext = context.environment().getContext().getParent();
         Object prev = truffleContext.enter(this);

diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/nodes/WasmFunctionNode.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/nodes/WasmFunctionNode.java
@@ -145,8 +145,6 @@ public final class WasmFunctionNode extends Node implements BytecodeOSRNode {
     @CompilationFinal private int bytecodeEndOffset;
     @CompilationFinal(dimensions = 1) private byte[] bytecode;
     @CompilationFinal private WasmNotifyFunction notifyFunction;
-    /** Bound module instance (single-context mode only). */
-    @CompilationFinal private WasmInstance boundInstance;
 
     public WasmFunctionNode(WasmModule module, WasmCodeEntry codeEntry, int bytecodeStartOffset, int bytecodeEndOffset) {
         this.module = module;
@@ -178,24 +176,6 @@ void updateBytecode(byte[] bytecode, int bytecodeStartOffset, int bytecodeEndOff
         this.notifyFunction = notifyFunction;
     }
 
-    WasmInstance instance(VirtualFrame frame) {
-        WasmInstance instance = boundInstance;
-        if (instance == null) {
-            instance = WasmArguments.getModuleInstance(frame.getArguments());
-        } else {
-            CompilerAsserts.partialEvaluationConstant(instance);
-            assert instance == WasmArguments.getModuleInstance(frame.getArguments());
-        }
-        assert instance == WasmContext.get(this).lookupModuleInstance(module);
-        return instance;
-    }
-
-    void setBoundModuleInstance(WasmInstance boundInstance) {
-        CompilerAsserts.neverPartOfCompilation();
-        assert this.boundInstance == null;
-        this.boundInstance = boundInstance;
-    }
-
     private WasmMemory memory(WasmInstance instance) {
         return memory(instance, 0);
     }
@@ -219,7 +199,8 @@ private static final class WasmOSRInterpreterState {
     public Object executeOSR(VirtualFrame osrFrame, int target, Object interpreterState) {
         WasmOSRInterpreterState state = (WasmOSRInterpreterState) interpreterState;
         WasmContext context = WasmContext.get(this);
-        return executeBodyFromOffset(context, osrFrame, target, state.stackPointer, state.line);
+        WasmInstance instance = ((WasmRootNode) getRootNode()).instance(osrFrame);
+        return executeBodyFromOffset(context, instance, osrFrame, target, state.stackPointer, state.line);
     }
 
     @Override
@@ -245,14 +226,14 @@ private static final class BackEdgeCounter {
         int count;
     }
 
-    public void execute(VirtualFrame frame, WasmContext context) {
-        executeBodyFromOffset(context, frame, bytecodeStartOffset, codeEntry.localCount(), -1);
+    public void execute(VirtualFrame frame, WasmContext context, WasmInstance instance) {
+        executeBodyFromOffset(context, instance, frame, bytecodeStartOffset, codeEntry.localCount(), -1);
     }
 
     @BytecodeInterpreterSwitch
     @ExplodeLoop(kind = ExplodeLoop.LoopExplosionKind.MERGE_EXPLODE)
     @SuppressWarnings({"UnusedAssignment", "hiding"})
-    public Object executeBodyFromOffset(WasmContext context, VirtualFrame frame, int startOffset, int startStackPointer, int startLine) {
+    public Object executeBodyFromOffset(WasmContext context, WasmInstance instance, VirtualFrame frame, int startOffset, int startStackPointer, int startLine) {
         final int localCount = codeEntry.localCount();
         final byte[] bytecode = this.bytecode;
 
@@ -266,7 +247,6 @@ public Object executeBodyFromOffset(WasmContext context, VirtualFrame frame, int
         int stackPointer = startStackPointer;
         int line = startLine;
 
-        final WasmInstance instance = instance(frame);
         // Note: The module may not have any memories.
         final WasmMemory zeroMemory = module.memoryCount() == 0 ? null : memory(instance);
 

diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/nodes/WasmInstrumentableFunctionNode.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/nodes/WasmInstrumentableFunctionNode.java
@@ -55,7 +55,6 @@
 import org.graalvm.wasm.debugging.representation.DebugObjectDisplayValue;
 import org.graalvm.wasm.memory.WasmMemory;
 
-import com.oracle.truffle.api.CompilerAsserts;
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 import com.oracle.truffle.api.frame.Frame;
 import com.oracle.truffle.api.frame.MaterializedFrame;
@@ -100,13 +99,8 @@ protected WasmInstrumentableFunctionNode(WasmInstrumentableFunctionNode node) {
         this.instrumentation = node.instrumentation;
     }
 
-    final WasmInstance instance(VirtualFrame frame) {
-        return functionNode.instance(frame);
-    }
-
-    public void setBoundModuleInstance(WasmInstance boundInstance) {
-        CompilerAsserts.neverPartOfCompilation();
-        functionNode.setBoundModuleInstance(boundInstance);
+    private WasmInstance instance(VirtualFrame frame) {
+        return ((WasmRootNode) getRootNode()).instance(frame);
     }
 
     private WasmMemory memory0(MaterializedFrame frame) {
@@ -125,8 +119,8 @@ int resultCount() {
         return codeEntry.resultCount();
     }
 
-    void execute(VirtualFrame frame, WasmContext context) {
-        functionNode.execute(frame, context);
+    void execute(VirtualFrame frame, WasmContext context, WasmInstance instance) {
+        functionNode.execute(frame, context, instance);
     }
 
     void enterErrorBranch() {

diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/nodes/WasmMemoryOverheadModeRootNode.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/nodes/WasmMemoryOverheadModeRootNode.java
@@ -45,6 +45,7 @@
 import com.oracle.truffle.api.frame.FrameDescriptor;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import org.graalvm.wasm.WasmContext;
+import org.graalvm.wasm.WasmInstance;
 import org.graalvm.wasm.exception.Failure;
 import org.graalvm.wasm.exception.WasmException;
 import org.graalvm.wasm.WasmOptions;
@@ -62,7 +63,7 @@ public WasmMemoryOverheadModeRootNode(TruffleLanguage<?> language, FrameDescript
     }
 
     @Override
-    public Object executeWithContext(VirtualFrame frame, WasmContext context) {
+    public Object executeWithContext(VirtualFrame frame, WasmContext context, WasmInstance instance) {
         throw WasmException.create(Failure.MEMORY_OVERHEAD_MODE, this, "WebAssembly functions cannot be executed with memory overhead mode enabled, since this can lead to unexpected behavior.");
     }
 }
diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/nodes/WasmRootNode.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/nodes/WasmRootNode.java
@@ -63,6 +63,7 @@
 import org.graalvm.wasm.memory.WasmMemory;
 
 import com.oracle.truffle.api.CompilerAsserts;
+import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 import com.oracle.truffle.api.TruffleLanguage;
 import com.oracle.truffle.api.frame.FrameDescriptor;
@@ -78,6 +79,8 @@ public class WasmRootNode extends RootNode {
     private SourceSection sourceSection;
     @Child private WasmInstrumentableFunctionNode functionNode;
     private final BranchProfile nonLinkedProfile = BranchProfile.create();
+    /** Bound module instance (single-context mode only). */
+    @CompilationFinal private WasmInstance boundInstance;
 
     public WasmRootNode(TruffleLanguage<?> language, FrameDescriptor frameDescriptor, WasmInstrumentableFunctionNode functionNode) {
         super(language, frameDescriptor);
@@ -106,11 +109,22 @@ public final void tryInitialize(WasmContext context, WasmInstance instance) {
         }
     }
 
-    /**
-     * Overridden by {@link org.graalvm.wasm.predefined.WasmBuiltinRootNode}.
-     */
-    protected WasmInstance instance(VirtualFrame frame) {
-        return functionNode.instance(frame);
+    protected final WasmInstance instance(VirtualFrame frame) {
+        WasmInstance instance = boundInstance;
+        if (instance == null) {
+            instance = WasmArguments.getModuleInstance(frame.getArguments());
+        } else {
+            CompilerAsserts.partialEvaluationConstant(instance);
+            assert instance == WasmArguments.getModuleInstance(frame.getArguments());
+        }
+        assert instance == WasmContext.get(this).lookupModuleInstance(module());
+        return instance;
+    }
+
+    public final void setBoundModuleInstance(WasmInstance boundInstance) {
+        CompilerAsserts.neverPartOfCompilation();
+        assert this.boundInstance == null;
+        this.boundInstance = boundInstance;
     }
 
     protected final WasmMemory memory(VirtualFrame frame) {
@@ -125,11 +139,12 @@ protected final WasmMemory memory(VirtualFrame frame, int index) {
     public final Object execute(VirtualFrame frame) {
         assert WasmArguments.isValid(frame.getArguments());
         final WasmContext context = getContext();
-        tryInitialize(context, instance(frame));
-        return executeWithContext(frame, context);
+        final WasmInstance instance = instance(frame);
+        tryInitialize(context, instance);
+        return executeWithContext(frame, context, instance);
     }
 
-    public Object executeWithContext(VirtualFrame frame, WasmContext context) {
+    public Object executeWithContext(VirtualFrame frame, WasmContext context, WasmInstance instance) {
         // WebAssembly structure dictates that a function's arguments are provided to the function
         // as local variables, followed by any additional local variables that the function
         // declares. A VirtualFrame contains a special array for the arguments, so we need to move
@@ -156,7 +171,7 @@ public Object executeWithContext(VirtualFrame frame, WasmContext context) {
         }
 
         try {
-            functionNode.execute(frame, context);
+            functionNode.execute(frame, context, instance);
         } catch (StackOverflowError e) {
             functionNode.enterErrorBranch();
             throw WasmException.create(Failure.CALL_STACK_EXHAUSTED);

diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/predefined/WasmBuiltinRootNode.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/predefined/WasmBuiltinRootNode.java
@@ -40,15 +40,10 @@
  */
 package org.graalvm.wasm.predefined;
 
-import org.graalvm.wasm.WasmArguments;
-import org.graalvm.wasm.WasmContext;
-import org.graalvm.wasm.WasmInstance;
 import org.graalvm.wasm.WasmLanguage;
 import org.graalvm.wasm.WasmModule;
 import org.graalvm.wasm.nodes.WasmRootNode;
 
-import com.oracle.truffle.api.frame.VirtualFrame;
-
 public abstract class WasmBuiltinRootNode extends WasmRootNode {
     protected final WasmModule module;
 
@@ -64,13 +59,6 @@ protected WasmModule module() {
         return module;
     }
 
-    @Override
-    protected WasmInstance instance(VirtualFrame frame) {
-        WasmInstance instance = WasmArguments.getModuleInstance(frame.getArguments());
-        assert instance == WasmContext.get(this).lookupModuleInstance(module);
-        return instance;
-    }
-
     @Override
     public String getName() {
         return "wasm-function:" + builtinNodeName();

diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/AbortNode.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/AbortNode.java
@@ -42,6 +42,7 @@
 
 import org.graalvm.wasm.WasmArguments;
 import org.graalvm.wasm.WasmContext;
+import org.graalvm.wasm.WasmInstance;
 import org.graalvm.wasm.WasmLanguage;
 import org.graalvm.wasm.WasmModule;
 import org.graalvm.wasm.exception.Failure;
@@ -57,7 +58,7 @@ public AbortNode(WasmLanguage language, WasmModule module) {
     }
 
     @Override
-    public Object executeWithContext(VirtualFrame frame, WasmContext context) {
+    public Object executeWithContext(VirtualFrame frame, WasmContext context, WasmInstance instance) {
         Object[] args = frame.getArguments();
         final int code = WasmArguments.getArgumentCount(args) > 0 ? (int) WasmArguments.getArgument(args, 0) : 0;
         throw fail(code);

diff --git a/....graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/AbortOnCannotGrowMemoryNode.java b/....graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/AbortOnCannotGrowMemoryNode.java
@@ -41,6 +41,7 @@
 package org.graalvm.wasm.predefined.emscripten;
 
 import org.graalvm.wasm.WasmContext;
+import org.graalvm.wasm.WasmInstance;
 import org.graalvm.wasm.WasmLanguage;
 import org.graalvm.wasm.WasmModule;
 
@@ -52,7 +53,7 @@ public AbortOnCannotGrowMemoryNode(WasmLanguage language, WasmModule module) {
     }
 
     @Override
-    public Object executeWithContext(VirtualFrame frame, WasmContext context) {
+    public Object executeWithContext(VirtualFrame frame, WasmContext context, WasmInstance instance) {
         return super.execute(frame);
     }
 

diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/AlignfaultNode.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/AlignfaultNode.java
@@ -41,6 +41,7 @@
 package org.graalvm.wasm.predefined.emscripten;
 
 import org.graalvm.wasm.WasmContext;
+import org.graalvm.wasm.WasmInstance;
 import org.graalvm.wasm.WasmLanguage;
 import org.graalvm.wasm.WasmModule;
 
@@ -52,7 +53,7 @@ public AlignfaultNode(WasmLanguage language, WasmModule module) {
     }
 
     @Override
-    public Object executeWithContext(VirtualFrame frame, WasmContext context) {
+    public Object executeWithContext(VirtualFrame frame, WasmContext context, WasmInstance instance) {
         return super.execute(frame);
     }
 

diff --git a/...rg.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/EmscriptenGetHeapSizeNode.java b/...rg.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/EmscriptenGetHeapSizeNode.java
@@ -41,6 +41,7 @@
 package org.graalvm.wasm.predefined.emscripten;
 
 import org.graalvm.wasm.WasmContext;
+import org.graalvm.wasm.WasmInstance;
 import org.graalvm.wasm.WasmLanguage;
 import org.graalvm.wasm.WasmModule;
 import org.graalvm.wasm.memory.WasmMemory;
@@ -54,7 +55,7 @@ public EmscriptenGetHeapSizeNode(WasmLanguage language, WasmModule module) {
     }
 
     @Override
-    public Object executeWithContext(VirtualFrame frame, WasmContext context) {
+    public Object executeWithContext(VirtualFrame frame, WasmContext context, WasmInstance instance) {
         WasmMemory memory = memory(frame);
         final long byteSize = memory.byteSize();
         return (int) byteSize;

diff --git a/.../org.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/EmscriptenMemcpyBigNode.java b/.../org.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/EmscriptenMemcpyBigNode.java
@@ -42,6 +42,7 @@
 
 import org.graalvm.wasm.WasmArguments;
 import org.graalvm.wasm.WasmContext;
+import org.graalvm.wasm.WasmInstance;
 import org.graalvm.wasm.WasmLanguage;
 import org.graalvm.wasm.WasmModule;
 import org.graalvm.wasm.memory.WasmMemory;
@@ -55,7 +56,7 @@ public EmscriptenMemcpyBigNode(WasmLanguage language, WasmModule module) {
     }
 
     @Override
-    public Object executeWithContext(VirtualFrame frame, WasmContext context) {
+    public Object executeWithContext(VirtualFrame frame, WasmContext context, WasmInstance instance) {
         Object[] args = frame.getArguments();
         assert WasmArguments.getArgumentCount(args) == 3;
 

diff --git a/...lvm.wasm/src/org/graalvm/wasm/predefined/emscripten/EmscriptenNotifyMemoryGrowthNode.java b/...lvm.wasm/src/org/graalvm/wasm/predefined/emscripten/EmscriptenNotifyMemoryGrowthNode.java
@@ -42,6 +42,7 @@
 
 import org.graalvm.wasm.WasmConstant;
 import org.graalvm.wasm.WasmContext;
+import org.graalvm.wasm.WasmInstance;
 import org.graalvm.wasm.WasmLanguage;
 import org.graalvm.wasm.WasmModule;
 import org.graalvm.wasm.predefined.WasmBuiltinRootNode;
@@ -54,7 +55,7 @@ public EmscriptenNotifyMemoryGrowthNode(WasmLanguage language, WasmModule module
     }
 
     @Override
-    public Object executeWithContext(VirtualFrame frame, WasmContext context) {
+    public Object executeWithContext(VirtualFrame frame, WasmContext context, WasmInstance instance) {
         return WasmConstant.VOID;
     }
 

diff --git a/...org.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/EmscriptenResizeHeapNode.java b/...org.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/EmscriptenResizeHeapNode.java
@@ -41,6 +41,7 @@
 package org.graalvm.wasm.predefined.emscripten;
 
 import org.graalvm.wasm.WasmContext;
+import org.graalvm.wasm.WasmInstance;
 import org.graalvm.wasm.WasmLanguage;
 import org.graalvm.wasm.WasmModule;
 
@@ -52,7 +53,7 @@ public EmscriptenResizeHeapNode(WasmLanguage language, WasmModule module) {
     }
 
     @Override
-    public Object executeWithContext(VirtualFrame frame, WasmContext context) {
+    public Object executeWithContext(VirtualFrame frame, WasmContext context, WasmInstance instance) {
         // Heap resizing is not supported by default by emscripten
         // (need to specify `-s ALLOW_MEMORY_GROWTH=1` on compilation).
         return super.execute(frame);

diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/GetTimeOfDayNode.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/predefined/emscripten/GetTimeOfDayNode.java
@@ -42,6 +42,7 @@
 
 import org.graalvm.wasm.WasmArguments;
 import org.graalvm.wasm.WasmContext;
+import org.graalvm.wasm.WasmInstance;
 import org.graalvm.wasm.WasmLanguage;
 import org.graalvm.wasm.WasmModule;
 import org.graalvm.wasm.memory.WasmMemory;
@@ -56,7 +57,7 @@ public GetTimeOfDayNode(WasmLanguage language, WasmModule module) {
     }
 
     @Override
-    public Object executeWithContext(VirtualFrame frame, WasmContext context) {
+    public Object executeWithContext(VirtualFrame frame, WasmContext context, WasmInstance instance) {
         Object[] args = frame.getArguments();
         assert WasmArguments.getArgumentCount(args) == 2;