diff --git a/spec/std/enumerable_spec.cr b/spec/std/enumerable_spec.cr index 9391c2298f04..7cb4a615d458 100644 --- a/spec/std/enumerable_spec.cr +++ b/spec/std/enumerable_spec.cr @@ -235,84 +235,96 @@ describe "Enumerable" do end end - describe "each_cons" do - it "returns running pairs" do - array = [] of Array(Int32) - [1, 2, 3, 4].each_cons(2) { |pair| array << pair } - array.should eq([[1, 2], [2, 3], [3, 4]]) - end + describe "#each_cons" do + context "iterator" do + it "iterates" do + iter = [1, 2, 3, 4, 5].each_cons(3) + iter.next.should eq([1, 2, 3]) + iter.next.should eq([2, 3, 4]) + iter.next.should eq([3, 4, 5]) + iter.next.should be_a(Iterator::Stop) + end - it "returns running triples" do - array = [] of Array(Int32) - [1, 2, 3, 4, 5].each_cons(3) { |triple| array << triple } - array.should eq([[1, 2, 3], [2, 3, 4], [3, 4, 5]]) - end + it "iterates with reuse = true" do + iter = [1, 2, 3, 4, 5].each_cons(3, reuse: true) - it "returns each_cons iterator" do - iter = [1, 2, 3, 4, 5].each_cons(3) - iter.next.should eq([1, 2, 3]) - iter.next.should eq([2, 3, 4]) - iter.next.should eq([3, 4, 5]) - iter.next.should be_a(Iterator::Stop) - end + a = iter.next + a.should eq([1, 2, 3]) - it "returns each_cons iterator with reuse = true" do - iter = [1, 2, 3, 4, 5].each_cons(3, reuse: true) + b = iter.next + b.should be(a) + end - a = iter.next - a.should eq([1, 2, 3]) + it "iterates with reuse = array" do + reuse = [] of Int32 + iter = [1, 2, 3, 4, 5].each_cons(3, reuse: reuse) - b = iter.next - b.should be(a) - end + a = iter.next + a.should eq([1, 2, 3]) + a.should be(reuse) + end - it "returns each_cons iterator with reuse = array" do - reuse = [] of Int32 - iter = [1, 2, 3, 4, 5].each_cons(3, reuse: reuse) + it "iterates with reuse = deque" do + reuse = Deque(Int32).new + iter = [1, 2, 3, 4, 5].each_cons(3, reuse: reuse) - a = iter.next - a.should eq([1, 2, 3]) - a.should be(reuse) + a = iter.next + a.should eq(Deque{1, 2, 3}) + a.should be(reuse) + end end - it "returns each_cons iterator with reuse = deque" do - reuse = Deque(Int32).new - iter = [1, 2, 3, 4, 5].each_cons(3, reuse: reuse) + context "yield" do + it "returns running pairs" do + array = [] of Array(Int32) + [1, 2, 3, 4].each_cons(2) { |pair| array << pair } + array.should eq([[1, 2], [2, 3], [3, 4]]) + end - a = iter.next - a.should eq(Deque{1, 2, 3}) - a.should be(reuse) - end + it "returns running triples" do + array = [] of Array(Int32) + [1, 2, 3, 4, 5].each_cons(3) { |triple| array << triple } + array.should eq([[1, 2, 3], [2, 3, 4], [3, 4, 5]]) + end - it "yields running pairs with reuse = true" do - array = [] of Array(Int32) - object_ids = Set(UInt64).new - [1, 2, 3, 4].each_cons(2, reuse: true) do |pair| - object_ids << pair.object_id - array << pair.dup + it "yields running pairs with reuse = true" do + array = [] of Array(Int32) + object_ids = Set(UInt64).new + [1, 2, 3, 4].each_cons(2, reuse: true) do |pair| + object_ids << pair.object_id + array << pair.dup + end + array.should eq([[1, 2], [2, 3], [3, 4]]) + object_ids.size.should eq(1) end - array.should eq([[1, 2], [2, 3], [3, 4]]) - object_ids.size.should eq(1) - end - it "yields running pairs with reuse = array" do - array = [] of Array(Int32) - reuse = [] of Int32 - [1, 2, 3, 4].each_cons(2, reuse: reuse) do |pair| - pair.should be(reuse) - array << pair.dup + it "yields running pairs with reuse = array" do + array = [] of Array(Int32) + reuse = [] of Int32 + [1, 2, 3, 4].each_cons(2, reuse: reuse) do |pair| + pair.should be(reuse) + array << pair.dup + end + array.should eq([[1, 2], [2, 3], [3, 4]]) end - array.should eq([[1, 2], [2, 3], [3, 4]]) - end - it "yields running pairs with reuse = deque" do - array = [] of Deque(Int32) - reuse = Deque(Int32).new - [1, 2, 3, 4].each_cons(2, reuse: reuse) do |pair| - pair.should be(reuse) - array << pair.dup + it "yields running pairs with reuse = deque" do + array = [] of Deque(Int32) + reuse = Deque(Int32).new + [1, 2, 3, 4].each_cons(2, reuse: reuse) do |pair| + pair.should be(reuse) + array << pair.dup + end + array.should eq([Deque{1, 2}, Deque{2, 3}, Deque{3, 4}]) end - array.should eq([Deque{1, 2}, Deque{2, 3}, Deque{3, 4}]) + end + end + + describe "#each_cons_pair" do + it "returns running pairs" do + array = [] of Array(Int32) + [1, 2, 3, 4].each_cons_pair { |a, b| array << [a, b] } + array.should eq([[1, 2], [2, 3], [3, 4]]) end end diff --git a/spec/std/iterator_spec.cr b/spec/std/iterator_spec.cr index 88a76367a2f3..5aa4a9e0e899 100644 --- a/spec/std/iterator_spec.cr +++ b/spec/std/iterator_spec.cr @@ -104,7 +104,7 @@ describe Iterator do end end - describe "cons" do + describe "#cons" do it "conses" do iter = (1..5).each.cons(3) iter.next.should eq([1, 2, 3]) @@ -168,6 +168,17 @@ describe Iterator do end end + describe "#cons_pair" do + it "conses" do + iter = (1..5).each.cons_pair + iter.next.should eq({1, 2}) + iter.next.should eq({2, 3}) + iter.next.should eq({3, 4}) + iter.next.should eq({4, 5}) + iter.next.should be_a(Iterator::Stop) + end + end + describe "cycle" do it "does cycle from range" do iter = (1..3).each.cycle diff --git a/src/enumerable.cr b/src/enumerable.cr index 6085d55f8c37..a6154414b67f 100644 --- a/src/enumerable.cr +++ b/src/enumerable.cr @@ -289,6 +289,10 @@ module Enumerable(T) # # This can be used to prevent many memory allocations when each slice of # interest is to be used in a read-only fashion. + # + # Chunks of two items can be iterated using `#each_cons_pair`, an optimized + # implementation for the special case of `size == 2` which avoids heap + # allocations. def each_cons(count : Int, reuse = false) raise ArgumentError.new "Invalid cons size: #{count}" if count <= 0 if reuse.nil? || reuse.is_a?(Bool) @@ -313,6 +317,40 @@ module Enumerable(T) nil end + # Iterates over the collection yielding pairs of adjacent items, + # but advancing one by one. + # + # ``` + # [1, 2, 3, 4, 5].each_cons do |a, b| + # puts "#{a}, #{b}" + # end + # ``` + # + # Prints: + # + # ```text + # 1, 2 + # 2, 3 + # 3, 4 + # 4, 5 + # ``` + # + # Chunks of more than two items can be iterated using `#each_cons`. + # This method is just an optimized implementation for the special case of + # `size == 2` to avoid heap allocations. + def each_cons_pair(& : (T, T) -> _) : Nil + last_elem = uninitialized T + first_iteration = true + each do |elem| + if first_iteration + first_iteration = false + else + yield last_elem, elem + end + last_elem = elem + end + end + # Iterates over the collection in slices of size *count*, # and runs the block for each of those. # diff --git a/src/iterator.cr b/src/iterator.cr index 4a2a478d349e..023be0b62a05 100644 --- a/src/iterator.cr +++ b/src/iterator.cr @@ -273,6 +273,10 @@ module Iterator(T) # # This can be used to prevent many memory allocations when each slice of # interest is to be used in a read-only fashion. + # + # Chunks oo two items can be iterated using `#cons_pair`, an optimized + # implementation for the special case of `size == 2` which avoids heap + # allocations. def cons(n : Int, reuse = false) raise ArgumentError.new "Invalid cons size: #{n}" if n <= 0 if reuse.nil? || reuse.is_a?(Bool) @@ -307,6 +311,49 @@ module Iterator(T) end end + # Returns an iterator that returns consecutive pairs of adjacent items. + # + # ``` + # iter = (1..5).each.cons + # iter.next # => {1, 2} + # iter.next # => {2, 3} + # iter.next # => {3, 4} + # iter.next # => {4, 5} + # iter.next # => Iterator::Stop::INSTANCE + # ``` + # + # Chunks of more than two items can be iterated using `#cons`. + # This method is just an optimized implementation for the special case of + # `size == 2` to avoid heap allocations. + def cons_pair : Iterator({T, T}) + ConsTuple(typeof(self), T).new(self) + end + + private struct ConsTuple(I, T) + include Iterator({T, T}) + include IteratorWrapper + + @last_elem : T | Iterator::Stop = Iterator::Stop::INSTANCE + + def initialize(@iterator : I) + end + + def next + elem = wrapped_next + return elem if elem.is_a?(Iterator::Stop) + + if @last_elem.is_a?(Iterator::Stop) + @last_elem = elem + + self.next + else + @last_elem, elem = elem, @last_elem + + {elem, @last_elem} + end + end + end + # Returns an iterator that repeatedly returns the elements of the original # iterator forever starting back at the beginning when the end was reached. #