From 47a24055802be4beca2a0f20623188cfc30ea083 Mon Sep 17 00:00:00 2001
From: "Brian J. Cardiff" <bcardiff@manas.tech>
Date: Wed, 29 Mar 2017 02:38:12 -0300
Subject: [PATCH] Allow Iterator.stop to be used inside Iterator.of block

---
 spec/std/iterator_spec.cr | 13 +++++++++++++
 src/iterator.cr           | 10 ++++++++--
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/spec/std/iterator_spec.cr b/spec/std/iterator_spec.cr
index 9240862c20a0..0885982156de 100644
--- a/spec/std/iterator_spec.cr
+++ b/spec/std/iterator_spec.cr
@@ -13,6 +13,19 @@ describe Iterator do
       iter = Iterator.of { a += 1 }
       iter.first(3).to_a.should eq([1, 2, 3])
     end
+
+    it "creates singleton from block can call Iterator.stop" do
+      a = 0
+      iter = Iterator.of do
+        if a >= 5
+          Iterator.stop
+        else
+          a += 1
+        end
+      end
+      iter.should be_a(Iterator(Int32))
+      iter.first(10).to_a.should eq([1, 2, 3, 4, 5])
+    end
   end
 
   describe "compact_map" do
diff --git a/src/iterator.cr b/src/iterator.cr
index db9e849b6f9a..4e2b77d29b90 100644
--- a/src/iterator.cr
+++ b/src/iterator.cr
@@ -131,13 +131,19 @@ module Iterator(T)
   end
 
   def self.of(&block : -> T)
-    SingletonProc(T).new(block)
+    SingletonProc(typeof(without_stop(&block))).new(block)
+  end
+
+  private def self.without_stop(&block : -> T)
+    e = block.call
+    raise "" if e.is_a?(Iterator::Stop)
+    e
   end
 
   private struct SingletonProc(T)
     include Iterator(T)
 
-    def initialize(@proc : -> T)
+    def initialize(@proc : (-> (T | Iterator::Stop)) | (-> T))
     end
 
     def next