Allow Iterator.stop to be used inside Iterator.of block (#4208)

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

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