Enumerators and Enumerable Methods in Ruby

One of the coolest features of Ruby is the Enumerable methods. They give you great flexibility to iterate over any sort of collection and, e.g., check whether any element in that collection returns true for a given block (#any), or reject items in a collection for which the block returns false (#reject).

But what if you don’t know how many elements you’ll need to iterate over? Suppose, for instance, you’re trying to write your own method to produce the prime factorization for a number. Ruby has it’s own method to do this in the Prime class, but it’s a good example to illustrate the usefulness of enumerators. One way to produce a prime factorization of a number is to iterate over an array of primes and see whether each prime divides evenly into the number you want to factorize. If it does divide evenly, the prime is added to the results, the number to be factorized is reset to itself divided by the prime, and the iteration is restarted. The process continues until the prime being tested is greater than the number to be factorized.

class PrimeFactors
  def self.for(number)
    primes.each_with_object([]) do |prime, a|
      return a if prime > number
      number1, mod = number.divmod(prime)
      next unless mod == 0
      number = number1
      a << prime
      redo
    end
  end
...
end

But how many primes should the array contain? Who knows? It needs to be as long as it needs to be, and we can’t tell that beforehand. So, how do we decide upon a set of prime numbers to iterate through? Well, we know that the largest prime factor of a number will be less than the number, so we could generate an array of primes from 0 to the number. The Sieve of Eratosthenes might be useful for doing just that. But for a sufficiently large number, that process is very slow and completely unnecessary.

What we need is something that will generate prime numbers one at a time, indefinitely. A simple loop could do that. Find a prime, yield it to a block, find another, yield it, etc. But, hey, wouldn’t it be nice the generator could respond to the enumerable methods? Instead of yielding up new primes to any old block, the loop can be wrapped up in an enumerator object that responds to all the enumerable methods. In this case, we want it to respond to #each_with_object.

Here’s what it looks like:

def self.primes
  Enumerator.new do |y|
    y << 2
    i = 3
    loop do
      y << i if prime?(i)
      i += 2
    end
  end
end
def self.prime?(num)
  2.upto(Math.sqrt(num)).each { |i| return false if num % i == 0 }
  true
end

There are more efficient ways to generate primes out there, but the thing to see is that the primes are being generated one at a time. Enumerator#new takes a block, and that block gets a single parameter, a “yielder.” When an object is shoveled into the yielder, that object gets passed to the block of the method the enumerator is calling. Let’s look at our prime factorization method again:

  def self.for(number)
    primes.each_with_object([]) do |prime, a|
      return a if prime > number
      number1, mod = number.divmod(prime)
      next unless mod == 0
      number = number1
      a << prime
      redo
    end
  end

Remember, `primes` returns an enumerator. It’s not an array. It has an indefinite number of elements — it will keep churning out prime numbers as long as you’d like. But each time it spits out a prime, it’s passed to each_with_object’s block as if it were the next element in an array. Now, we don’t have to guess how many prime numbers we’ll need. It doesn’t matter! We’ll stop generating them when we don’t need any more, not before and not after. And we get all the fancy additions that enumerable methods give us. Look at that redo! It does exactly what you’d expect it to do if primes were an array — in this case, it starts the iteration over at 2. The enumerator acts sort of like an ever-expanding array, growing one item larger with each iteration.

The use case for enumerators might be limited, but it’s helpful to know they’re there. When you want to iterate over a collection with an enumerable method but you don’t know how many items that collection will need to contain, reach for an enumerator.