diff --git a/spec/std/iterator_spec.cr b/spec/std/iterator_spec.cr index 80683ec6ef78..2fde318e8986 100644 --- a/spec/std/iterator_spec.cr +++ b/spec/std/iterator_spec.cr @@ -758,4 +758,72 @@ describe Iterator do iter.rewind.to_a.should eq([1, 2, 2, 3, 3]) end end + + describe "#slice_after" do + it "slices after" do + ary = [1, 3, 5, 8, 10, 11, 13, 15, 16, 17] + iter = ary.slice_after(&.even?) + iter.next.should eq([1, 3, 5, 8]) + iter.next.should eq([10]) + iter.next.should eq([11, 13, 15, 16]) + iter.next.should eq([17]) + iter.next.should be_a(Iterator::Stop) + end + + it "slices after: #to_a" do + ary = [1, 3, 5, 8, 10, 11, 13, 15, 16, 17] + ary.slice_after(&.even?).to_a.should eq([ + [1, 3, 5, 8], + [10], + [11, 13, 15, 16], + [17], + ]) + end + + it "slices after: #rewind" do + ary = [1, 3, 5, 8, 10, 11, 13, 15, 16, 17] + iter = ary.slice_after(&.even?) + iter.next.should eq([1, 3, 5, 8]) + iter.next.should eq([10]) + + iter.rewind + iter.next.should eq([1, 3, 5, 8]) + end + + it "slices after with reuse = true" do + ary = [1, 3, 5, 8, 10, 11, 13, 15, 16, 17] + iter = ary.slice_after(reuse: true, &.even?) + a = iter.next + a.should eq([1, 3, 5, 8]) + + b = iter.next + b.should eq([10]) + + a.should be(b) + end + + it "slices after with reuse = array" do + reuse = [] of Int32 + ary = [1, 3, 5, 8, 10, 11, 13, 15, 16, 17] + iter = ary.slice_after(reuse: reuse, &.even?) + a = iter.next + a.should eq([1, 3, 5, 8]) + + b = iter.next + b.should eq([10]) + + a.should be(b) + a.should be(reuse) + end + + it "slices after: non-bool block" do + ary = [1, nil, nil, 2, 3, nil] + iter = ary.slice_after(&.itself) + iter.next.should eq([1]) + iter.next.should eq([nil, nil, 2]) + iter.next.should eq([3]) + iter.next.should eq([nil]) + iter.next.should be_a(Iterator::Stop) + end + end end diff --git a/src/iterable.cr b/src/iterable.cr index 765eb0aeb9b7..f8a5ac5a01b5 100644 --- a/src/iterable.cr +++ b/src/iterable.cr @@ -45,4 +45,9 @@ module Iterable(T) def each_with_object(obj) each.with_object(obj) end + + # Same as `each.slice_after(reuse, &block)`. + def slice_after(reuse : Bool | Array(T) = false, &block : T -> B) forall B + each.slice_after(reuse, &block) + end end diff --git a/src/iterator.cr b/src/iterator.cr index 3ff165a57e93..fc3508d94cef 100644 --- a/src/iterator.cr +++ b/src/iterator.cr @@ -1290,4 +1290,91 @@ module Iterator(T) self end end + + # Returns an iterator for each chunked elements where the ends + # of chunks are defined by the block, when the block's value + # is _truthy_. Thus, the enumerable is sliced just **after** each + # time the block returns a _truthy_ value. + # + # For example, to get chunks that end at each uppercase letter: + # + # ``` + # ary = ['a', 'b', 'C', 'd', 'E', 'F', 'g', 'h'] + # iter = ary.slice_after(&.uppercase?) + # iter.next # => ['a', 'b', 'C'] + # iter.next # => ['d', 'E'] + # iter.next # => ['F'] + # iter.next # => ['g', 'h'] + # iter.next # => Iterator::Stop + # ``` + # + # By default, a new array is created and yielded for each slice when invoking `next`. + # * If *reuse* is `false`, the method will create a new array for each chunk + # * If *reuse* is `true`, the method will create a new array and reuse it. + # * If *reuse* is an `Array`, that array will be reused + # + # This can be used to prevent many memory allocations when each slice of + # interest is to be used in a read-only fashion. + def slice_after(reuse : Bool | Array(T) = false, &block : T -> B) forall B + SliceAfter(typeof(self), T, B).new(self, block, reuse) + end + + # :nodoc: + class SliceAfter(I, T, B) + include Iterator(Array(T)) + + def initialize(@iterator : I, @block : T -> B, reuse) + @end = false + @clear_on_next = false + + if reuse + if reuse.is_a?(Array) + @values = reuse + else + @values = [] of T + end + @reuse = true + else + @values = [] of T + @reuse = false + end + end + + def next + return stop if @end + + if @clear_on_next + @values.clear + @clear_on_next = false + end + + while true + value = @iterator.next + + if value.is_a?(Stop) + @end = true + if @values.empty? + return stop + else + return @reuse ? @values : @values.dup + end + end + + @values << value + + if @block.call(value) + @clear_on_next = true + return @reuse ? @values : @values.dup + end + end + end + + def rewind + @iterator.rewind + @values.clear + @end = false + @clear_on_next = false + self + end + end end