Provide a fast path for getting and setting backrefs from C.

Author: aardvark179

lib/truffle/truffle/cext.rb

@@ -1734,11 +1734,15 @@ def rb_class_inherited_p(ruby_module, object)
   end
 
   def rb_backref_get
-    Primitive.regexp_last_match_get(Truffle::ThreadOperations.ruby_caller_special_variables([Truffle::CExt, Truffle::Interop.singleton_class]))
+    vars = Primitive.cext_special_variables_from_stack
+    vars = Truffle::ThreadOperations.ruby_caller_special_variables([Truffle::CExt, Truffle::Interop.singleton_class]) if vars.nil?
+    Primitive.regexp_last_match_get(vars)
   end
 
   def rb_backref_set(value)
-    Truffle::RegexpOperations::LAST_MATCH_SET.call(value, Truffle::ThreadOperations.ruby_caller_special_variables([Truffle::CExt, Truffle::Interop.singleton_class]))
+    vars = Primitive.cext_special_variables_from_stack
+    vars = Truffle::ThreadOperations.ruby_caller_special_variables([Truffle::CExt, Truffle::Interop.singleton_class]) if vars.nil?
+    Truffle::RegexpOperations::LAST_MATCH_SET.call(value, vars)
   end
 
   def rb_gv_set(name, value)

lib/truffle/truffle/cext_ruby.rb

@@ -38,7 +38,7 @@ def rb_define_method(mod, name, function, argc)
       # We must set block argument if given here so that the
       # `rb_block_*` functions will be able to find it by walking the
       # stack.
-      res = Primitive.cext_unwrap(Primitive.call_with_c_mutex_and_frame(function, args, block))
+      res = Primitive.cext_unwrap(Primitive.call_with_c_mutex_and_frame(function, args, Primitive.caller_special_variables_if_fast, block))
       Primitive.thread_set_exception(exc)
       res
     end

src/main/java/org/truffleruby/cext/CExtNodes.java

@@ -124,10 +124,10 @@ public abstract static class CallWithCExtLockAndFrameNode extends PrimitiveArray
         @Child protected CallWithCExtLockNode callCextNode = CallWithCExtLockNodeFactory.create(RubyNode.EMPTY_ARRAY);
 
         @Specialization
-        protected Object callWithCExtLockAndFrame(Object receiver, RubyArray argsArray, Object block,
+        protected Object callWithCExtLockAndFrame(Object receiver, RubyArray argsArray, Object variables, Object block,
                 @Cached MarkingServiceNodes.GetMarkerThreadLocalDataNode getDataNode) {
             final ExtensionCallStack extensionStack = getDataNode.execute().getExtensionCallStack();
-            extensionStack.push(block);
+            extensionStack.push(variables, block);
             try {
                 return callCextNode.execute(receiver, argsArray);
             } finally {
@@ -698,6 +698,16 @@ protected Object block(
         }
     }
 
+    @Primitive(name = "cext_special_variables_from_stack")
+    public abstract static class VarsFromStackNode extends PrimitiveArrayArgumentsNode {
+
+        @Specialization
+        protected Object variables(
+                @Cached MarkingServiceNodes.GetMarkerThreadLocalDataNode getDataNode) {
+            return getDataNode.execute().getExtensionCallStack().getVariables();
+        }
+    }
+
     @CoreMethod(names = "rb_check_frozen", onSingleton = true, required = 1)
     public abstract static class CheckFrozenNode extends CoreMethodArrayArgumentsNode {
 
@@ -1557,9 +1567,9 @@ protected Object createMarker(RubyDynamicObject object, RubyProc marker) {
     public abstract static class PushPreservingFrame extends CoreMethodArrayArgumentsNode {
 
         @Specialization
-        protected Object pushFrame(RubyProc block,
+        protected Object pushFrame(Object variables, RubyProc block,
                 @Cached MarkingServiceNodes.GetMarkerThreadLocalDataNode getDataNode) {
-            getDataNode.execute().getExtensionCallStack().push(block);
+            getDataNode.execute().getExtensionCallStack().push(variables, block);
             return nil;
         }
     }

src/main/java/org/truffleruby/core/FinalizationService.java

@@ -93,7 +93,7 @@ protected void processReference(RubyContext context, ProcessingReference<?> fina
 
     protected void processReferenceInternal(RubyContext context, FinalizerReference finalizerReference) {
         ExtensionCallStack stack = context.getMarkingService().getThreadLocalData().getExtensionCallStack();
-        stack.push(stack.getBlock());
+        stack.push(stack.getVariables(), stack.getBlock());
         try {
             while (!context.isFinalizing()) {
                 final Finalizer finalizer;

src/main/java/org/truffleruby/core/MarkingService.java

@@ -87,7 +87,7 @@ protected void processReference(RubyContext context, ProcessingReference<?> refe
         @TruffleBoundary
         public void runAllMarkers(RubyContext context) {
             ExtensionCallStack stack = markingService.getThreadLocalData().getExtensionCallStack();
-            stack.push(stack.getBlock());
+            stack.push(stack.getVariables(), stack.getBlock());
             try {
                 // TODO (eregon, 15 Sept 2020): there seems to be no synchronization here while walking the list of
                 // markingService, and concurrent mutations seem to be possible.
@@ -117,7 +117,7 @@ public static class MarkerThreadLocalData {
         private final ExtensionCallStack extensionCallStack;
 
         public MarkerThreadLocalData(MarkingService service) {
-            this.extensionCallStack = new ExtensionCallStack(Nil.INSTANCE);
+            this.extensionCallStack = new ExtensionCallStack(Nil.INSTANCE, Nil.INSTANCE);
         }
 
         public ExtensionCallStack getExtensionCallStack() {
@@ -129,18 +129,20 @@ protected static class ExtensionCallStackEntry {
         protected final ExtensionCallStackEntry previous;
         protected final ArrayList<Object> preservedObjects = new ArrayList<>();
         protected final Object block;
+        protected Object variables;
 
-        protected ExtensionCallStackEntry(ExtensionCallStackEntry previous, Object block) {
+        protected ExtensionCallStackEntry(ExtensionCallStackEntry previous, Object variables, Object block) {
             this.previous = previous;
             this.block = block;
+            this.variables = variables;
         }
     }
 
     public static class ExtensionCallStack {
         protected ExtensionCallStackEntry current;
 
-        public ExtensionCallStack(Object block) {
-            current = new ExtensionCallStackEntry(null, block);
+        public ExtensionCallStack(Object variables, Object block) {
+            current = new ExtensionCallStackEntry(null, variables, block);
         }
 
         public ArrayList<Object> getKeptObjects() {
@@ -151,8 +153,16 @@ public void pop() {
             current = current.previous;
         }
 
-        public void push(Object block) {
-            current = new ExtensionCallStackEntry(current, block);
+        public void push(Object variable, Object block) {
+            current = new ExtensionCallStackEntry(current, variable, block);
+        }
+
+        public Object getVariables() {
+            return current.variables;
+        }
+
+        public void setVariables(Object variables) {
+            current.variables = variables;
         }
 
         public Object getBlock() {

src/main/java/org/truffleruby/core/kernel/TruffleKernelNodes.java

@@ -38,6 +38,7 @@
 import org.truffleruby.language.RubyBaseNode;
 import org.truffleruby.language.RubyNode;
 import org.truffleruby.language.RubyRootNode;
+import org.truffleruby.language.arguments.MaybeReadCallerVariablesNode;
 import org.truffleruby.language.arguments.ReadCallerVariablesNode;
 import org.truffleruby.language.arguments.RubyArguments;
 import org.truffleruby.language.control.RaiseException;
@@ -346,6 +347,23 @@ protected Object storage(VirtualFrame frame) {
         }
     }
 
+    @Primitive(name = "caller_special_variables_if_fast")
+    public abstract static class GetCallerSpecialVariableStorageIfFast extends PrimitiveArrayArgumentsNode {
+
+        @Child MaybeReadCallerVariablesNode callerVariablesNode = new MaybeReadCallerVariablesNode();
+
+        @Specialization
+        protected Object storage(VirtualFrame frame,
+                @Cached ConditionProfile nullProfile) {
+            Object variables = callerVariablesNode.execute(frame);
+            if (nullProfile.profile(variables == null)) {
+                return nil;
+            } else {
+                return variables;
+            }
+        }
+    }
+
     /* When getting special variables from the wrong side of a C call we know it's going to be slow. */
     @Primitive(name = "ruby_caller_special_variables")
     public abstract static class GetSlowCallerSpecialVariableStorage extends PrimitiveArrayArgumentsNode {

src/main/java/org/truffleruby/core/thread/TruffleThreadNodes.java

@@ -16,6 +16,9 @@
 import org.truffleruby.core.array.RubyArray;
 import org.truffleruby.core.array.library.ArrayStoreLibrary;
 import org.truffleruby.core.kernel.TruffleKernelNodes.GetSpecialVariableStorage;
+import org.truffleruby.core.MarkingServiceNodes;
+import org.truffleruby.language.FrameAndVariablesSendingNode;
+import org.truffleruby.language.RubyContextNode;
 
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 import com.oracle.truffle.api.dsl.Cached;
@@ -24,10 +27,24 @@
 import com.oracle.truffle.api.frame.Frame;
 import com.oracle.truffle.api.frame.FrameInstance.FrameAccess;
 import com.oracle.truffle.api.library.CachedLibrary;
+import com.oracle.truffle.api.nodes.DirectCallNode;
+import com.oracle.truffle.api.nodes.IndirectCallNode;
+import com.oracle.truffle.api.nodes.Node;
 
 @CoreModule("Truffle::ThreadOperations")
 public class TruffleThreadNodes {
 
+    private static class FrameAndCallNode {
+
+        public final Frame frame;
+        public final Node callNode;
+
+        public FrameAndCallNode(Frame frame, Node callNode) {
+            this.frame = frame;
+            this.callNode = callNode;
+        }
+    }
+
     @CoreMethod(names = "ruby_caller_special_variables", onSingleton = true, required = 1)
     @ImportStatic(ArrayGuards.class)
     public abstract static class FindRubyCallerSpecialStorage extends CoreMethodArrayArgumentsNode {
@@ -36,18 +53,41 @@ public abstract static class FindRubyCallerSpecialStorage extends CoreMethodArra
         @Specialization(limit = "storageStrategyLimit()")
         protected Object findRubyCaller(RubyArray modules,
                 @CachedLibrary("modules.store") ArrayStoreLibrary stores,
-                @Cached GetSpecialVariableStorage storageNode) {
+                @Cached GetSpecialVariableStorage storageNode,
+                @Cached MarkingServiceNodes.GetMarkerThreadLocalDataNode getDataNode) {
             final int modulesSize = modules.size;
             Object[] moduleArray = stores.boxedCopyOfRange(modules.store, 0, modulesSize);
-            Frame rubyCaller = getContext()
+            FrameAndCallNode data = getContext()
                     .getCallStack()
-                    .getCallerFrameNotInModules(FrameAccess.MATERIALIZE, moduleArray);
-            if (rubyCaller == null) {
+                    .iterateFrameNotInModules(
+                            moduleArray,
+                            f -> new FrameAndCallNode(f.getFrame(FrameAccess.MATERIALIZE), f.getCallNode()));
+            if (data == null) {
                 return nil;
             } else {
-                return storageNode.execute(rubyCaller.materialize());
+                notifyToStartSendingStorage(data.callNode);
+                Object variables = storageNode.execute(data.frame.materialize());
+                getDataNode.execute().getExtensionCallStack().setVariables(variables);
+                return variables;
             }
         }
 
+        private static boolean notifyToStartSendingStorage(Node callerNode) {
+            if (callerNode instanceof DirectCallNode || callerNode instanceof IndirectCallNode) {
+                Node parent = callerNode.getParent();
+                while (parent != null) {
+                    if (parent instanceof FrameAndVariablesSendingNode) {
+                        ((FrameAndVariablesSendingNode) parent).startSendingOwnVariables();
+                        return true;
+                    }
+                    if (parent instanceof RubyContextNode) {
+                        return false;
+                    }
+                    parent = parent.getParent();
+                }
+            }
+
+            return false;
+        }
     }
 }

src/main/java/org/truffleruby/language/CallStackManager.java

@@ -90,6 +90,22 @@ public Frame getCallerFrameNotInModules(FrameAccess frameAccess, Object[] module
         }, frameAccess);
     }
 
+    @TruffleBoundary
+    public <R> R iterateFrameNotInModules(Object[] modules, Function<FrameInstance, R> action) {
+        final Memo<Boolean> skippedFirstFrameFound = new Memo<>(false);
+
+        return iterateFrames(1, frameInstance -> {
+            final InternalMethod method = tryGetMethod(frameInstance.getFrame(FrameAccess.READ_ONLY));
+            if (method != null && !ArrayUtils.contains(modules, method.getDeclaringModule())) {
+                if (skippedFirstFrameFound.get()) {
+                    return true;
+                }
+                skippedFirstFrameFound.set(true);
+            }
+            return false;
+        }, action);
+    }
+
     // Node
 
     @TruffleBoundary

src/main/java/org/truffleruby/language/arguments/MaybeReadCallerVariablesNode.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright (c) 2015, 2020 Oracle and/or its affiliates. All rights reserved. This
+ * code is released under a tri EPL/GPL/LGPL license. You can use it,
+ * redistribute it and/or modify it under the terms of the:
+ *
+ * Eclipse Public License version 2.0, or
+ * GNU General Public License version 2, or
+ * GNU Lesser General Public License version 2.1.
+ */
+package org.truffleruby.language.arguments;
+
+public class MaybeReadCallerVariablesNode extends ReadCallerVariablesNode {
+
+    public static MaybeReadCallerVariablesNode create() {
+        return new MaybeReadCallerVariablesNode();
+    }
+
+    @Override
+    protected Object getCallerData() {
+        return null;
+    }
+}

src/main/java/org/truffleruby/language/arguments/ReadCallerDataNode.java

@@ -45,7 +45,7 @@ public Object execute(Frame frame) {
     protected abstract Object getData(Frame frame);
 
     @TruffleBoundary
-    private Object getCallerData() {
+    protected Object getCallerData() {
         if (!notifyCallerToSendData()) {
             // If we fail to notify the call node (e.g., because it is a UncachedDispatchNode which is not handled yet),
             // we don't want to deoptimize this CallTarget on every call.