[WIP] Implement ruby-style Enumerator

spec/std/enumerator_spec.cr

@@ -0,0 +1,228 @@
+require "spec"
+require "enumerator"
+
+describe Enumerator do
+  describe "initialize" do
+    it "returns an instance of Enumerator" do
+      instance = Enumerator(String).new { }
+      instance.should be_a(Enumerator(String))
+    end
+  end
+
+  describe "yielder" do
+    it "accepts values of enumerator type" do
+      enumerator = Enumerator(String).new do |yielder|
+        yielder << "hello"
+        yielder << "world"
+      end
+
+      enumerator.next.should eq("hello")
+      enumerator.next.should eq("world")
+    end
+
+    it "allows to chain yielder calls" do
+      enumerator = Enumerator(String).new do |yielder|
+        yielder << "hello" << "world"
+      end
+
+      enumerator.next.should eq("hello")
+      enumerator.next.should eq("world")
+    end
+  end
+
+  describe "next" do
+    it "returns the values in the order they were yielded" do
+      enumerator = Enumerator(String?).new do |yielder|
+        yielder << "hello"
+        yielder << nil
+        yielder << "world"
+      end
+
+      enumerator.next.should eq("hello")
+      enumerator.next.should be_nil
+      enumerator.next.should eq("world")
+      enumerator.next.should be_a(Iterator::Stop)
+    end
+
+    it "works when the type is Nil" do
+      enumerator = Enumerator(Nil).new do |yielder|
+        yielder << nil
+        yielder << nil
+      end
+
+      enumerator.next.should be_nil
+      enumerator.next.should be_nil
+      enumerator.next.should be_a(Iterator::Stop)
+    end
+
+    it "starts the yielder block at the first call" do
+      started = false
+      enumerator = Enumerator(Int32).new do |y|
+        started = true
+        y << 1
+      end
+
+      started.should be_false
+      enumerator.next.should eq(1)
+      started.should be_true
+    end
+
+    it "cooperates with other fibers outside the enumerator" do
+      enumerator = Enumerator(Int32).new do |y|
+        y << 1
+      end
+
+      value = 0
+      spawn do
+        value = enumerator.next
+      end
+
+      Fiber.yield
+
+      value.should eq(1)
+    end
+
+    it "cooperates with other fibers inside the enumerator" do
+      enumerator = Enumerator(Int32).new do |y|
+        spawn do
+          y << 1
+        end
+        Fiber.yield
+        y << 2
+      end
+
+      enumerator.next.should eq(1)
+      enumerator.next.should eq(2)
+    end
+  end
+
+  describe "peek" do
+    it "peeks at the next value without affecting next" do
+      enumerator = Enumerator(String).new do |yielder|
+        yielder << "hello"
+        yielder << "world"
+      end
+
+      enumerator.peek.should eq("hello")
+      enumerator.peek.should eq("hello")
+      enumerator.next.should eq("hello")
+      enumerator.peek.should eq("world")
+      enumerator.peek.should eq("world")
+      enumerator.next.should eq("world")
+      enumerator.peek.should be_a(Iterator::Stop)
+      enumerator.next.should be_a(Iterator::Stop)
+    end
+
+    it "does not affect rewind" do
+      enumerator = Enumerator(String).new do |yielder|
+        yielder << "hello"
+        yielder << "world"
+      end
+
+      enumerator.next.should eq("hello")
+      enumerator.peek.should eq("world")
+      enumerator.rewind
+      enumerator.peek.should eq("hello")
+      enumerator.next.should eq("hello")
+    end
+
+    it "works when the type is Nil" do
+      enumerator = Enumerator(Nil).new do |yielder|
+        yielder << nil
+        yielder << nil
+      end
+
+      enumerator.peek.should be_nil
+      enumerator.next.should be_nil
+      enumerator.peek.should be_nil
+      enumerator.next.should be_nil
+      enumerator.peek.should be_a(Iterator::Stop)
+      enumerator.next.should be_a(Iterator::Stop)
+    end
+  end
+
+  describe "each" do
+    it "iterates the yielded values" do
+      enumerator = Enumerator(String).new do |yielder|
+        yielder << "hello"
+        yielder << "world"
+      end
+
+      values = [] of String
+      enumerator.each { |value| values << value }
+      values.should eq(["hello", "world"])
+    end
+  end
+
+  describe "rewind" do
+    it "rewinds the iterator after full iteration" do
+      enumerator = Enumerator(String).new do |yielder|
+        yielder << "hello"
+        yielder << "world"
+      end
+
+      enumerator.next.should eq("hello")
+      enumerator.next.should eq("world")
+      enumerator.next.should be_a(Iterator::Stop)
+
+      enumerator.rewind.should be_a(Enumerator(String))
+
+      enumerator.next.should eq("hello")
+    end
+
+    it "rewinds the iterator after partial iteration" do
+      enumerator = Enumerator(String).new do |yielder|
+        yielder << "hello"
+        yielder << "world"
+      end
+      enumerator.next.should eq("hello")
+
+      enumerator.rewind.should be_a(Enumerator(String))
+      enumerator.next.should eq("hello")
+    end
+
+    it "can be called before the first iteration" do
+      enumerator = Enumerator(String).new do |yielder|
+        yielder << "hello"
+        yielder << "world"
+      end
+
+      enumerator.rewind.should be_a(Enumerator(String))
+
+      enumerator.next.should eq("hello")
+      enumerator.next.should eq("world")
+      enumerator.next.should be_a(Iterator::Stop)
+    end
+
+    it "can be rewound multiple times" do
+      enumerator = Enumerator(String).new do |yielder|
+        yielder << "hello"
+        yielder << "world"
+      end
+
+      enumerator.next.should eq("hello")
+      enumerator.next.should eq("world")
+      enumerator.next.should be_a(Iterator::Stop)
+
+      enumerator.rewind.should be_a(Enumerator(String))
+      enumerator.rewind.should be_a(Enumerator(String))
+
+      enumerator.next.should eq("hello")
+    end
+  end
+
+  describe "Fibonacci enumerator" do
+    it "generates the first 10 numbers in the Fibonacci sequence" do
+      # Make sure example from the Enumerator docs works.
+      fib = Enumerator(Int32).new do |y|
+        a = b = 1
+        loop do
+          y << a
+          a, b = b, a + b
+        end
+      end
+
+      fib.first(10).to_a.should eq([1, 1, 2, 3, 5, 8, 13, 21, 34, 55])
+    end
+  end
+end

src/enumerator.cr

@@ -0,0 +1,146 @@
+require "./weak_ref"
+
+# An `Enumerator` allows to lazily yield values from a block and consume them as `Iterator`.
+#
+# It is useful, if you have a block that generates values and needs to block in the middle
+# of execution until the next value is consumed.
+#
+# You should use an `Iterator` instead, if you do not need this blocking behavior, since it
+# is much more efficient.
+#
+# As an example, let's build a generator that yields numbers from the Fibonacci Sequence:
+#
+# ```
+# fib = Enumerator(Int32).new do |y|
+#   a = b = 1
+#   loop do
+#     y << a
+#     a, b = b, a + b
+#   end
+# end
+#
+# fib.first(10).to_a # => [1, 1, 2, 3, 5, 8, 13, 21, 34, 55]
+# ```
+#
+# Note that the above example works, but it would be much faster using an `Iterator`:
+#
+# ```
+# a = b = 1
+# fib = Iterator.of do
+#   a.tap { a, b = b, a + b }
+# end
+#
+# fib.first(10).to_a # => [1, 1, 2, 3, 5, 8, 13, 21, 34, 55]
+# ```
+# It is generally advisable to avoid `Enumerator` in cases where you have a large number
+# of iterations and the computation inside the block is trivial, because the additional
+# overhead of the blocking behavior that is implemented with `Fiber` under the hood will
+# negatively impact performance.
+#
+# However for more complex examples you will find that it might not be possible to replace
+# the loop with suspendable state and you would have to build your own blocking behavior
+# using for example an `Iterator` and a `Channel`.
+#
+# Using `Enumerator` is the right fit for those cases and it handles a lot of the tricky
+# edge cases like ensuring the block is not started until you consume the first value or
+# aborted if you rewind or abandon the iterator.
+class Enumerator(T)
+  include Iterator(T)
+
+  # :nodoc:
+  getter! current_fiber
+
+  @fiber : Fiber?
+  @done : Bool?
+
+  def initialize(&@block : Yielder(T) ->)
+    @current_fiber = Fiber.current
+    @stack = Deque(T).new(1)
+    @yielder = Yielder(T).new(self)
+    run
+  end
+
+  # Returns the next element yielded from the block or `Iterator::Stop::INSTANCE`
+  # if there are no more elements.
+  def next
+    fetch_next
+    stack.shift { stop }
+  end
+
+  # Rewinds the iterator and restarts the yielder block.
+  def rewind
+    fiber.kill
+    run
+    self
+  end
+
+  # Peeks at the next value without forwarding the iterator.
+  def peek
+    fetch_next
+    if stack.empty?
+      stop
+    else
+      stack[0]
+    end
+  end
+
+  # :nodoc:
+  def finalize
+    fiber.kill
+  end
+
+  protected def <<(value : T)
+    stack.push(value)
+  end
+
+  private getter stack
+  private getter! fiber
+  private getter? done
+
+  private def fetch_next : Nil
+    @current_fiber = Fiber.current
+    fiber.resume if fiber.alive? && stack.empty?
+  end
+
+  private def run
+    stack.clear
+
+    block = @block
+    yielder = @yielder.not_nil!
+
+    @fiber = Fiber.new do
+      block.call(yielder)
+    end.on_finish do
+      yielder.resume_current_fiber
+    end
+  end
+
+  private struct Yielder(T)
+    @enumerator : WeakRef(Enumerator(T))
+
+    def initialize(enumerator : Enumerator(T))
+      @enumerator = WeakRef.new(enumerator)
+    end
+
+    # Yields the next value to the enumerator.
+    def <<(value : T)
+      pass_to_enumerator(value)
+      resume_current_fiber
+      self
+    end
+
+    # :nodoc:
+    def resume_current_fiber
+      current_fiber = get_current_fiber
+      current_fiber.try &.resume
+    end
+
+    private def pass_to_enumerator(value : T)
+      @enumerator.value.try &.<< value
+    end
+
+    private def get_current_fiber
+      @enumerator.value.try &.current_fiber
+    end
+  end
+end