Add product method to enumerable. (#2488)

crystal-lang · May 3, 2016 · b6c212c · b6c212c
1 parent a90921c
commit b6c212c
Showing 3 changed files with 73 additions and 2 deletions.
diff --git a/spec/std/enumerable_spec.cr b/spec/std/enumerable_spec.cr
@@ -573,6 +573,21 @@ describe "Enumerable" do
     end
   end
 
+  describe "product" do
+    assert { ([] of Int32).product.should eq(1) }
+    assert { [1, 2, 3].product.should eq(6) }
+    assert { [1, 2, 3].product(4).should eq(24) }
+    assert { [1, 2, 3].product(4.5).should eq(27) }
+    assert { (1..3).product { |x| x * 2 }.should eq(48) }
+    assert { (1..3).product(1.5) { |x| x * 2 }.should eq(72) }
+
+    it "uses zero from type" do
+      typeof([1, 2, 3].product).should eq(Int32)
+      typeof([1.5, 2.5, 3.5].product).should eq(Float64)
+      typeof([1, 2, 3].product(&.to_f)).should eq(Float64)
+    end
+  end
+
   describe "first" do
     assert { (1..3).first(1).should eq([1]) }
     assert { (1..3).first(4).should eq([1, 2, 3]) }

diff --git a/src/array.cr b/src/array.cr
@@ -1521,8 +1521,8 @@ class Array(T)
     result
   end
 
-  def product(ary, &block)
-    self.each { |a| ary.each { |b| yield a, b } }
+  def product(enumerable : Enumerable(U), &block)
+    self.each { |a| enumerable.each { |b| yield a, b } }
   end
 
   # Append. Pushes one value to the end of `self`, given that the type of the value is *T*

diff --git a/src/enumerable.cr b/src/enumerable.cr
@@ -848,6 +848,9 @@ module Enumerable(T)
   end
 
   # Adds *initial* and all the elements in the collection together.
+  # The type of *initial* will be the type of the sum, so use this if
+  # (for instance) you need to specify a large enough type to avoid
+  # overflow.
   #
   # Only collections of numbers (objects that can be added via an `+` method) are supported.
   #
@@ -884,6 +887,59 @@ module Enumerable(T)
     reduce(initial) { |memo, e| memo + (yield e) }
   end
 
+  # Multiplies all the elements in the collection together.
+  #
+  # Only collections of numbers (objects that can be multiplied via a `*` method) are supported.
+  #
+  #     [1, 2, 3, 4, 5, 6].product  #=> 720
+  #
+  # If the collection is empty, returns 1.
+  #
+  #     ([] of Int32).product #=> 1
+  def product
+    product Reflect(T).first.zero + 1
+  end
+
+  # Multiplies *initial* and all the elements in the collection
+  # together.  The type of *initial* will be the type of the product,
+  # so use this if (for instance) you need to specify a large enough
+  # type to avoid overflow.
+  #
+  # Only collections of numbers (objects that can be multiplied via a `*` method) are supported.
+  #
+  #     [1, 2, 3, 4, 5, 6].product(7)  #=> 5040
+  #
+  # If the collection is empty, returns *initial*.
+  #
+  #     ([] of Int32).product(7) #=> 7
+  def product(initial : Number)
+    product initial, &.itself
+  end
+
+  # Multiplies all results of the passed block for each element in the collection.
+  #
+  #     ["Alice", "Bob"].product { |name| name.size }  #=> 15 (5 * 3)
+  #
+  # If the collection is empty, returns 1.
+  #
+  #     ([] of Int32).product { |x| x + 1 } #=> 1
+  def product(&block)
+    product(Reflect(typeof(yield first)).first.zero + 1) do |value|
+      yield value
+    end
+  end
+
+  # Multiplies *initial* and all results of the passed block for each element in the collection.
+  #
+  #     ["Alice", "Bob"].product(2) { |name| name.size }  #=> 30 (2 * 5 * 3)
+  #
+  # If the collection is empty, returns one.
+  #
+  #     ([] of String).product(1) { |name| name.size } #=> 1
+  def product(initial : Number, &block)
+    reduce(initial) { |memo, e| memo * (yield e) }
+  end
+
   # Returns an array with the first *count* elements in the collection.
   #
   # If *count* is bigger than the number of elements in the collection, returns as many as possible. This