Rewrite Array#each in Ruby using Primitive

ruby · k0kubun · Jan 23, 2024 · Jan 12, 2024 · Jan 21, 2024 · Jan 23, 2024
commit 3aaa6de3ec1476ceb9ce1b46295183b1523f9937
diff --git a/array.c b/array.c
@@ -28,6 +28,7 @@
 #include "ruby/encoding.h"
 #include "ruby/st.h"
 #include "ruby/util.h"
+#include "vm_core.h"
 #include "builtin.h"
 
 #if !ARRAY_DEBUG
@@ -2484,50 +2485,19 @@ ary_enum_length(VALUE ary, VALUE args, VALUE eobj)
     return rb_ary_length(ary);
 }
 
-/*
- *  call-seq:
- *    array.each {|element| ... } -> self
- *    array.each -> Enumerator
- *
- *  Iterates over array elements.
- *
- *  When a block given, passes each successive array element to the block;
- *  returns +self+:
- *
- *    a = [:foo, 'bar', 2]
- *    a.each {|element|  puts "#{element.class} #{element}" }
- *
- *  Output:
- *
- *    Symbol foo
- *    String bar
- *    Integer 2
- *
- *  Allows the array to be modified during iteration:
- *
- *    a = [:foo, 'bar', 2]
- *    a.each {|element| puts element; a.clear if element.to_s.start_with?('b') }
- *
- *  Output:
- *
- *    foo
- *    bar
- *
- *  When no block given, returns a new Enumerator:
- *    a = [:foo, 'bar', 2]
- *
- *    e = a.each
- *    e # => #<Enumerator: [:foo, "bar", 2]:each>
- *    a1 = e.each {|element|  puts "#{element.class} #{element}" }
- *
- *  Output:
- *
- *    Symbol foo
- *    String bar
- *    Integer 2
- *
- *  Related: #each_index, #reverse_each.
- */
+// Primitive to avoid a race condition in Array#each.
+// Return `true` and write `value` and `index` if the element exists.
+static VALUE
+ary_fetch_next(VALUE self, VALUE *index, VALUE *value)
+{
+    long i = NUM2LONG(*index);
+    if (i >= RARRAY_LEN(self)) {
+        return Qfalse;
+    }
+    *value = RARRAY_AREF(self, i);
+    *index = LONG2NUM(i + 1);
+    return Qtrue;
+}
 
 VALUE
 rb_ary_each(VALUE ary)
@@ -8644,7 +8614,6 @@ Init_Array(void)
     rb_define_method(rb_cArray, "unshift", rb_ary_unshift_m, -1);
     rb_define_alias(rb_cArray,  "prepend", "unshift");
     rb_define_method(rb_cArray, "insert", rb_ary_insert, -1);
-    rb_define_method(rb_cArray, "each", rb_ary_each, 0);
     rb_define_method(rb_cArray, "each_index", rb_ary_each_index, 0);
     rb_define_method(rb_cArray, "reverse_each", rb_ary_reverse_each, 0);
     rb_define_method(rb_cArray, "length", rb_ary_length, 0);

diff --git a/array.rb b/array.rb
@@ -1,4 +1,58 @@
 class Array
+  # call-seq:
+  #   array.each {|element| ... } -> self
+  #   array.each -> Enumerator
+  #
+  # Iterates over array elements.
+  #
+  # When a block given, passes each successive array element to the block;
+  # returns +self+:
+  #
+  #   a = [:foo, 'bar', 2]
+  #   a.each {|element|  puts "#{element.class} #{element}" }
+  #
+  # Output:
+  #
+  #   Symbol foo
+  #   String bar
+  #   Integer 2
+  #
+  # Allows the array to be modified during iteration:
+  #
+  #   a = [:foo, 'bar', 2]
+  #   a.each {|element| puts element; a.clear if element.to_s.start_with?('b') }
+  #
+  # Output:
+  #
+  #   foo
+  #   bar
+  #
+  # When no block given, returns a new Enumerator:
+  #   a = [:foo, 'bar', 2]
+  #
+  #   e = a.each
+  #   e # => #<Enumerator: [:foo, "bar", 2]:each>
+  #   a1 = e.each {|element|  puts "#{element.class} #{element}" }
+  #
+  # Output:
+  #
+  #   Symbol foo
+  #   String bar
+  #   Integer 2
+  #
+  # Related: #each_index, #reverse_each.
+  def each
+    unless defined?(yield)
+      return to_enum(:each) { self.length }
+    end
+    _i = 0
+    value = nil
+    while Primitive.cexpr!(%q{ ary_fetch_next(self, LOCAL_PTR(_i), LOCAL_PTR(value)) })
+      yield value
+    end
+    self
+  end
+
   # call-seq:
   #    array.shuffle!(random: Random) -> array
   #

diff --git a/benchmark/loop_each.yml b/benchmark/loop_each.yml
@@ -0,0 +1,4 @@
+prelude: |
+  arr = [nil] * 30_000_000
+benchmark:
+  loop_each: arr.each{|e|}
diff --git a/builtin.h b/builtin.h
@@ -106,6 +106,8 @@ rb_vm_lvar(rb_execution_context_t *ec, int index)
 #endif
 }
 
+#define LOCAL_PTR(local) local ## __ptr
+
 // dump/load
 
 struct builtin_binary {

diff --git a/gems/bundled_gems b/gems/bundled_gems
@@ -19,7 +19,7 @@ matrix          0.4.2   https://github.com/ruby/matrix
 prime           0.1.2   https://github.com/ruby/prime
 rbs             3.4.2   https://github.com/ruby/rbs
 typeprof        0.21.9  https://github.com/ruby/typeprof
-debug           1.9.1   https://github.com/ruby/debug
+debug           1.9.1   https://github.com/ruby/debug 19b91b14ce814a0eb615abb8d2bef0594c61c5c8
 racc            1.7.3   https://github.com/ruby/racc
 mutex_m         0.2.0   https://github.com/ruby/mutex_m
 getoptlong      0.2.1   https://github.com/ruby/getoptlong

diff --git a/test/ruby/test_gc_compact.rb b/test/ruby/test_gc_compact.rb
@@ -373,7 +373,7 @@ def add_ivars
 
       Fiber.new {
         $ary = OBJ_COUNT.times.map { Foo.new }
-        $ary.each(&:add_ivars)
+        $ary.reverse_each(&:add_ivars)
 
         GC.start
         Foo.new.add_ivars

diff --git a/test/ruby/test_settracefunc.rb b/test/ruby/test_settracefunc.rb
@@ -504,7 +504,7 @@ def trace_by_tracepoint *trace_events
     1: trace = TracePoint.trace(*trace_events){|tp| next if !target_thread?
     2:   events << [tp.event, tp.lineno, tp.path, _defined_class.(tp), tp.method_id, tp.self, tp.binding&.eval("_local_var"), _get_data.(tp)] if tp.path == 'xyzzy'
     3: }
-    4: [1].each{|;_local_var| _local_var = :inner
+    4: [1].reverse_each{|;_local_var| _local_var = :inner
     5:   tap{}
     6: }
     7: class XYZZY
@@ -531,10 +531,10 @@ def trace_by_tracepoint *trace_events
     answer_events = [
      #
      [:line,     4, 'xyzzy', self.class,  method,           self,        :outer, :nothing],
-     [:c_call,   4, 'xyzzy', Array,       :each,            [1],         nil,    :nothing],
+     [:c_call,   4, 'xyzzy', Array,       :reverse_each,    [1],         nil,    :nothing],
      [:line,     4, 'xyzzy', self.class,  method,           self,        nil,    :nothing],
      [:line,     5, 'xyzzy', self.class,  method,           self,        :inner, :nothing],
-     [:c_return, 4, "xyzzy", Array,       :each,            [1],         nil, [1]],
+     [:c_return, 4, "xyzzy", Array,       :reverse_each,    [1],         nil, [1]],
      [:line,     7, 'xyzzy', self.class,  method,           self,        :outer, :nothing],
      [:c_call,   7, "xyzzy", Module,      :const_added,     TestSetTraceFunc, nil, :nothing],
      [:c_return, 7, "xyzzy", Module,      :const_added,     TestSetTraceFunc, nil, nil],

diff --git a/tool/mk_builtin_loader.rb b/tool/mk_builtin_loader.rb
@@ -274,7 +274,8 @@ def generate_cexpr(ofile, lineno, line_file, body_lineno, text, locals, func_nam
   locals&.reverse_each&.with_index{|param, i|
     next unless Symbol === param
     next unless local_candidates.include?(param.to_s)
-    f.puts "MAYBE_UNUSED(const VALUE) #{param} = rb_vm_lvar(ec, #{-3 - i});"
+    f.puts "VALUE *const #{param}__ptr = (VALUE *)&ec->cfp->ep[#{-3 - i}];"
+    f.puts "MAYBE_UNUSED(const VALUE) #{param} = *#{param}__ptr;"
     lineno += 1
   }
   f.puts "#line #{body_lineno} \"#{line_file}\""
-Original file line number
+Diff line change
@@ Expand Up / @@ -106,6 +106,8 @@ rb_vm_lvar(rb_execution_context_t *ec, int index) @@
     #endif
     }
+    #define LOCAL_PTR(local) local ## __ptr
     // dump/load
     struct builtin_binary {
@@ Expand Down @@