jruby · Feb 6, 2015 · Feb 6, 2015
diff --git a/spec/truffle/tags/core/array/flatten_tags.txt b/spec/truffle/tags/core/array/flatten_tags.txt
@@ -1,28 +1,5 @@
-fails:Array#flatten returns a one-dimensional flattening recursively
-fails:Array#flatten takes an optional argument that determines the level of recursion
-fails:Array#flatten returns dup when the level of recursion is 0
-fails:Array#flatten ignores negative levels
-fails:Array#flatten tries to convert passed Objects to Integers using #to_int
-fails:Array#flatten raises a TypeError when the passed Object can't be converted to an Integer
-fails:Array#flatten does not call flatten on elements
-fails:Array#flatten raises an ArgumentError on recursive arrays
-fails:Array#flatten flattens any element which responds to #to_ary, using the return value of said method
 fails:Array#flatten returns subclass instance for Array subclasses
 fails:Array#flatten returns a tainted array if self is tainted
 fails:Array#flatten returns an untrusted array if self is untrusted
-fails:Array#flatten with a non-Array object in the Array ignores the return value of #to_ary if it is nil
-fails:Array#flatten with a non-Array object in the Array raises a TypeError if the return value of #to_ary is not an Array
-fails:Array#flatten! modifies array to produce a one-dimensional flattening recursively
-fails:Array#flatten! returns self if made some modifications
-fails:Array#flatten! returns nil if no modifications took place
-fails:Array#flatten! should not check modification by size
-fails:Array#flatten! takes an optional argument that determines the level of recursion
-fails:Array#flatten! returns nil when the level of recursion is 0
-fails:Array#flatten! treats negative levels as no arguments
-fails:Array#flatten! tries to convert passed Objects to Integers using #to_int
-fails:Array#flatten! raises a TypeError when the passed Object can't be converted to an Integer
-fails:Array#flatten! does not call flatten! on elements
-fails:Array#flatten! raises an ArgumentError on recursive arrays
-fails:Array#flatten! flattens any elements which responds to #to_ary, using the return value of said method
 fails:Array#flatten! raises a RuntimeError on frozen arrays when the array is modified
 fails:Array#flatten! raises a RuntimeError on frozen arrays when the array would not be modified
diff --git a/truffle/src/main/java/org/jruby/truffle/nodes/core/ArrayNodes.java b/truffle/src/main/java/org/jruby/truffle/nodes/core/ArrayNodes.java
@@ -1385,37 +1385,6 @@ public Object eachWithIndexObject(VirtualFrame frame, RubyArray array, RubyProc
 
     }
 
-    @CoreMethod(names = "flatten")
-    public abstract static class FlattenNode extends ArrayCoreMethodNode {
-
-        public FlattenNode(RubyContext context, SourceSection sourceSection) {
-            super(context, sourceSection);
-        }
-
-        public FlattenNode(FlattenNode prev) {
-            super(prev);
-        }
-
-        @Specialization
-        public RubyArray flatten(RubyArray array) {
-            notDesignedForCompilation();
-            final List<Object> flattened = new ArrayList<>();
-            flatten(flattened, array.slowToArray());
-            return new RubyArray(getContext().getCoreLibrary().getArrayClass(), flattened.toArray(), array.getSize());
-        }
-
-        private void flatten(List<Object> flattened, Object[] store) {
-            for (Object value : store) {
-                if (value instanceof RubyArray) {
-                    flatten(flattened, ((RubyArray) value).slowToArray());
-                } else {
-                    flattened.add(value);
-                }
-            }
-        }
-
-    }
-
     @CoreMethod(names = "include?", required = 1)
     public abstract static class IncludeNode extends ArrayCoreMethodNode {
 
@@ -3276,24 +3245,6 @@ public int compare(Object a, Object b) {
 
     }
 
-    @CoreMethod(names = "to_a")
-    public abstract static class ToANode extends CoreMethodNode {
-
-        public ToANode(RubyContext context, SourceSection sourceSection) {
-            super(context, sourceSection);
-        }
-
-        public ToANode(ToANode prev) {
-            super(prev);
-        }
-
-        @Specialization
-        public RubyArray toA(RubyArray array) {
-            return array;
-        }
-
-    }
-
     @CoreMethod(names = "uniq")
     public abstract static class UniqNode extends CoreMethodNode {
 

diff --git a/truffle/src/main/ruby/jruby/truffle/core/rubinius/api/shims/array.rb b/truffle/src/main/ruby/jruby/truffle/core/rubinius/api/shims/array.rb
@@ -33,3 +33,12 @@ def start
     end
   end
 end
+
+class Array
+
+  def new_reserved(count)
+    # TODO CS 6-Feb-15 do we want to reserve space or allow the runtime to optimise for us?
+    []
+  end
+
+end
diff --git a/truffle/src/main/ruby/jruby/truffle/core/rubinius/kernel/common/array.rb b/truffle/src/main/ruby/jruby/truffle/core/rubinius/kernel/common/array.rb
@@ -309,4 +309,76 @@ def empty?
     @total == 0
   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 = Thread.detect_recursion(array) do
+      m = Rubinius::Mirror::Array.reflect array
+
+      i = m.start
+      total = i + m.total
+      tuple = m.tuple
+
+      while i < total
+        o = tuple.at i
+
+        if ary = Rubinius::Type.check_convert_type(o, Array, :to_ary)
+          modified = true
+          recursively_flatten(ary, out, max_levels)
+        else
+          out << o
+        end
+
+        i += 1
+      end
+    end
+
+    raise ArgumentError, "tried to flatten recursive array" if recursion
+    modified
+  end
+
+  private :recursively_flatten
+
+  def flatten(level=-1)
+    level = Rubinius::Type.coerce_to_collection_index level
+    return self.dup if level == 0
+
+    out = new_reserved size
+    recursively_flatten(self, out, level)
+    Rubinius::Type.infect(out, self)
+    out
+  end
+
+  def flatten!(level=-1)
+    Rubinius.check_frozen
+
+    level = Rubinius::Type.coerce_to_collection_index level
+    return nil if level == 0
+
+    out = new_reserved size
+    if recursively_flatten(self, out, level)
+      replace(out)
+      return self
+    end
+
+    nil
+  end
+
+  def to_a
+    if self.instance_of? Array
+      self
+    else
+      Array.new(self)
+    end
+  end
+
 end