add enumerable#chunk

crystal-lang · Oct 25, 2016 · ac83b6c · ac83b6c
1 parent 06639e6
commit ac83b6c
Showing 3 changed files with 237 additions and 0 deletions.
diff --git a/spec/std/enumerable_spec.cr b/spec/std/enumerable_spec.cr
@@ -92,6 +92,118 @@ describe "Enumerable" do
     end
   end
 
+  describe "chunk" do
+    it "works" do
+      [1].chunk { true }.to_a.should eq [{true, [1]}]
+      [1, 2].chunk { false }.to_a.should eq [{false, [1, 2]}]
+      [1, 1, 2, 3, 3].chunk(&.itself).to_a.should eq [{1, [1, 1]}, {2, [2]}, {3, [3, 3]}]
+      [1, 1, 2, 3, 3].chunk(&.<=(2)).to_a.should eq [{true, [1, 1, 2]}, {false, [3, 3]}]
+      (0..10).chunk(&./(3)).to_a.should eq [{0, [0, 1, 2]}, {1, [3, 4, 5]}, {2, [6, 7, 8]}, {3, [9, 10]}]
+    end
+
+    it "work with class" do
+      [1, 1, 2, 3, 3].chunk(&.class).to_a.should eq [{Int32, [1, 1, 2, 3, 3]}]
+    end
+
+    it "works with block" do
+      res = [] of Tuple(Bool, Array(Int32))
+      [1, 2, 3].chunk { |x| x < 3 }.each { |(k, v)| res << {k, v} }
+      res.should eq [{true, [1, 2]}, {false, [3]}]
+    end
+
+    it "rewind" do
+      i = (0..10).chunk(&./(3))
+      i.next.should eq({0, [0, 1, 2]})
+      i.next.should eq({1, [3, 4, 5]})
+      i.rewind
+      i.next.should eq({0, [0, 1, 2]})
+      i.next.should eq({1, [3, 4, 5]})
+    end
+
+    it "returns elements of the Enumerable in an Array of Tuple, {v, ary}, where 'ary' contains the consecutive elements for which the block returned the value 'v'" do
+      result = [1, 2, 3, 2, 3, 2, 1].chunk { |x| x < 3 ? 1 : 0 }.to_a
+      result.should eq [{1, [1, 2]}, {0, [3]}, {1, [2]}, {0, [3]}, {1, [2, 1]}]
+    end
+
+    it "returns elements for which the block returns Enumerable::Chunk::Alone in separate Arrays" do
+      result = [1, 2, 3, 2, 1].chunk { |x| x < 2 ? Enumerable::Chunk::Alone : false }.to_a
+      result.should eq [{Enumerable::Chunk::Alone, [1]}, {false, [2, 3, 2]}, {Enumerable::Chunk::Alone, [1]}]
+    end
+
+    it "alone all" do
+      result = [1, 2].chunk { Enumerable::Chunk::Alone }.to_a
+      result.should eq [{Enumerable::Chunk::Alone, [1]}, {Enumerable::Chunk::Alone, [2]}]
+    end
+
+    it "does not return elements for which the block returns Enumerable::Chunk::Drop" do
+      result = [1, 2, 3, 3, 2, 1].chunk { |x| x == 2 ? Enumerable::Chunk::Drop : 1 }.to_a
+      result.should eq [{1, [1]}, {1, [3, 3]}, {1, [1]}]
+    end
+
+    it "drop all" do
+      result = [1, 2].chunk { Enumerable::Chunk::Drop }.to_a
+      result.size.should eq 0
+    end
+
+    it "nil allowed as value" do
+      result = [1, 2, 3, 2, 1].chunk { |x| x == 2 ? nil : 1 }.to_a
+      result.should eq [{1, [1]}, {nil, [2]}, {1, [3]}, {nil, [2]}, {1, [1]}]
+    end
+
+    it "nil 2 case" do
+      result = [nil, nil, 1, 1, nil].chunk(&.itself).to_a
+      result.should eq [{nil, [nil, nil]}, {1, [1, 1]}, {nil, [nil]}]
+    end
+  end
+
+  describe "chunks" do
+    it "works" do
+      [1].chunks { true }.should eq [{true, [1]}]
+      [1, 2].chunks { false }.should eq [{false, [1, 2]}]
+      [1, 1, 2, 3, 3].chunks(&.itself).should eq [{1, [1, 1]}, {2, [2]}, {3, [3, 3]}]
+      [1, 1, 2, 3, 3].chunks(&.<=(2)).should eq [{true, [1, 1, 2]}, {false, [3, 3]}]
+      (0..10).chunks(&./(3)).should eq [{0, [0, 1, 2]}, {1, [3, 4, 5]}, {2, [6, 7, 8]}, {3, [9, 10]}]
+    end
+
+    it "work with class" do
+      [1, 1, 2, 3, 3].chunks(&.class).should eq [{Int32, [1, 1, 2, 3, 3]}]
+    end
+
+    it "work with pure enumerable" do
+      SpecEnumerable.new.chunks(&./(2)).should eq [{0, [1]}, {1, [2, 3]}]
+    end
+
+    it "returns elements for which the block returns Enumerable::Chunk::Alone in separate Arrays" do
+      result = [1, 2, 3, 2, 1].chunks { |x| x < 2 ? Enumerable::Chunk::Alone : false }
+      result.should eq [{Enumerable::Chunk::Alone, [1]}, {false, [2, 3, 2]}, {Enumerable::Chunk::Alone, [1]}]
+    end
+
+    it "alone all" do
+      result = [1, 2].chunks { Enumerable::Chunk::Alone }
+      result.should eq [{Enumerable::Chunk::Alone, [1]}, {Enumerable::Chunk::Alone, [2]}]
+    end
+
+    it "does not return elements for which the block returns Enumerable::Chunk::Drop" do
+      result = [1, 2, 3, 3, 2, 1].chunks { |x| x == 2 ? Enumerable::Chunk::Drop : 1 }
+      result.should eq [{1, [1]}, {1, [3, 3]}, {1, [1]}]
+    end
+
+    it "drop all" do
+      result = [1, 2].chunks { Enumerable::Chunk::Drop }
+      result.size.should eq 0
+    end
+
+    it "nil allowed as value" do
+      result = [1, 2, 3, 2, 1].chunks { |x| x == 2 ? nil : 1 }
+      result.should eq [{1, [1]}, {nil, [2]}, {1, [3]}, {nil, [2]}, {1, [1]}]
+    end
+
+    it "nil 2 case" do
+      result = [nil, nil, 1, 1, nil].chunks(&.itself)
+      result.should eq [{nil, [nil, nil]}, {1, [1, 1]}, {nil, [nil]}]
+    end
+  end
+
   describe "each_cons" do
     it "returns running pairs" do
       array = [] of Array(Int32)

diff --git a/src/enumerable.cr b/src/enumerable.cr
@@ -74,6 +74,82 @@ module Enumerable(T)
     any? &.itself
   end
 
+  # Returns an iterator with chunks
+  #
+  #     (0..7).chunk(&./(3)).to_a => [{0, [0, 1, 2]}, {1, [3, 4, 5]}, {2, [6, 7]}]
+  #
+  def chunk(&block : T -> U)
+    self.each.chunk &block
+  end
+
+  # Returns an array of chunks (TODO: rewrite to use enumerator)
+  #
+  #     (0..7).chunks(&./(3)) => [{0, [0, 1, 2]}, {1, [3, 4, 5]}, {2, [6, 7]}]
+  #
+  def chunks(&block : T -> U)
+    res = [] of Tuple(U, Array(T))
+    chunks_internal(block) { |k, v| res << {k, v} }
+    res
+  end
+
+  # :nodoc:
+  module Chunk
+    record Drop
+    record Alone
+
+    # :nodoc:
+    class Accumulator(T, U)
+      def initialize
+        @key = uninitialized U
+        @initialized = false
+        @data = [] of T
+      end
+
+      def init(key, val)
+        return if key == Drop
+        @key = key
+        @data = [val]
+        @initialized = true
+      end
+
+      def add(d)
+        @data << d
+      end
+
+      def fetch
+        if @initialized
+          {@key, @data}.tap { @initialized = false }
+        end
+      end
+
+      def same_as?(key)
+        return false unless @initialized
+        return false if key == Alone || key == Drop
+        @key == key
+      end
+    end
+  end
+
+  # :nodoc:
+  private def chunks_internal(original_block : T -> U)
+    acc = Chunk::Accumulator(T, U).new
+    each do |val|
+      key = original_block.call(val)
+      if acc.same_as?(key)
+        acc.add(val)
+      else
+        if tuple = acc.fetch
+          yield(*tuple)
+        end
+        acc.init(key, val)
+      end
+    end
+
+    if tuple = acc.fetch
+      yield(*tuple)
+    end
+  end
+
   # Returns an array with the results of running the block against each element of the collection, removing `nil` values.
   #
   #     ["Alice", "Bob"].map { |name| name.match(/^A./) }         #=> [#<Regex::MatchData "Al">, nil]

diff --git a/src/iterator.cr b/src/iterator.cr
@@ -975,4 +975,53 @@ module Iterator(T)
       self
     end
   end
+
+  # Returns an iterator that returns a chunks.
+  #
+  #     iter = (1..3).chunk(&./(2))
+  #     iter.next # => {0, [1]}
+  #     iter.next # => {1, [2, 3]}
+  #     iter.next # => Iterator::Stop::INSTANCE
+  #
+  def chunk(&block : T -> U)
+    Chunk(typeof(self), T, U).new(self, &block)
+  end
+
+  # :nodoc:
+  class Chunk(I, T, U)
+    include Iterator(Tuple(U, Array(T)))
+    @iterator : I
+
+    def initialize(@iterator : Iterator(T), &@original_block : T -> U)
+      @acc = Enumerable::Chunk::Accumulator(T, U).new
+    end
+
+    def next
+      @iterator.each do |val|
+        key = @original_block.call(val)
+
+        if @acc.same_as?(key)
+          @acc.add(val)
+        else
+          tuple = @acc.fetch
+          @acc.init(key, val)
+          return tuple if tuple
+        end
+      end
+
+      if tuple = @acc.fetch
+        return tuple
+      end
+      stop
+    end
+
+    def rewind
+      @iterator.rewind
+      init_state
+    end
+
+    private def init_state
+      @acc = Enumerable::Chunk::Accumulator(T, U).new
+    end
+  end
 end