Refactor of rb_block handling.

Author: aardvark179

lib/truffle/truffle/cext.rb

@@ -962,7 +962,7 @@ def rb_funcallv(recv, meth, args)
 
   def rb_funcall(recv, meth, n, *args)
     # see #call_with_thread_locally_stored_block
-    thread_local_block           = Thread.current[:__C_BLOCK__]
+    thread_local_block = Thread.current[:__C_BLOCK__]
     Thread.current[:__C_BLOCK__] = nil
     recv.__send__(meth, *args, &thread_local_block)
   ensure

lib/truffle/truffle/cext_ruby.rb

@@ -30,7 +30,7 @@ def rb_define_method(mod, name, function, argc)
       end
 
       # Using raw execute instead of #call here to avoid argument conversion
-      Truffle::CExt.push_preserving_frame
+      Truffle::CExt.push_extension_call_frame(block)
       begin
         # We must call explicitly with the block argument if given
         # here so that the `rb_block_*` functions will be able to find
@@ -41,7 +41,7 @@ def rb_define_method(mod, name, function, argc)
           Truffle::CExt.rb_tr_unwrap(Truffle.invoke_primitive(:call_with_c_mutex, function, args))
         end
       ensure
-        Truffle::CExt.pop_preserving_frame
+        Truffle::CExt.pop_extension_call_frame
       end
     end
 
@@ -51,19 +51,27 @@ def rb_define_method(mod, name, function, argc)
   private
 
   def rb_iterate_call_block(callback, block_arg, callback_arg, &block)
-    rb_tr_unwrap(Truffle.invoke_primitive(:call_with_c_mutex, callback, [rb_tr_wrap(block_arg), rb_tr_wrap(callback_arg)]))
+    Truffle::CExt.push_extension_call_frame(block)
+    begin
+      rb_tr_unwrap(Truffle.invoke_primitive(:call_with_c_mutex, callback, [rb_tr_wrap(block_arg), rb_tr_wrap(callback_arg)]))
+    ensure
+      Truffle::CExt.pop_extension_call_frame
+    end
   end
 
   def call_with_thread_locally_stored_block(function, *args, &block)
     # MRI puts the block to a thread local th->passed_block and later rb_funcall reads it,
     # we have to do the same
     # TODO (pitr-ch 14-Dec-2017): This is fixed just for rb_iterate with a rb_funcall in it combination
+
     previous_block = Thread.current[:__C_BLOCK__]
+    Truffle::CExt.push_extension_call_frame(block)
     begin
       Thread.current[:__C_BLOCK__] = block
       rb_tr_unwrap(Truffle.invoke_primitive(:call_with_c_mutex, function, args.map! { |arg| Truffle::CExt.rb_tr_wrap(arg) }, &block))
     ensure
       Thread.current[:__C_BLOCK__] = previous_block
+      Truffle::CExt.pop_extension_call_frame
     end
   end
 

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

@@ -590,31 +590,11 @@ public DynamicObject rbStrResize(DynamicObject string, int newByteLength,
     @CoreMethod(names = "rb_block_proc", onSingleton = true)
     public abstract static class BlockProcNode extends CoreMethodArrayArgumentsNode {
 
-        // TODO (pitr-ch 04-Dec-2017): needs optimising
-        @TruffleBoundary
         @Specialization
-        public DynamicObject blockProc() {
-            return Truffle.getRuntime().iterateFrames(frameInstance -> {
-                final Node callNode = frameInstance.getCallNode();
-
-                if (callNode != null) {
-                    final RootNode rootNode = callNode.getRootNode();
-                    // Skip Ruby frames in cext.rb file since they are implementing methods which are implemented
-                    // with C in MRI, and therefore are also implicitly skipped when when looking up the block passed
-                    // to a C API function.
-                    if (rootNode instanceof RubyRootNode &&
-                            rootNode.getSourceSection().isAvailable() &&
-                            !rootNode.getSourceSection().getSource().getName().endsWith("truffle/cext.rb")) {
-
-                        final DynamicObject block = RubyArguments.getBlock(frameInstance.getFrame(FrameAccess.READ_ONLY));
-                        return block == null ? nil() : block;
-                    }
-                }
-
-                return null;
-            });
+        public DynamicObject block(VirtualFrame frame,
+                @Cached("create()") MarkingServiceNodes.GetMarkerThreadLocalDataNode getDataNode) {
+            return getDataNode.execute(frame).getExtensionCallStack().getBlock();
         }
-
     }
 
     @CoreMethod(names = "rb_check_frozen", onSingleton = true, required = 1)
@@ -1520,24 +1500,24 @@ protected void addObjectToMarkingService(DynamicObject object, DynamicObject mar
         }
     }
 
-    @CoreMethod(names = "push_preserving_frame", onSingleton = true, required = 0)
+    @CoreMethod(names = "push_extension_call_frame", onSingleton = true, required = 1)
     public abstract static class PushPreservingFrame extends CoreMethodArrayArgumentsNode {
 
         @Specialization
-        public DynamicObject pushFrame(VirtualFrame frame,
+        public DynamicObject pushFrame(VirtualFrame frame, DynamicObject block,
                 @Cached("create()") MarkingServiceNodes.GetMarkerThreadLocalDataNode getDataNode) {
-            getDataNode.execute(frame).getPreservationStack().push();
+            getDataNode.execute(frame).getExtensionCallStack().push(block);
             return nil();
         }
     }
 
-    @CoreMethod(names = "pop_preserving_frame", onSingleton = true, required = 0)
+    @CoreMethod(names = "pop_extension_call_frame", onSingleton = true, required = 0)
     public abstract static class PopPreservingFrame extends CoreMethodArrayArgumentsNode {
 
         @Specialization
         public DynamicObject popFrame(VirtualFrame frame,
                 @Cached("create()") MarkingServiceNodes.GetMarkerThreadLocalDataNode getDataNode) {
-            getDataNode.execute(frame).getPreservationStack().pop();
+            getDataNode.execute(frame).getExtensionCallStack().pop();
             return nil();
         }
     }

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

@@ -126,19 +126,19 @@ public void runAllMarkers() {
 
     public static class MarkerThreadLocalData {
         private final MarkerKeptObjects keptObjects;
-        private final MarkerStack preservationStack;
+        private final ExtensionCallStack extensionCallStack;
 
         public MarkerThreadLocalData(MarkingService service) {
-            this.preservationStack = new MarkerStack();
+            this.extensionCallStack = new ExtensionCallStack(service.context.getCoreLibrary().getNil());
             this.keptObjects = new MarkerKeptObjects(service);
         }
 
         public MarkerKeptObjects getKeptObjects() {
             return keptObjects;
         }
 
-        public MarkerStack getPreservationStack() {
-            return preservationStack;
+        public ExtensionCallStack getExtensionCallStack() {
+            return extensionCallStack;
         }
     }
 
@@ -226,28 +226,38 @@ public void keepObjects(MarkerKeptObjects otherObjects) {
 
     }
 
-    protected static class MarkerStackEntry {
-        protected final MarkerStackEntry previous;
-        protected final ArrayList<Object> entries = new ArrayList<>();
+    protected static class ExtensionCallStackEntry {
+        protected final ExtensionCallStackEntry previous;
+        protected final ArrayList<Object> preservedObjects = new ArrayList<>();
+        protected final DynamicObject block;
 
-        protected MarkerStackEntry(MarkerStackEntry previous) {
+        protected ExtensionCallStackEntry(ExtensionCallStackEntry previous, DynamicObject block) {
             this.previous = previous;
+            this.block = block;
         }
     }
 
-    public static class MarkerStack {
-        protected MarkerStackEntry current = new MarkerStackEntry(null);
+    public static class ExtensionCallStack {
+        protected ExtensionCallStackEntry current;
 
-        public ArrayList<Object> get() {
-            return current.entries;
+        public ExtensionCallStack(DynamicObject block) {
+            current  = new ExtensionCallStackEntry(null, block);
+        }
+
+        public ArrayList<Object> getKeptObjects() {
+            return current.preservedObjects;
         }
 
         public void pop() {
             current = current.previous;
         }
 
-        public void push() {
-            current = new MarkerStackEntry(current);
+        public void push(DynamicObject block) {
+            current = new ExtensionCallStackEntry(current, block);
+        }
+
+        public DynamicObject getBlock() {
+            return current.block;
         }
     }
 

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

@@ -29,7 +29,7 @@ public static abstract class KeepAliveNode extends RubyBaseNode {
         public void keepObjectAlive(VirtualFrame frame, Object object,
                 @Cached("create()") GetMarkerThreadLocalDataNode getThreadLocalDataNode) {
             MarkerThreadLocalData data = getThreadLocalDataNode.execute(frame);
-            addToList(data.getPreservationStack().get(), object);
+            addToList(data.getExtensionCallStack().getKeptObjects(), object);
             data.getKeptObjects().keepObject(object);
         }