Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
A mostly equivalent version of Array#flatten in Ruby, including checks for recursive arrays
require 'set'
class MyArray < Array
def my_flatten(level = -1)
level = Integer(level)
return self.dup if level == 0
flattened_array = self.class.new
recursively_flatten(self, flattened_array, level)
flattened_array
end
def my_flatten!(level = -1)
level = Integer(level)
return nil if level == 0
flattened_array = self.class.new
if recursively_flatten(self, flattened_array, level)
self.replace flattened_array
end
end
private
# Adapted and simplified from Array#recursively_flatten of Rubinius
# https://github.com/rubinius/rubinius/blob/b7a755c83f3dd3f0c1f5e546f0e58fb61851ea44/core/array.rb#L2089
def recursively_flatten(array, out, max_levels = -1, memo = Set.new)
# Strict equality since < 0 means 'infinite'
if max_levels == 0
out.concat(array)
return false
end
max_levels -= 1
modified = false
id = array.object_id
raise ArgumentError, 'tried to flatten recursive array' if memo.include?(id)
memo.add id
array.each do |element|
if element.respond_to?(:to_ary)
modified = true
recursively_flatten(element.to_ary, out, max_levels, memo)
else
out << element
end
end
memo.delete id
modified
end
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment