[GR-26883] Rewrite Array#recursively_flatten to not use Thread#detect_recursion and reduce splitting.

PullRequest: truffleruby/2054

Author: Nicolas Laurent

bench/micro/array/flatten-recursive.rb

@@ -0,0 +1,19 @@
+# Copyright (c) 2020 Oracle and/or its affiliates. All rights reserved. This
+# code is released under a tri EPL/GPL/LGPL license. You can use it,
+# redistribute it and/or modify it under the terms of the:
+#
+# Eclipse Public License version 2.0, or
+# GNU General Public License version 2, or
+# GNU Lesser General Public License version 2.1.
+
+array = []
+current_array = array
+100.times do
+  next_array = []
+  current_array.append(0, 1, next_array, 3, 4)
+  current_array = next_array
+end
+
+benchmark 'core-array-flatten-recursive' do
+  array.flatten
+end

bench/micro/array/flatten-simple.rb

@@ -0,0 +1,13 @@
+# Copyright (c) 2020 Oracle and/or its affiliates. All rights reserved. This
+# code is released under a tri EPL/GPL/LGPL license. You can use it,
+# redistribute it and/or modify it under the terms of the:
+#
+# Eclipse Public License version 2.0, or
+# GNU General Public License version 2, or
+# GNU Lesser General Public License version 2.1.
+
+array = [ [[1, 2], [3,4]], [[5,6], [7,8]] ]
+
+benchmark 'core-array-flatten-simple' do
+  array.flatten
+end

src/main/java/org/truffleruby/collections/SimpleStack.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright (c) 2020 Oracle and/or its affiliates. All rights reserved. This
+ * code is released under a tri EPL/GPL/LGPL license. You can use it,
+ * redistribute it and/or modify it under the terms of the:
+ *
+ * Eclipse Public License version 2.0, or
+ * GNU General Public License version 2, or
+ * GNU Lesser General Public License version 2.1.
+ */
+package org.truffleruby.collections;
+
+import org.truffleruby.core.array.ArrayUtils;
+
+/** Simplistic array stack implementation that will partial-evaluate nicely, unlike {@link java.util.ArrayDeque}. */
+@SuppressWarnings("unchecked")
+public class SimpleStack<T> {
+
+    Object[] storage;
+    int index = -1;
+
+    public SimpleStack() {
+        this(16);
+    }
+
+    public SimpleStack(int length) {
+        this.storage = new Object[length];
+    }
+
+    public boolean isEmpty() {
+        return index == -1;
+    }
+
+    public void push(T value) {
+        if (++index == storage.length) {
+            storage = ArrayUtils.grow(storage, storage.length * 2);
+        }
+        storage[index] = value;
+    }
+
+    public T peek() {
+        return (T) storage[index];
+    }
+
+    public T pop() {
+        T out = (T) storage[index];
+        storage[index] = null;
+        --index;
+        return out;
+    }
+
+    public int size() {
+        return index + 1;
+    }
+}

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

@@ -17,6 +17,8 @@
 
 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;
@@ -26,6 +28,7 @@
 import org.truffleruby.builtins.Primitive;
 import org.truffleruby.builtins.PrimitiveArrayArgumentsNode;
 import org.truffleruby.builtins.YieldingCoreMethodNode;
+import org.truffleruby.collections.SimpleStack;
 import org.truffleruby.core.CoreLibrary;
 import org.truffleruby.core.Hashing;
 import org.truffleruby.core.array.ArrayBuilderNode.BuilderState;
@@ -2338,4 +2341,98 @@ 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);
+            class Entry {
+                final RubyArray array;
+                final int index;
+
+                Entry(RubyArray array, int index) {
+                    this.array = array;
+                    this.index = index;
+                }
+            }
+            final SimpleStack<Entry> workStack = new SimpleStack<>();
+            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
-    recursively_flatten(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 recursively_flatten(self, out, level)
+    if Primitive.array_flatten_helper(self, out, level)
       Primitive.steal_array_storage(self, out)
       return self
     end
@@ -1251,44 +1251,6 @@ def zip(*others)
     end
   end
 
-  # Helper to recurse through flattening since the method
-  # is not allowed to recurse itself. Detects recursive structures.
-  def recursively_flatten(array, out, max_levels = -1)
-    modified = false
-
-    # Strict equality since < 0 means 'infinite'
-    if max_levels == 0
-      out.concat(array)
-      return false
-    end
-
-    max_levels -= 1
-    recursion = Truffle::ThreadOperations.detect_recursion(array) do
-      array = Truffle::Type.coerce_to(array, Array, :to_ary)
-
-      i = 0
-      size = array.size
-
-      while i < size
-        o = array.at i
-
-        tmp = Truffle::Type.rb_check_convert_type(o, Array, :to_ary)
-        if Primitive.nil? tmp
-          out << o
-        else
-          modified = true
-          recursively_flatten tmp, out, max_levels
-        end
-
-        i += 1
-      end
-    end
-
-    raise ArgumentError, 'tried to flatten recursive array' if recursion
-    modified
-  end
-  private :recursively_flatten
-
   private def sort_fallback(&block)
     # Use this instead of #dup as we want an instance of Array
     sorted = Array.new(self)