Merge pull request #5569 from dearblue/array-product

Added `Array#product` method

mrbgems/mruby-array-ext/mrblib/array.rb

@@ -896,4 +896,51 @@ class Array
   alias append push
   alias prepend unshift
   alias filter! select!
+
+  ##
+  # call-seq:
+  #   ary.product(*arys)                  ->   array
+  #   ary.product(*arys) { |item| ... }   ->   self
+  def product(*arys, &block)
+    size = arys.size
+    i = size
+    while i > 0
+      i -= 1
+      unless arys[i].kind_of?(Array)
+        raise TypeError, "no implicit conversion into Array"
+      end
+    end
+
+    i = size
+    total = self.size
+    total *= arys[i -= 1].size while i > 0
+
+    if block
+      result = self
+      list = ->(*, e) { block.call e }
+      class << list; alias []= call; end
+    else
+      result = [nil] * total
+      list = result
+    end
+
+    i = 0
+    while i < total
+      group = [nil] * (size + 1)
+      j = size
+      n = i
+      while j > 0
+        j -= 1
+        a = arys[j]
+        b = a.size
+        group[j + 1] = a[n % b]
+        n /= b
+      end
+      group[0] = self[n]
+      list[i] = group
+      i += 1
+    end
+
+    result
+  end
 end

mrbgems/mruby-array-ext/test/array.rb

@@ -421,3 +421,25 @@ assert('Array#transpose') do
   assert_raise(TypeError) { [1].transpose }
   assert_raise(IndexError) { [[1], [2,3,4]].transpose }
 end
+
+assert "Array#product" do
+  assert_equal [[1], [2], [3]], [1, 2, 3].product
+  assert_equal [], [1, 2, 3].product([])
+  assert_equal [], [1, 2, 3].product([4, 5, 6], [])
+
+  expect = [[1, 5, 8], [1, 5, 9], [1, 6, 8], [1, 6, 9], [1, 7, 8], [1, 7, 9],
+            [2, 5, 8], [2, 5, 9], [2, 6, 8], [2, 6, 9], [2, 7, 8], [2, 7, 9],
+            [3, 5, 8], [3, 5, 9], [3, 6, 8], [3, 6, 9], [3, 7, 8], [3, 7, 9],
+            [4, 5, 8], [4, 5, 9], [4, 6, 8], [4, 6, 9], [4, 7, 8], [4, 7, 9]]
+  assert_equal expect, [1, 2, 3, 4].product([5, 6, 7], [8, 9])
+
+  expect = [[1, 4, 7], [1, 4, 8], [1, 4, 9], [1, 5, 7], [1, 5, 8], [1, 5, 9], [1, 6, 7], [1, 6, 8], [1, 6, 9],
+            [2, 4, 7], [2, 4, 8], [2, 4, 9], [2, 5, 7], [2, 5, 8], [2, 5, 9], [2, 6, 7], [2, 6, 8], [2, 6, 9],
+            [3, 4, 7], [3, 4, 8], [3, 4, 9], [3, 5, 7], [3, 5, 8], [3, 5, 9], [3, 6, 7], [3, 6, 8], [3, 6, 9]]
+
+  assert_equal expect, [1, 2, 3].product([4, 5, 6], [7, 8, 9])
+  base = [1, 2, 3]
+  x = []
+  assert_equal base, base.product([4, 5, 6], [7, 8, 9]) { |e| x << e }
+  assert_equal expect, x
+end