Merge pull request #8332 from straight-shoota/feature/each_cons-2

Add Enumerable#each_cons_pair and Iterator#cons_pair yielding a tuple

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
 

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

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.
   #

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.
   #