Programming in Lua : 9.3

This first edition was written for Lua 5.0. While still largely relevant for later versions, there are some differences.
The fourth edition targets Lua 5.3 and is available at Amazon and other bookstores.
By buying the book, you also help to support the Lua project.

Programming in Lua

Part I. The Language Chapter 9. Coroutines

9.3 – Coroutines as Iterators

We can see loop iterators as a quite specific example of the producer-consumer pattern. An iterator produces items to be consumed by the loop body. Therefore, it seems appropriate to use coroutines to write iterators. Actually, coroutines provide a powerful tool for this task. Again, the key feature is their ability to turn upside-down the relationship between caller and callee. With this feature, we can write iterators without worrying about how to keep state between successive calls to the iterator.

To illustrate this kind of use, let us write an iterator to traverse all permutations of a given array. It is not an easy task to write directly such iterator, but it is not so difficult to write a recursive function that generates all those permutations. The idea is simple: Put each array element in the last position, in turn, and recursively generate all permutations of the remaining elements. The code is as follows:

    function permgen (a, n)
      if n == 0 then
        printResult(a)
      else
        for i=1,n do
    
          -- put i-th element as the last one
          a[n], a[i] = a[i], a[n]
    
          -- generate all permutations of the other elements
          permgen(a, n - 1)
    
          -- restore i-th element
          a[n], a[i] = a[i], a[n]
    
        end
      end
    end

To see it working, we should define an appropriate printResult function and call permgen with proper arguments:

    function printResult (a)
      for i,v in ipairs(a) do
        io.write(v, " ")
      end
      io.write("\n")
    end
    
    permgen ({1,2,3,4}, 4)

After we have the generator ready, it is an automatic task to convert it to an iterator. First, we change printResult to yield:

    function permgen (a, n)
      if n == 0 then
        coroutine.yield(a)
      else
      ...

Then, we define a factory that arranges for the generator to run inside a coroutine, and then create the iterator function. The iterator simply resumes the coroutine to produce the next permutation:

    function perm (a)
      local n = table.getn(a)
      local co = coroutine.create(function () permgen(a, n) end)
      return function ()   -- iterator
        local code, res = coroutine.resume(co)
        return res
      end
    end

With that machinery in place, it is trivial to iterate over all permutations of an array with a for statement:

    for p in perm{"a", "b", "c"} do
      printResult(p)
    end
      --> b c a
      --> c b a
      --> c a b
      --> a c b
      --> b a c
      --> a b c

The perm function uses a common pattern in Lua, which packs a call to resume with its corresponding coroutine inside a function. This pattern is so common that Lua provides a special function for it: coroutine.wrap. Like create, wrap creates a new coroutine. Unlike create, wrap does not return the coroutine itself; instead, it returns a function that, when called, resumes the coroutine. Unlike the original resume, that function does not return an error code as its first result; instead, it raises the error in case of errors. Using wrap, we can write perm as follows:

    function perm (a)
      local n = table.getn(a)
      return coroutine.wrap(function () permgen(a, n) end)
    end

Usually, coroutine.wrap is simpler to use than coroutine.create. It gives us exactly what we need from a coroutine: a function to resume it. However, it is also less flexible. There is no way to check the status of a coroutine created with wrap. Moreover, we cannot check for errors.