rewrite flatten_array_helper in java

Author: Nicolas Laurent

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

@@ -1049,7 +1049,6 @@ public boolean isTruffleBootMainMethod(SharedMethodInfo info) {
             "/core/truffle/internal.rb",
             "/core/kernel.rb",
             "/core/lazy_rubygems.rb",
-            "/core/truffle/array_operations.rb",
             "/core/truffle/boot.rb",
             "/core/truffle/debug.rb",
             "/core/truffle/encoding_operations.rb",

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

@@ -13,10 +13,13 @@
 import static org.truffleruby.core.array.ArrayHelpers.setStoreAndSize;
 import static org.truffleruby.language.dispatch.DispatchNode.PUBLIC;
 
+import java.util.ArrayDeque;
 import java.util.Arrays;
 
 import com.oracle.truffle.api.dsl.CachedLanguage;
 import com.oracle.truffle.api.profiles.LoopConditionProfile;
+import org.graalvm.collections.EconomicSet;
+import org.graalvm.collections.Equivalence;
 import org.truffleruby.Layouts;
 import org.truffleruby.RubyLanguage;
 import org.truffleruby.builtins.CoreMethod;
@@ -2338,4 +2341,99 @@ protected boolean isStoreNative(RubyArray array,
             return stores.isNative(array.store);
         }
     }
+
+    @Primitive(name = "array_flatten_helper", lowerFixnum = 2)
+    public abstract static class FlattenHelperNode extends PrimitiveArrayArgumentsNode {
+
+        @Specialization(guards = "!canContainObject.execute(array)")
+        protected boolean flattenHelperPrimitive(RubyArray array, RubyArray out, int maxLevels,
+                @Cached ArrayAppendManyNode concat,
+                @Cached TypeNodes.CanContainObjectNode canContainObject) {
+            concat.executeAppendMany(out, array);
+            return false;
+        }
+
+        @Specialization(replaces = "flattenHelperPrimitive")
+        protected boolean flattenHelper(RubyArray array, RubyArray out, int maxLevels,
+                @CachedLanguage RubyLanguage language,
+                @Cached TypeNodes.CanContainObjectNode canContainObject,
+                @Cached ArrayAppendManyNode concat,
+                @Cached AtNode at,
+                @Cached DispatchNode convert,
+                @Cached ArrayAppendOneNode append) {
+
+            boolean modified = false;
+            final EconomicSet<RubyArray> visited = EconomicSet.create(Equivalence.IDENTITY);
+            // final HashSet<RubyArray> visited = new HashSet<>();
+            class Entry {
+                final RubyArray array;
+                final int index;
+
+                Entry(RubyArray array, int index) {
+                    this.array = array;
+                    this.index = index;
+                }
+            }
+            final ArrayDeque<Entry> workStack = new ArrayDeque<>();
+            workStack.push(new Entry(array, 0));
+
+            while (!workStack.isEmpty()) {
+                final Entry e = workStack.pop();
+
+                if (e.index == 0) {
+                    if (!canContainObject.execute(e.array)) {
+                        concat.executeAppendMany(out, e.array);
+                        continue;
+                    } else if (contains(visited, e.array)) {
+                        throw new RaiseException(
+                                getContext(),
+                                coreExceptions().argumentError("tried to flatten recursive array", this));
+                    } else if (maxLevels == workStack.size()) {
+                        concat.executeAppendMany(out, e.array);
+                        continue;
+                    }
+                    add(visited, e.array);
+                }
+
+                int i = e.index;
+                for (; i < e.array.size; ++i) {
+                    final Object obj = at.executeAt(e.array, i);
+                    final Object converted = convert.call(
+                            coreLibrary().truffleTypeModule,
+                            "rb_check_convert_type",
+                            obj,
+                            coreLibrary().arrayClass,
+                            language.coreSymbols.TO_ARY);
+                    if (converted == nil) {
+                        append.executeAppendOne(out, obj);
+                    } else {
+                        modified = true;
+                        workStack.push(new Entry(e.array, i + 1));
+                        workStack.push(new Entry((RubyArray) converted, 0));
+                        break;
+                    }
+                }
+                if (i == e.array.size) {
+                    remove(visited, e.array);
+                }
+            }
+
+            return modified;
+        }
+
+        @TruffleBoundary
+        private static boolean contains(EconomicSet<RubyArray> set, RubyArray array) {
+            return set.contains(array);
+        }
+
+        @TruffleBoundary
+        private static void remove(EconomicSet<RubyArray> set, RubyArray array) {
+            set.remove(array);
+        }
+
+        @TruffleBoundary
+        private static void add(EconomicSet<RubyArray> set, RubyArray array) {
+            set.add(array);
+        }
+    }
 }

src/main/java/org/truffleruby/core/support/TypeNodes.java

@@ -261,6 +261,12 @@ protected Object objectHiddenVarSet(RubyDynamicObject object, Object identifier,
     @ImportStatic(ArrayGuards.class)
     public abstract static class CanContainObjectNode extends PrimitiveArrayArgumentsNode {
 
+        public static CanContainObjectNode create() {
+            return TypeNodesFactory.CanContainObjectNodeFactory.create(null);
+        }
+
+        abstract public boolean execute(RubyArray array);
+
         @Specialization(
                 guards = {
                         "stores.accepts(array.store)",

src/main/ruby/truffleruby/core/array.rb

@@ -444,7 +444,7 @@ def flatten(level=-1)
     return self.dup if level == 0
 
     out = self.class.allocate # new_reserved size
-    Truffle::ArrayOperations.flatten_helper(self, out, level)
+    Primitive.array_flatten_helper(self, out, level)
     Primitive.infect(out, self)
     out
   end
@@ -456,7 +456,7 @@ def flatten!(level=-1)
     return nil if level == 0
 
     out = self.class.allocate # new_reserved size
-    if Truffle::ArrayOperations.flatten_helper(self, out, level)
+    if Primitive.array_flatten_helper(self, out, level)
       Primitive.steal_array_storage(self, out)
       return self
     end