Merge pull request #287 from Shopify/split-allocations

Optimize `Money#split(n)` for large `n`

lib/money.rb

@@ -5,6 +5,7 @@
 require_relative 'money/currency'
 require_relative 'money/null_currency'
 require_relative 'money/allocator'
+require_relative 'money/splitter'
 require_relative 'money/config'
 require_relative 'money/money'
 require_relative 'money/errors'

lib/money/money.rb

@@ -296,12 +296,12 @@ def allocate_max_amounts(maximums)
   #
   # @param [2] number of parties.
   #
-  # @return [Array<Money, Money, Money>]
+  # @return [Enumerable<Money, Money, Money>]
   #
   # @example
-  #   Money.new(100, "USD").split(3) #=> [Money.new(34), Money.new(33), Money.new(33)]
+  #   Money.new(100, "USD").split(3) #=> Enumerable[Money.new(34), Money.new(33), Money.new(33)]
   def split(num)
-    calculate_splits(num).sum([]) { |value, count| Array.new(count, value) }
+    Splitter.new(self, num)
   end
 
   # Calculate the splits evenly without losing pennies.
@@ -316,17 +316,7 @@ def split(num)
   # @example
   #   Money.new(100, "USD").calculate_splits(3) #=> {Money.new(34) => 1, Money.new(33) => 2}
   def calculate_splits(num)
-    raise ArgumentError, "need at least one party" if num < 1
-    subunits = self.subunits
-    low = Money.from_subunits(subunits / num, currency)
-    high = Money.from_subunits(low.subunits + 1, currency)
-
-    num_high = subunits % num
-
-    {}.tap do |result|
-      result[high] = num_high if num_high > 0
-      result[low] = num - num_high
-    end
+    Splitter.new(self, num).split.dup
   end
 
   # Clamps the value to be within the specified minimum and maximum. Returns

lib/money/splitter.rb

@@ -0,0 +1,115 @@
+# frozen_string_literal: true
+
+class Money
+  class Splitter
+    include Enumerable
+
+    def initialize(money, num)
+      @num = Integer(num)
+      raise ArgumentError, "need at least one party" if num < 1
+      @money = money
+      @split = nil
+    end
+
+    protected attr_writer :split
+
+    def split
+      @split ||= begin
+        subunits = @money.subunits
+        low = Money.from_subunits(subunits / @num, @money.currency)
+        high = Money.from_subunits(low.subunits + 1, @money.currency)
+
+        num_high = subunits % @num
+
+        split = {}
+        split[high] = num_high if num_high > 0
+        split[low] = @num - num_high
+        split.freeze
+      end
+    end
+
+    alias_method :to_ary, :to_a
+
+    def first(count = (count_undefined = true))
+      if count_undefined
+        each do |money|
+          return money
+        end
+      else
+        if count >= size
+          to_a
+        else
+          result = Array.new(count)
+          index = 0
+          each do |money|
+            result[index] = money
+            index += 1
+            break if index == count
+          end
+          result
+        end
+      end
+    end
+
+    def last(count = (count_undefined = true))
+      if count_undefined
+        reverse_each do |money|
+          return money
+        end
+      else
+        if count >= size
+          to_a
+        else
+          result = Array.new(count)
+          index = 0
+          reverse_each do |money|
+            result[index] = money
+            index += 1
+            break if index == count
+          end
+          result.reverse!
+          result
+        end
+      end
+    end
+
+    def [](index)
+      offset = 0
+      split.each do |money, count|
+        offset += count
+        if index < offset
+          return money
+        end
+      end
+      nil
+    end
+
+    def reverse_each(&block)
+      split.reverse_each do |money, count|
+        count.times do
+          yield money
+        end
+      end
+    end
+
+    def each(&block)
+      split.each do |money, count|
+        count.times do
+          yield money
+        end
+      end
+    end
+
+    def reverse
+      copy = dup
+      copy.split = split.reverse_each.to_h.freeze
+      copy
+    end
+
+    def size
+      count = 0
+      split.each_value { |c| count += c }
+      count
+    end
+  end
+end

spec/money_spec.rb

@@ -728,22 +728,32 @@
     specify "#split needs at least one party" do
       expect {Money.new(1).split(0)}.to raise_error(ArgumentError)
       expect {Money.new(1).split(-1)}.to raise_error(ArgumentError)
+      expect {Money.new(1).split(0.1)}.to raise_error(ArgumentError)
+      expect(Money.new(1).split(BigDecimal("0.1e1")).to_a).to eq([Money.new(1)])
+    end
+
+    specify "#split can be zipped" do
+      expect(Money.new(100).split(3).zip(Money.new(50).split(3)).to_a).to eq([
+        [Money.new(33.34), Money.new(16.67)],
+        [Money.new(33.33), Money.new(16.67)],
+        [Money.new(33.33), Money.new(16.66)],
+      ])
     end
 
     specify "#gives 1 cent to both people if we start with 2" do
-      expect(Money.new(0.02).split(2)).to eq([Money.new(0.01), Money.new(0.01)])
+      expect(Money.new(0.02).split(2).to_a).to eq([Money.new(0.01), Money.new(0.01)])
     end
 
     specify "#split may distribute no money to some parties if there isnt enough to go around" do
-      expect(Money.new(0.02).split(3)).to eq([Money.new(0.01), Money.new(0.01), Money.new(0)])
+      expect(Money.new(0.02).split(3).to_a).to eq([Money.new(0.01), Money.new(0.01), Money.new(0)])
     end
 
     specify "#split does not lose pennies" do
-      expect(Money.new(0.05).split(2)).to eq([Money.new(0.03), Money.new(0.02)])
+      expect(Money.new(0.05).split(2).to_a).to eq([Money.new(0.03), Money.new(0.02)])
     end
 
     specify "#split does not lose dollars with non-decimal currencies" do
-      expect(Money.new(5, 'JPY').split(2)).to eq([Money.new(3, 'JPY'), Money.new(2, 'JPY')])
+      expect(Money.new(5, 'JPY').split(2).to_a).to eq([Money.new(3, 'JPY'), Money.new(2, 'JPY')])
     end
 
     specify "#split a dollar" do
@@ -759,6 +769,41 @@
       expect(moneys[1].value).to eq(33)
       expect(moneys[2].value).to eq(33)
     end
+
+    specify "#split return respond to #first" do
+      expect(Money.new(100).split(3).first).to eq(Money.new(33.34))
+      expect(Money.new(100).split(3).first(2)).to eq([Money.new(33.34), Money.new(33.33)])
+
+      expect(Money.new(100).split(10).first).to eq(Money.new(10))
+      expect(Money.new(100).split(10).first(2)).to eq([Money.new(10), Money.new(10)])
+      expect(Money.new(20).split(2).first(4)).to eq([Money.new(10), Money.new(10)])
+    end
+
+    specify "#split return respond to #last" do
+      expect(Money.new(100).split(3).last).to eq(Money.new(33.33))
+      expect(Money.new(100).split(3).last(2)).to eq([Money.new(33.33), Money.new(33.33)])
+      expect(Money.new(20).split(2).last(4)).to eq([Money.new(10), Money.new(10)])
+    end
+
+    specify "#split return supports destructuring" do
+      first, second = Money.new(100).split(3)
+      expect(first).to eq(Money.new(33.34))
+      expect(second).to eq(Money.new(33.33))
+
+      first, *rest = Money.new(100).split(3)
+      expect(first).to eq(Money.new(33.34))
+      expect(rest).to eq([Money.new(33.33), Money.new(33.33)])
+    end
+
+    specify "#split return can be reversed" do
+      list = Money.new(100).split(3)
+      expect(list.first).to eq(Money.new(33.34))
+      expect(list.last).to eq(Money.new(33.33))
+
+      list = list.reverse
+      expect(list.first).to eq(Money.new(33.33))
+      expect(list.last).to eq(Money.new(33.34))
+    end
   end
 
   describe "calculate_splits" do

spec/splitter_spec.rb

@@ -0,0 +1,104 @@
+# frozen_string_literal: true
+require 'spec_helper'
+require 'yaml'
+
+RSpec.describe "Money::Splitter" do
+  specify "#split needs at least one party" do
+    expect {Money.new(1).split(0)}.to raise_error(ArgumentError)
+    expect {Money.new(1).split(-1)}.to raise_error(ArgumentError)
+    expect {Money.new(1).split(0.1)}.to raise_error(ArgumentError)
+    expect(Money.new(1).split(BigDecimal("0.1e1")).to_a).to eq([Money.new(1)])
+  end
+
+  specify "#split can be zipped" do
+    expect(Money.new(100).split(3).zip(Money.new(50).split(3)).to_a).to eq([
+      [Money.new(33.34), Money.new(16.67)],
+      [Money.new(33.33), Money.new(16.67)],
+      [Money.new(33.33), Money.new(16.66)],
+    ])
+  end
+
+  specify "#gives 1 cent to both people if we start with 2" do
+    expect(Money.new(0.02).split(2).to_a).to eq([Money.new(0.01), Money.new(0.01)])
+  end
+
+  specify "#split may distribute no money to some parties if there isnt enough to go around" do
+    expect(Money.new(0.02).split(3).to_a).to eq([Money.new(0.01), Money.new(0.01), Money.new(0)])
+  end
+
+  specify "#split does not lose pennies" do
+    expect(Money.new(0.05).split(2).to_a).to eq([Money.new(0.03), Money.new(0.02)])
+  end
+
+  specify "#split does not lose dollars with non-decimal currencies" do
+    expect(Money.new(5, 'JPY').split(2).to_a).to eq([Money.new(3, 'JPY'), Money.new(2, 'JPY')])
+  end
+
+  specify "#split a dollar" do
+    moneys = Money.new(1).split(3)
+    expect(moneys[0].subunits).to eq(34)
+    expect(moneys[1].subunits).to eq(33)
+    expect(moneys[2].subunits).to eq(33)
+  end
+
+  specify "#split a 100 yen" do
+    moneys = Money.new(100, 'JPY').split(3)
+    expect(moneys[0].value).to eq(34)
+    expect(moneys[1].value).to eq(33)
+    expect(moneys[2].value).to eq(33)
+  end
+
+  specify "#split return respond to #first" do
+    expect(Money.new(100).split(3).first).to eq(Money.new(33.34))
+    expect(Money.new(100).split(3).first(2)).to eq([Money.new(33.34), Money.new(33.33)])
+
+    expect(Money.new(100).split(10).first).to eq(Money.new(10))
+    expect(Money.new(100).split(10).first(2)).to eq([Money.new(10), Money.new(10)])
+    expect(Money.new(20).split(2).first(4)).to eq([Money.new(10), Money.new(10)])
+  end
+
+  specify "#split return respond to #last" do
+    expect(Money.new(100).split(3).last).to eq(Money.new(33.33))
+    expect(Money.new(100).split(3).last(2)).to eq([Money.new(33.33), Money.new(33.33)])
+    expect(Money.new(20).split(2).last(4)).to eq([Money.new(10), Money.new(10)])
+  end
+
+  specify "#split return supports destructuring" do
+    first, second = Money.new(100).split(3)
+    expect(first).to eq(Money.new(33.34))
+    expect(second).to eq(Money.new(33.33))
+
+    first, *rest = Money.new(100).split(3)
+    expect(first).to eq(Money.new(33.34))
+    expect(rest).to eq([Money.new(33.33), Money.new(33.33)])
+  end
+
+  specify "#split return can be reversed" do
+    list = Money.new(100).split(3)
+    expect(list.first).to eq(Money.new(33.34))
+    expect(list.last).to eq(Money.new(33.33))
+
+    list = list.reverse
+    expect(list.first).to eq(Money.new(33.33))
+    expect(list.last).to eq(Money.new(33.34))
+  end
+
+  describe "calculate_splits" do
+    specify "#calculate_splits gives 1 cent to both people if we start with 2" do
+      actual = Money.new(0.02, 'CAD').calculate_splits(2)
+
+      expect(actual).to eq({
+        Money.new(0.01, 'CAD') => 2,
+      })
+    end
+
+    specify "#calculate_splits gives an extra penny to one" do
+      actual = Money.new(0.04, 'CAD').calculate_splits(3)
+
+      expect(actual).to eq({
+        Money.new(0.02, 'CAD') => 1,
+        Money.new(0.01, 'CAD') => 2,
+      })
+    end
+  end
+end