A memoized enumerator

enumerator-memoized.rb

# frozen_string_literal: true

class Enumerator
  class Memoized < Enumerator
    INTERRUPT = defined?(IRB::Abort) ? IRB::Abort : Interrupt
    private_constant :INTERRUPT

    module Refinement
      refine Enumerable do
        def memoized
          Enumerator::Memoized.new(self)
        end
      end
    end

    def initialize(enum)
      @enum = case enum
      when Enumerator::ArithmeticSequence
        Range.new(enum.begin, enum.end, enum.exclude_end?) % enum.step
      when Enumerator
        enum.clone.rewind
      else
        enum.clone.each
      end
      @memoized_values = []

      super(@enum.size) do |yielder|
        @memoized_values.each do |memoized_value|
          yielder << memoized_value
        end

        loop do
          @enum.next.tap do |next_value|
            @memoized_values << next_value
            yielder << next_value
          end
        end
      rescue INTERRUPT
        rewind
        raise
      end
    end

    def rewind
      @memoized_values.clear
      @enum.rewind
      super
    end

    alias eager itself
  end
end

using Enumerator::Memoized::Refinement

counter = 1.step.lazy.map { |n| puts "Eagerly producing: #{n} ..."; sleep 1; n }.memoized

p counter.first(2)
p counter.first(2)
p counter.first(3)