Replace ary_fetch_next with separate array iteration primitives

2 files changed, 34 insertions, 25 deletions

diff --git a/array.c b/array.c
index 90e83fd3ad..876642dff7 100644
--- a/array.c
+++ b/array.c
@@ -2694,18 +2694,25 @@ ary_enum_length(VALUE ary, VALUE args, VALUE eobj)
     return rb_ary_length(ary);
 }
 
-// Primitive to avoid a race condition in Array#each.
-// Return `true` and write `value` and `index` if the element exists.
+// Return true if the index is at or past the end of the array.
 static VALUE
-ary_fetch_next(VALUE self, VALUE *index, VALUE *value)
+rb_jit_ary_at_end(rb_execution_context_t *ec, VALUE self, VALUE index)
 {
-    long i = NUM2LONG(*index);
-    if (i >= RARRAY_LEN(self)) {
-        return Qfalse;
-    }
-    *value = RARRAY_AREF(self, i);
-    *index = LONG2NUM(i + 1);
-    return Qtrue;
+    return FIX2LONG(index) >= RARRAY_LEN(self) ? Qtrue : Qfalse;
+}
+
+// Return the element at the given fixnum index.
+static VALUE
+rb_jit_ary_at(rb_execution_context_t *ec, VALUE self, VALUE index)
+{
+    return RARRAY_AREF(self, FIX2LONG(index));
+}
+
+// Increment a fixnum by 1.
+static VALUE
+rb_jit_fixnum_inc(rb_execution_context_t *ec, VALUE self, VALUE num)
+{
+    return LONG2FIX(FIX2LONG(num) + 1);
 }
 
 /*
diff --git a/array.rb b/array.rb
index 4fee5ecaeb..e21de1a5b2 100644
--- a/array.rb
+++ b/array.rb
@@ -222,10 +222,10 @@ class Array
         unless defined?(yield)
           return Primitive.cexpr! 'SIZED_ENUMERATOR(self, 0, 0, ary_enum_length)'
         end
-        _i = 0
-        value = nil
-        while Primitive.cexpr!(%q{ ary_fetch_next(self, LOCAL_PTR(_i), LOCAL_PTR(value)) })
-          yield value
+        i = 0
+        until Primitive.rb_jit_ary_at_end(i)
+          yield Primitive.rb_jit_ary_at(i)
+          i = Primitive.rb_jit_fixnum_inc(i)
         end
         self
       end
@@ -241,12 +241,12 @@ class Array
           return Primitive.cexpr! 'SIZED_ENUMERATOR(self, 0, 0, ary_enum_length)'
         end
 
-        _i = 0
-        value = nil
+        i = 0
         result = Primitive.ary_sized_alloc
-        while Primitive.cexpr!(%q{ ary_fetch_next(self, LOCAL_PTR(_i), LOCAL_PTR(value)) })
-          value = yield(value)
-          Primitive.cexpr!(%q{ rb_ary_push(result, value) })
+        until Primitive.rb_jit_ary_at_end(i)
+          _value = yield(Primitive.rb_jit_ary_at(i))
+          Primitive.cexpr!(%q{ rb_ary_push(result, _value) })
+          i = Primitive.rb_jit_fixnum_inc(i)
         end
         result
       end
@@ -267,13 +267,14 @@ class Array
           return Primitive.cexpr! 'SIZED_ENUMERATOR(self, 0, 0, ary_enum_length)'
         end
 
-        _i = 0
-        value = nil
+        i = 0
         result = Primitive.ary_sized_alloc
-        while Primitive.cexpr!(%q{ ary_fetch_next(self, LOCAL_PTR(_i), LOCAL_PTR(value)) })
+        until Primitive.rb_jit_ary_at_end(i)
+          value = Primitive.rb_jit_ary_at(i)
           if yield value
             Primitive.cexpr!(%q{ rb_ary_push(result, value) })
           end
+          i = Primitive.rb_jit_fixnum_inc(i)
         end
         result
       end
@@ -293,10 +294,11 @@ class Array
         unless defined?(yield)
           return Primitive.cexpr! 'SIZED_ENUMERATOR(self, 0, 0, ary_enum_length)'
         end
-        _i = 0
-        value = nil
-        while Primitive.cexpr!(%q{ ary_fetch_next(self, LOCAL_PTR(_i), LOCAL_PTR(value)) })
+        i = 0
+        until Primitive.rb_jit_ary_at_end(i)
+          value = Primitive.rb_jit_ary_at(i)
           return value if yield(value)
+          i = Primitive.rb_jit_fixnum_inc(i)
         end
         if_none_proc&.call
       end