enum.rb

# frozen_string_literal: true

require 'fiber'

class Enum
  class Yielder
    def initialize(&block)
      raise LocalJumpError, 'no block given' unless block_given?

      @block = block
    end

    def yield(*args) = @block.call(*args)

    def <<(*args)
      @block.call(*args)

      self
    end

    def to_proc = ->(*args) { @block.call(*args) }
  end

  InitialNotProvided = Object.new

  def self.produce(initial = InitialNotProvided)
    raise ArgumentError, 'No block given' unless block_given?

    new Float::INFINITY do |yielder|
      yielded = initial == InitialNotProvided ? yield : initial

      loop do
        yielder << yielded

        yielded = yield yielded
      end
    end
  end

  include Enumerable

  attr_reader :size

  NoPeek = Object.new
  NoFeed = Object.new

  def initialize(size = nil, &block)
    @size = size
    @block = block
    @fiber = fiber
    @peek = NoPeek
    @feed = NoFeed
  end

  def each
    return self unless block_given?

    yield_interim_fiber do
      loop { yield self.next }
    end
  end

  def next
    unless @peek == NoPeek
      peek = @peek
      @peek = NoPeek

      return peek
    end

    raise StopIteration, 'iteration reached an end' unless @fiber.alive?

    @fiber.resume
  end

  def next_values = [self.next]

  def peek
    @peek = self.next if @peek == NoPeek

    @peek
  end

  def peek_values = [peek]

  def rewind
    @fiber = fiber
    @peek = NoPeek

    self
  end

  # TODO: Return an Enum here?
  def +(other) = Enumerator::Chain.new(self, other)

  def inspect = "#<#{self.class.name}: ...>"

  alias with_object each_with_object

  def with_index(offset = 0)
    # TODO: Return an Enum here?
    return to_enum(__method__, offset) { @size } unless block_given?

    each do
      yield _1, Integer(offset)

      offset += 1
    end
  end

  # TODO: Implement #feed, which is a strange one.

  private

  def yield_interim_fiber
    existing_fiber = @fiber
    rewind
    yield

    @fiber = existing_fiber
  end

  def fiber
    Fiber.new do
      @block.call Yielder.new { |*args| Fiber.yield(*args) }

      raise StopIteration, 'iteration reached an end'
    end
  end
end