Enumerators and Lazy Evaluation

# frozen_string_literal: true

require "shelf"

# << returns the yielder (so it chains); yield passes several values.
log = []
numbers = Enumerator.new do |y|
  log << :start
  y << 1 << 2
  y.yield 3, 4
  log << :end
  :finished
end
p numbers.to_a # => [1, 2, [3, 4]]
p log # => [:start, :end]
p numbers.size # => nil
p(numbers.each { nil }) # => :finished

# The block runs from the top for every iteration, and only as far as the consumer asks.
log.clear
p numbers.first(1) # => [1]
p log # => [:start]
p numbers.next # => 1
p log # => [:start, :start]

# State kept outside the block is shared by every iteration; state created inside it starts fresh.
counter = 0
leaky = Enumerator.new { |y| 3.times { y << (counter += 1) } }
p leaky.to_a # => [1, 2, 3]
p leaky.to_a # => [4, 5, 6]

# Infinite generators are fine because the consumer decides when to stop.
fibonacci = Enumerator.new do |y|
  a = 0
  b = 1
  loop do
    y << a
    a, b = b, a + b
  end
end
p fibonacci.take(10) # => [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]
p fibonacci.lazy.select(&:even?).first(5) # => [0, 2, 8, 34, 144]
p(fibonacci.each_with_index.find { |n, _| n > 100 }) # => [144, 12]
p fibonacci.size # => nil

# The optional argument is the size: a number, Float::INFINITY, or a callable evaluated on demand.
sized = Enumerator.new(3) { |y| 3.times { y << it } }
p sized.size # => 3
p(Enumerator.new(Float::INFINITY) { |y| loop { y << 1 } }.size) # => Infinity
calls = 0
dynamic = Enumerator.new(-> { calls += 1 }) { |y| y << 1 }
p [dynamic.size, dynamic.size] # => [1, 2]

# A Yielder can be passed where a block is expected: y.to_proc is what &y calls.
catalog = Shelf::Catalog.new
catalog.add(Shelf::Book.new(Shelf::Isbn.new("9780441013593"), "Dune", Shelf::Author.new("Frank Herbert")))
titles = Enumerator.new { |y| catalog.each { |book| y << book.title } }
direct = Enumerator.new { |y| catalog.map(&:title).each(&y) }
p titles.to_a # => ["Dune"]
p direct.to_a # => ["Dune"]

# Generators are values: pass them to methods that expect any Enumerable.
def numbered(source) = source.each_with_index.map { |item, index| "#{index + 1}. #{item}" }
p numbered(titles) # => ["1. Dune"]

# frozen_string_literal: true

def outcome
  yield.inspect
rescue StandardError => e
  "#{e.class}: #{e.message}"
end

# produce(initial) yields initial, then block(initial), block(block(initial)), and so on.
p(Enumerator.produce(1) { it * 2 }.take(6)) # => [1, 2, 4, 8, 16, 32]
p Enumerator.produce(1, &:succ).lazy.select(&:even?).first(3) # => [2, 4, 6]

# Without an initial value the first element is the first call of the block (called with no argument).
counter = 0
ticks = Enumerator.produce { counter += 1 }
p ticks.take(3) # => [1, 2, 3]

# Raising StopIteration inside the block ends the sequence: this is how finite sequences are made.
countdown = Enumerator.produce(3) { it.zero? ? raise(StopIteration) : it - 1 }
p countdown.to_a # => [3, 2, 1, 0]

# The Collatz sequence: each element depends on the previous one, which is what produce is for.
collatz = Enumerator.produce(6) do |n|
  raise StopIteration if n == 1

  n.even? ? n / 2 : (3 * n) + 1
end
p collatz.to_a # => [6, 3, 10, 5, 16, 8, 4, 2, 1]
p collatz.size # => Infinity

# Due dates of a loan that is renewed every 14 days, as Time values.
lent_on = Time.utc(2026, 10, 1)
due_dates = Enumerator.produce(lent_on + (14 * 86_400)) { it + (14 * 86_400) }
p(due_dates.first(3).map { it.strftime("%Y-%m-%d") }) # => ["2026-10-15", "2026-10-29", "2026-11-12"]
labels = due_dates.lazy.map { it.strftime("%b %-d") }
p(labels.take_while { it != "Dec 24" }.to_a) # => ["Oct 15", "Oct 29", "Nov 12", "Nov 26", "Dec 10"]

# Walking up a parent chain: stop by raising StopIteration when there is no parent.
Node = Struct.new(:name, :parent)
root = Node.new("library")
floor = Node.new("first floor", root)
aisle = Node.new("aisle 7", floor)
path = Enumerator.produce(aisle) { it.parent || raise(StopIteration) }
p path.map(&:name) # => ["aisle 7", "first floor", "library"]

# Ruby 4.0: the size keyword. The default is Float::INFINITY, nil means "unknown".
p Enumerator.produce(1, &:succ).size # => Infinity
p Enumerator.produce(1, size: nil, &:succ).size # => nil
p(Enumerator.produce(3, size: 4) { it.zero? ? raise(StopIteration) : it - 1 }.size) # => 4
p(Enumerator.produce(1, size: -> { 42 }, &:succ).size) # => 42
puts(outcome { Enumerator.produce }) # => ArgumentError: no block given

# frozen_string_literal: true

fiction = %w[Dune Emma]
reference = %w[Atlas]

both = fiction.each + reference
p both # => #<Enumerator::Chain: [#<Enumerator: ["Dune", "Emma"]:each>, ["Atlas"]]>
p both.class # => Enumerator::Chain
p both.to_a # => ["Dune", "Emma", "Atlas"]
p both.size # => 3
p both.with_index(1).to_a # => [["Dune", 1], ["Emma", 2], ["Atlas", 3]]
p fiction.chain(reference, 1..2) # => #<Enumerator::Chain: [["Dune", "Emma"], ["Atlas"], 1..2]>
p fiction.chain(reference, 1..2).to_a # => ["Dune", "Emma", "Atlas", 1, 2]
p fiction.chain.to_a # => ["Dune", "Emma"]
p Enumerator::Chain.new(fiction, reference).to_a # => ["Dune", "Emma", "Atlas"]

# The size is the sum when every part knows its size, nil when one does not, Infinity when one is infinite.
p fiction.chain(reference).size # => 3
p(fiction.chain(Enumerator.new { |y| y << 1 }).size) # => nil
p fiction.chain(1..).size # => Infinity

# A chain is lazy about its parts: they are iterated one after the other and only as far as needed.
log = []
second = Enumerator.new do |y|
  log << :second_started
  y << 3
  y << 4
end
chain = [1, 2].each + second
p chain.first(2) # => [1, 2]
p log # => []
p chain.first(3) # => [1, 2, 3]
p log # => [:second_started]

# Enumerable#chain on a lazy enumerator keeps the result lazy.
p((1..Float::INFINITY).lazy.chain([0]).class) # => Enumerator::Lazy
p([1, 2].lazy.chain([3]).map { it * 2 }.to_a) # => [2, 4, 6]

# A chain is not a cursor: next, peek and feed are undefined on Enumerator::Chain. Wrap it to get one.
p both.respond_to?(:next) # => false
p both.respond_to?(:rewind) # => true
cursor = both.to_enum
p [cursor.next, cursor.next] # => ["Dune", "Emma"]
p cursor.peek # => "Atlas"

# frozen_string_literal: true

def outcome
  yield.inspect
rescue StandardError => e
  "#{e.class}: #{e.message}"
end

formats = %i[paperback ebook]
languages = %w[en es]
editions = Enumerator.product(formats, languages)
p editions # => #<Enumerator::Product: [[:paperback, :ebook], ["en", "es"]]>
p editions.class # => Enumerator::Product
p editions.to_a # => [[:paperback, "en"], [:paperback, "es"], [:ebook, "en"], [:ebook, "es"]]
p editions.size # => 4
p(editions.map { |format, language| "#{format}/#{language}" }.first(2)) # => ["paperback/en", "paperback/es"]
p Enumerator.product(formats, languages, 1..2).first(2) # => [[:paperback, "en", 1], [:paperback, "en", 2]]

# Compare with Array#product: the array is built immediately, the enumerator is not.
p formats.product(languages) == editions.to_a # => true

# Edge cases: no arguments gives one empty combination; an empty argument gives none.
p Enumerator.product.to_a # => [[]]
p Enumerator.product(formats, []).to_a # => []
p Enumerator::Product.new(formats, languages).to_a == editions.to_a # => true

# The first enumerable is the outermost loop, and the last one varies fastest.
log = []
inner = Enumerator.new do |y|
  log << :inner_started
  y << :a
  y << :b
end
Enumerator.product([1, 2], inner).each { |number, letter| log << [number, letter] }
p log # => [:inner_started, [1, :a], [1, :b], :inner_started, [2, :a], [2, :b]]

# Infinite arguments: the size is Infinity and the combinations are produced on demand.
members = Enumerator.product(1.., %w[Ana Luis])
p members.first(3) # => [[1, "Ana"], [1, "Luis"], [2, "Ana"]]
p members.size # => Infinity
p Enumerator.product(1..3, 1..).first(4) # => [[1, 1], [1, 2], [1, 3], [1, 4]]

# The arguments only need each_entry, so a hash, a range or a Set all work.
p(Enumerator.product({ a: 1 }, Set[10, 20]).to_a) # => [[[:a, 1], 10], [[:a, 1], 20]]
puts(outcome { Enumerator.product(nil).to_a }) # => NoMethodError: undefined method 'each_entry' for nil

# On Ruby 3.3 to 4.0 with_index and each_with_index fail directly on a Product, so wrap it first.
puts(outcome { editions.with_index.to_a }) # => TypeError: wrong argument type product (expected enumerator)
p editions.to_enum.with_index(1).first(2) # => [[[:paperback, "en"], 1], [[:paperback, "es"], 2]]
p editions.lazy.each_with_index.first(2) # => [[[:paperback, "en"], 0], [[:paperback, "es"], 1]]

# frozen_string_literal: true

def outcome
  yield.inspect
rescue StandardError => e
  "#{e.class}: #{e.message}"
end

# Three ways to get an Enumerator::ArithmeticSequence (a number sequence with a fixed difference).
by_numeric = 1.step(10, 3)
by_range = (1..10).step(3)
by_percent = (1..10) % 3
by_keywords = 1.step(by: 3, to: 10)
p by_numeric # => (1.step(10, 3))
p by_range # => ((1..10).step(3))
p by_percent # => ((1..10).%(3))
p by_keywords # => (1.step(by: 3, to: 10))
p by_numeric.class # => Enumerator::ArithmeticSequence
p [by_numeric, by_range, by_percent, by_keywords].map(&:to_a).uniq # => [[1, 4, 7, 10]]
p by_numeric == by_range # => true
p by_numeric.eql?(by_range) # => true

# It answers questions about itself without iterating.
sequence = (1...20).step(4)
p [sequence.begin, sequence.end, sequence.step, sequence.exclude_end?] # => [1, 20, 4, true]
p [sequence.first, sequence.first(2), sequence.last, sequence.last(2), sequence.size] # => [1, [1, 5], 17, [13, 17], 5]
p sequence.to_a # => [1, 5, 9, 13, 17]
p [sequence.minmax, sequence.sum, sequence.include?(9)] # => [[1, 17], 45, true]

# Descending, endless and floating-point sequences.
p 10.step(1, -3).to_a # => [10, 7, 4, 1]
p 1.step(by: -1).first(3) # => [1, 0, -1]
p(1.step(by: 2).lazy.map { it * it }.first(3)) # => [1, 9, 25]
p 1.step.size # => Infinity
p 1.0.step(2.0, 0.25).to_a # => [1.0, 1.25, 1.5, 1.75, 2.0]
p 1.0.step(2.0, 0.1).size # => 11
p 1.0.step(2.0, 0.1).last # => 2.0
p 1r.step(2r, 1/3r).to_a # => [(1/1), (4/3), (5/3), (2/1)]

# Rules and errors.
p 3.step(1).to_a # => []
puts(outcome { 1.step(10, 0) }) # => ArgumentError: step can't be 0
puts(outcome { 1.step.last }) # => RangeError: cannot get the last element of endless arithmetic sequence
puts(outcome { (1..10).step("x").to_a }) # => TypeError: String can't be coerced into Integer
p(1.step(10, 3) { nil }) # => 1

# An ArithmeticSequence is a cursor too, unlike Chain and Product.
cursor = 1.step(10, 3)
p [cursor.next, cursor.next, cursor.peek] # => [1, 4, 7]

# Arrays accept a sequence as an index (Ruby 3.0) and take every n-th element.
shelves = %w[A B C D E F]
p shelves[(0..).step(2)] # => ["A", "C", "E"]
p shelves[(1..4).step(2)] # => ["B", "D"]
p shelves[(-1..0).step(-2)] # => ["F", "D", "B"]

# frozen_string_literal: true

# A paged remote catalogue simulated in memory: fetch_page counts how many "requests" were made.
class PagedSource
  PAGE_SIZE = 2

  attr_reader :requests

  def initialize(titles)
    @titles = titles
    @requests = 0
  end

  def fetch_page(number)
    @requests += 1
    @titles.each_slice(PAGE_SIZE).to_a.fetch(number - 1, [])
  end

  # Hide the paging behind one enumerator: callers see a flat, lazy stream of titles.
  def each_title(&)
    return to_enum(__method__) unless block_given?

    (1..).each do |number|
      page = fetch_page(number)
      break if page.empty?

      page.each(&)
    end
  end
end

source = PagedSource.new(["Dune", "The Hobbit", "Foundation", "Emma", "Ulysses"])
p source.each_title.first(3) # => ["Dune", "The Hobbit", "Foundation"]
p source.requests # => 2

source = PagedSource.new(["Dune", "The Hobbit", "Foundation", "Emma", "Ulysses"])
p(source.each_title.lazy.select { it.length > 4 }.map(&:upcase).first(2)) # => ["THE HOBBIT", "FOUNDATION"]
p source.requests # => 2

source = PagedSource.new(["Dune", "The Hobbit", "Foundation", "Emma", "Ulysses"])
p source.each_title.to_a.size # => 5
p source.requests # => 4

# An identifier generator: produce for the sequence, next for the cursor. Keep a cursor in one thread.
ids = Enumerator.produce(1000, &:succ)
book_ids = Array.new(3) { ids.next }
p book_ids # => [1000, 1001, 1002]
p ids.peek # => 1003

# frozen_string_literal: true

def outcome
  yield.inspect
rescue StandardError => e
  "#{e.class}: #{e.message}"
end

# Enumerator#each with a block calls the original method with that block, so the return value is the
# one of the underlying method. This is what makes map.with_index and select.with_index work.
titles = %w[Dune Emma Ulysses]
p titles.map.each(&:upcase) # => ["DUNE", "EMMA", "ULYSSES"]
p(titles.select.each { it.length > 4 }) # => ["Ulysses"]
p(titles.each_slice(2).each { nil }) # => ["Dune", "Emma", "Ulysses"]

# feed sets the value the next yield receives. For map that value becomes the mapped element.
mapper = titles.map
mapper.next
mapper.feed("D")
mapper.next
mapper.feed("E")
mapper.next
mapper.feed("U")
begin
  mapper.next
rescue StopIteration => e
  p e.result # => ["D", "E", "U"]
end
second = [1].each
second.feed(1)
puts(outcome { second.feed(2) }) # => TypeError: feed value already set

# Constructor errors.
puts(outcome { Enumerator.new }) # => ArgumentError: tried to create Proc object without a block
puts(outcome { Enumerator.new(1, 2) { nil } }) # => ArgumentError: wrong number of arguments (given 2, expected 0..1)
puts(outcome { [1].to_enum(:missing).to_a }) # => NoMethodError: undefined method 'missing' for an instance of Array

# A StopIteration raised inside the generator is not swallowed by to_a, only by loop.
broken = Enumerator.new do |y|
  y << 1
  raise StopIteration, "stopped early"
end
puts(outcome { broken.to_a }) # => StopIteration: stopped early
p broken.next # => 1
puts(outcome { broken.next }) # => StopIteration: stopped early

# with_index accepts any offset that converts to an Integer, and nil means 0.
p titles.each.with_index(10).to_a.last # => ["Ulysses", 12]
p titles.each.with_index(nil).to_a.last # => ["Ulysses", 2]
p titles.each_with_index.with_index(1).to_a.last # => [["Ulysses", 2], 3]

# rewind forwards to the receiver when it defines rewind.
class Tape
  def each
    yield :a
    yield :b
  end

  def rewind = puts("tape rewound")
end
reader = Tape.new.to_enum
p reader.next # => :a
reader.rewind # => tape rewound
p reader.next # => :a