[GR-19220] Fix crash caused by mutating an array within the .each call block (#2587)

PullRequest: truffleruby/3170

Author: eregon

CHANGELOG.md

@@ -28,6 +28,7 @@ Compatibility:
 * Set `@gem_prelude_index` variable on the default load paths (#2586 , @bjfish)
 * Do not call `IO#flush` dynamically from `IO#close` (#2594, @gogainda).
 * Implement `rb_str_new_static` for C extensions that use it (@aardvark179).
+* Rewrote `ArrayEachIteratorNode` and re-introduced `each` specs for MRI parity when mutating arrays whilst iterating, rather than crashing (#2587, @MattAlp)
 
 Performance:
 
@@ -36,6 +37,7 @@ Performance:
 * Removed extra array allocations for method calls in the interpreter to improve warmup performance (@aardvark179).
 * Optimize `Dir[]` by sorting entries as they are found and grouping syscalls (#2092, @aardvark179).
 * Reduce memory footprint by tracking `VALUE`s created during C extension init separately (@aardvark179).
+* Rewrote `ArrayEachIteratorNode` to optimize performance for a constant-sized array and reduce specializations to 1 general case (#2587, @MattAlp)
 
 Changes:
 

spec/ruby/core/array/each_spec.rb

@@ -3,9 +3,10 @@
 require_relative 'shared/enumeratorize'
 require_relative '../enumerable/shared/enumeratorized'
 
-# Modifying a collection while the contents are being iterated
-# gives undefined behavior. See
-# http://blade.nagaokaut.ac.jp/cgi-bin/scat.rb/ruby/ruby-core/23633
+# Mutating the array while it is being iterated is discouraged as it can result in confusing behavior.
+# Yet a Ruby implementation must not crash in such a case, and following the simple CRuby behavior makes sense.
+# CRuby simply reads the array storage and checks the size for every iteration;
+# like `i = 0; while i < size; yield self[i]; end`
 
 describe "Array#each" do
   it "yields each element to the block" do
@@ -15,6 +16,34 @@
     a.should == [1, 2, 3]
   end
 
+  it "yields each element to the block even if the array is changed during iteration" do
+    a = [1, 2, 3, 4, 5]
+    iterated = []
+    a.each { |x| iterated << x; a << x+5 if x.even? }
+    iterated.should == [1, 2, 3, 4, 5, 7, 9]
+  end
+
+  it "yields only elements that are still in the array" do
+    a = [0, 1, 2, 3, 4]
+    iterated = []
+    a.each { |x| iterated << x; a.pop if x.even? }
+    iterated.should == [0, 1, 2]
+  end
+
+  it "yields elements based on an internal index" do
+    a = [0, 1, 2, 3, 4]
+    iterated = []
+    a.each { |x| iterated << x; a.shift if x.even? }
+    iterated.should == [0, 2, 4]
+  end
+
+  it "yields the same element multiple times if inserting while iterating" do
+    a = [1, 2]
+    iterated = []
+    a.each { |x| iterated << x; a.unshift(0) if a.size == 2 }
+    iterated.should == [1, 1, 2]
+  end
+
   it "yields each element to a block that takes multiple arguments" do
     a = [[1, 2], :a, [3, 4]]
     b = []

src/main/java/org/truffleruby/core/array/ArrayEachIteratorNode.java

@@ -10,12 +10,12 @@
 package org.truffleruby.core.array;
 
 import com.oracle.truffle.api.TruffleSafepoint;
+import com.oracle.truffle.api.profiles.IntValueProfile;
 import org.truffleruby.core.array.library.ArrayStoreLibrary;
 import org.truffleruby.core.proc.RubyProc;
 import org.truffleruby.language.RubyBaseNode;
 
 import com.oracle.truffle.api.CompilerDirectives;
-import com.oracle.truffle.api.dsl.Bind;
 import com.oracle.truffle.api.dsl.Cached;
 import com.oracle.truffle.api.dsl.ImportStatic;
 import com.oracle.truffle.api.dsl.ReportPolymorphism;
@@ -42,35 +42,19 @@ public static ArrayEachIteratorNode create() {
     public abstract RubyArray execute(RubyArray array, RubyProc block, int startAt,
             ArrayElementConsumerNode consumerNode);
 
-    @Specialization(
-            guards = { "array.size == 1", "startAt == 0" },
-            limit = "storageStrategyLimit()")
-    protected RubyArray iterateOne(RubyArray array, RubyProc block, int startAt, ArrayElementConsumerNode consumerNode,
-            @Bind("array.store") Object store,
-            @CachedLibrary("store") ArrayStoreLibrary stores) {
-
-        consumerNode.accept(array, block, stores.read(store, 0), 0);
-
-        if (array.size > 1) {
-            // Implicitly profiles through lazy node creation
-            return getRecurseNode().execute(array, block, 1, consumerNode);
-        }
-
-        return array;
-    }
-
-    @Specialization(
-            guards = { "array.size != 1" },
-            limit = "storageStrategyLimit()")
+    @Specialization(limit = "storageStrategyLimit()")
     protected RubyArray iterateMany(RubyArray array, RubyProc block, int startAt, ArrayElementConsumerNode consumerNode,
-            @Bind("array.store") Object store,
-            @CachedLibrary("store") ArrayStoreLibrary stores,
+            // Checkstyle: stop -- Verified @Bind is not necessary here due to using `Library#accepts()`.
+            @CachedLibrary("array.store") ArrayStoreLibrary stores,
+            // Checkstyle: resume
             @Cached LoopConditionProfile loopProfile,
+            @Cached("createIdentityProfile()") IntValueProfile arraySizeProfile,
             @Cached ConditionProfile strategyMatchProfile) {
         int i = startAt;
         try {
-            for (; loopProfile.inject(i < array.size); i++) {
-                if (strategyMatchProfile.profile(stores.accepts(array.store))) {
+            for (; loopProfile.inject(i < arraySizeProfile.profile(array.size)); i++) {
+                Object store = array.store;
+                if (strategyMatchProfile.profile(stores.accepts(store))) {
                     consumerNode.accept(array, block, stores.read(store, i), i);
                 } else {
                     return getRecurseNode().execute(array, block, i, consumerNode);