Enumerable#memoized so Enumerable#lazy doesn't get lonely

memoized.rb

# frozen_string_literal: true

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

    def initialize(enum)
      @original_inspect = enum.inspect
      @enum = cloned(enum)
      @cached_values = []

      super(@enum.size) do |yielder|
        @cached_values.each do |cached_value|
          yielder << cached_value
        end

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

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

    def inspect
      "#<#{self.class.name}: #{@original_inspect}>"
    end

    alias eager itself

    private

    def cloned(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
    end
  end

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

using Enumerator::Memoized::Refinement

# Pretend it takes a second to multiply n by 3.
counter = 1.step.lazy.map do |n|
  puts "Eagerly multiplying `#{n}' by `3' ..."
  sleep 1
  n * 3
end.memoized

p counter.first(2)
# Two seconds later...
#=> [3, 6]
p counter.first(2)
#=> [3, 6]
p counter.first(3)
# One second later...
#=> [3, 6, 9]