Make the CallTarget for Method#to_proc context-independent

Author: eregon

src/main/java/org/truffleruby/core/method/MethodNodes.java

@@ -45,6 +45,7 @@
 import org.truffleruby.language.control.RaiseException;
 import org.truffleruby.language.methods.CallInternalMethodNode;
 import org.truffleruby.language.methods.InternalMethod;
+import org.truffleruby.language.methods.SharedMethodInfo;
 import org.truffleruby.language.objects.AllocationTracing;
 import org.truffleruby.language.objects.MetaClassNode;
 import org.truffleruby.language.threadlocal.SpecialVariableStorage;
@@ -279,28 +280,29 @@ protected RubyUnboundMethod unbind(RubyMethod method) {
     @CoreMethod(names = "to_proc")
     public abstract static class ToProcNode extends CoreMethodArrayArgumentsNode {
 
-        @Specialization(guards = "methodObject == cachedMethodObject", limit = "getCacheLimit()")
-        protected RubyProc toProcCached(RubyMethod methodObject,
+        @Specialization(guards = { "isSingleContext()", "methodObject == cachedMethodObject" },
+                limit = "getCacheLimit()")
+        protected RubyProc toProcCachedSingleContext(RubyMethod methodObject,
                 @Cached("methodObject") RubyMethod cachedMethodObject,
                 @Cached("toProcUncached(cachedMethodObject)") RubyProc proc) {
             return proc;
         }
 
         @Specialization(
-                guards = "cachedMethod == methodObject.method",
+                guards = "methodObject.method.getCallTarget() == methodCallTarget",
                 limit = "getCacheLimit()",
-                replaces = "toProcCached")
+                replaces = "toProcCachedSingleContext")
         protected RubyProc toProcCachedTarget(RubyMethod methodObject,
-                @Cached("methodObject.method") InternalMethod cachedMethod,
-                @Cached("methodCallTarget(cachedMethod)") RootCallTarget callTarget) {
-            return createProc(callTarget, cachedMethod, methodObject.receiver);
+                @Cached("methodObject.method.getCallTarget()") RootCallTarget methodCallTarget,
+                @Cached("procCallTargetToCallRubyMethod(methodCallTarget)") RootCallTarget procCallTarget) {
+            return createProc(procCallTarget, methodObject.method, methodObject.receiver);
         }
 
-        @Specialization
+        @Specialization(replaces = { "toProcCachedSingleContext", "toProcCachedTarget" })
         protected RubyProc toProcUncached(RubyMethod methodObject) {
             final InternalMethod method = methodObject.method;
-            final RootCallTarget callTarget = methodCallTarget(method);
             final Object receiver = methodObject.receiver;
+            final RootCallTarget callTarget = procCallTargetToCallRubyMethod(method.getCallTarget());
             return createProc(callTarget, method, receiver);
         }
 
@@ -324,26 +326,26 @@ private RubyProc createProc(RootCallTarget callTarget, InternalMethod method, Ob
         }
 
         @TruffleBoundary
-        protected RootCallTarget methodCallTarget(InternalMethod method) {
+        protected RootCallTarget procCallTargetToCallRubyMethod(RootCallTarget callTarget) {
             // translate to something like:
             // lambda { |same args list| method.call(args) }
             // We need to preserve the method receiver and we want to have the same argument list.
             // We create a new CallTarget for the Proc that calls the method CallTarget and passes the correct receiver.
 
-            final SourceSection sourceSection = method.getSharedMethodInfo().getSourceSection();
-            final RubyRootNode methodRootNode = RubyRootNode.of(method.getCallTarget());
+            final RubyRootNode methodRootNode = RubyRootNode.of(callTarget);
+            final SharedMethodInfo sharedMethodInfo = methodRootNode.getSharedMethodInfo();
 
-            var callWithRubyMethodReceiverNode = new CallWithRubyMethodReceiverNode(method);
+            var callWithRubyMethodReceiverNode = new CallWithRubyMethodReceiverNode();
             final RubyLambdaRootNode wrapRootNode = new RubyLambdaRootNode(
                     getLanguage(),
-                    sourceSection,
+                    sharedMethodInfo.getSourceSection(),
                     methodRootNode.getFrameDescriptor(),
-                    method.getSharedMethodInfo(),
+                    sharedMethodInfo,
                     callWithRubyMethodReceiverNode,
                     methodRootNode.getSplit(),
                     methodRootNode.returnID,
                     BreakID.INVALID,
-                    method.getSharedMethodInfo().getArity());
+                    sharedMethodInfo.getArity());
             return wrapRootNode.getCallTarget();
         }
 
@@ -352,16 +354,13 @@ protected int getCacheLimit() {
         }
 
         private static class CallWithRubyMethodReceiverNode extends RubyContextSourceNode {
-            private final InternalMethod method;
             @Child private CallInternalMethodNode callInternalMethodNode = CallInternalMethodNode.create();
 
-            public CallWithRubyMethodReceiverNode(InternalMethod method) {
-                this.method = method;
-            }
-
             @Override
             public Object execute(VirtualFrame frame) {
-                final Object receiver = RubyArguments.getSelf(RubyArguments.getDeclarationFrame(frame));
+                final MaterializedFrame declarationFrame = RubyArguments.getDeclarationFrame(frame);
+                final Object receiver = RubyArguments.getSelf(declarationFrame);
+                final InternalMethod method = RubyArguments.getMethod(declarationFrame);
                 return CallNode.callBoundMethod(frame, method, receiver, frame.getArguments(), callInternalMethodNode);
             }
         }

src/main/java/org/truffleruby/language/methods/InternalMethod.java

@@ -22,6 +22,7 @@
 import org.truffleruby.language.LexicalScope;
 import org.truffleruby.language.Nil;
 import org.truffleruby.language.RubyBaseNode;
+import org.truffleruby.language.RubyRootNode;
 import org.truffleruby.language.Visibility;
 import org.truffleruby.language.objects.ObjectGraphNode;
 
@@ -160,6 +161,7 @@ private InternalMethod(
         assert declaringModule != null;
         assert lexicalScope != null;
         assert !sharedMethodInfo.isBlock() : sharedMethodInfo;
+        assert callTarget == null || RubyRootNode.of(callTarget).getSharedMethodInfo() == sharedMethodInfo;
         assert capturedBlock instanceof Nil || capturedBlock instanceof RubyProc : capturedBlock;
         this.sharedMethodInfo = sharedMethodInfo;
         this.lexicalScope = lexicalScope;
@@ -235,6 +237,7 @@ public RootCallTarget getCallTarget() {
         if (callTarget == null) {
             CompilerDirectives.transferToInterpreterAndInvalidate();
             callTarget = callTargetSupplier.get();
+            assert RubyRootNode.of(callTarget).getSharedMethodInfo() == sharedMethodInfo;
         }
         return callTarget;
     }