[GR-32062] Enter TruffleContext in call_indirect when it is needed only.

PullRequest: graal/9403
mttcr · Jul 22, 2021 · 2bea06d · 2bea06d
2 parents cdc7c79 + e862256
commit 2bea06d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 14 deletions.
diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/WasmContext.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/WasmContext.java
@@ -51,7 +51,6 @@
 import java.util.Map;
 
 public final class WasmContext {
-    private final Uid uid;
     private final Env env;
     private final WasmLanguage language;
     private final Map<SymbolTable.FunctionType, Integer> equivalenceClasses;
@@ -69,7 +68,6 @@ public static WasmContext getCurrent() {
     }
 
     public WasmContext(Env env, WasmLanguage language) {
-        this.uid = new Uid();
         this.env = env;
         this.language = language;
         this.equivalenceClasses = new HashMap<>();
@@ -84,10 +82,6 @@ public WasmContext(Env env, WasmLanguage language) {
         instantiateBuiltinInstances();
     }
 
-    public Uid uid() {
-        return uid;
-    }
-
     public Env environment() {
         return env;
     }
@@ -204,6 +198,4 @@ public void reinitInstance(WasmInstance instance, boolean reinitMemory) {
         }
     }
 
-    public class Uid {
-    }
 }
diff --git a/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/nodes/WasmBlockNode.java b/wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/nodes/WasmBlockNode.java
@@ -637,12 +637,12 @@ public int execute(WasmContext context, VirtualFrame frame, long[] stacklocals)
                         final WasmFunctionInstance functionInstance;
                         final WasmFunction function;
                         final CallTarget target;
-                        final WasmContext.Uid functionInstanceContextUid;
+                        final WasmContext functionInstanceContext;
                         if (element instanceof WasmFunctionInstance) {
                             functionInstance = (WasmFunctionInstance) element;
                             function = functionInstance.function();
                             target = functionInstance.target();
-                            functionInstanceContextUid = functionInstance.context().uid();
+                            functionInstanceContext = functionInstance.context();
                         } else {
                             CompilerDirectives.transferToInterpreter();
                             throw WasmException.format(Failure.UNSPECIFIED_TRAP, this, "Unknown table element type: %s", element);
@@ -663,7 +663,8 @@ public int execute(WasmContext context, VirtualFrame frame, long[] stacklocals)
                         offset += 1;
 
                         // Validate that the function type matches the expected type.
-                        if (functionInstanceContextUid == context.uid()) {
+                        boolean functionFromCurrentContext = (functionInstanceContext == context);
+                        if (functionFromCurrentContext) {
                             // We can do a quick equivalence-class check.
                             if (expectedTypeEquivalenceClass != function.typeEquivalenceClass()) {
                                 CompilerDirectives.transferToInterpreter();
@@ -685,13 +686,25 @@ public int execute(WasmContext context, VirtualFrame frame, long[] stacklocals)
                         Object[] args = createArgumentsForCall(stacklocals, expectedFunctionTypeIndex, numArgs, stackPointer);
                         stackPointer -= args.length;
 
-                        TruffleContext truffleContext = functionInstance.context().environment().getContext();
-                        Object prev = truffleContext.enter(this);
+                        // Enter function's context when it is not from the current one
+                        final boolean enterContext = !functionFromCurrentContext;
+                        TruffleContext truffleContext;
+                        Object prev;
+                        if (enterContext) {
+                            truffleContext = functionInstanceContext.environment().getContext();
+                            prev = truffleContext.enter(this);
+                        } else {
+                            truffleContext = null;
+                            prev = null;
+                        }
+
                         final Object result;
                         try {
                             result = executeIndirectCallNode(childrenOffset, target, args);
                         } finally {
-                            truffleContext.leave(this, prev);
+                            if (enterContext) {
+                                truffleContext.leave(this, prev);
+                            }
                         }
 
                         childrenOffset++;