Lua Programming/How to Lua/iterator

An iterator is a special function which can be used within a generic for (or for each) statement. In most cases, an iterator is used to loop over or traverse a data-structure. Examples of this are the iterators returned by the pairs() and ipairs() functions which are used, respectively, to traverse the elements of a table or an and array.

fruit = { "Apple", "Banana", "Orange" }
colors = { ["Apple"] = "Red", ["Banana"] = "Yellow", ["Orange"] = "Orange" }
 
for idx,f in ipairs(fruit) do
  print(idx .. ": " .. f)
end
 
for f,col in pairs(colors) do
  print("The " .. f .. " is " .. col)
end

However, there is no requirement for an iterator to work on a data structure: an iterator can be designed for any case where looping is required. For example, the io.lines() function returns an iterator which returns a line from stdin at each iteration.

for line in io.lines() do
  print("Line: " .. line)
end

Building Iterators

An iterator consists of three things:

1. A transformation function.
2. A state value.
3. One or more loop variables.

The transformation function is used to modify the values of the loop variables (the variables that appear between the for and the in) for each iteration of the loop. This function is invoked before every iteration and takes, as arguments, the value the loop variables were set to during the last iteration. The function is expected to return a tuple containing the new values for these variables. The loop variables will be set to the components of the returned tuple and the loop will go through an iteration. Once that iteration has completed (so long as it has not been interrupted by a break or return statement) the transformation function will be called again and will return another set of values, which the loop variables will be set to for the next iteration, and so on and so forth. This cycle of calling the transformation function and iterating over the statements of the loop will continue until the transformation function returns nil.

Along with the loop variables, the transformation function is also passed in a state value which will remain constant throughout the loop. The can be used, for example, to maintain a reference to the data structure, file handle or resource the transformation function is iterating over.

An example of a transformation function which will produce a sequence of numbers up to 10 is given below. This transformation function requires only one loop variable which will have it's value stored in x.

function seq(state, x)
  if (x >= 10) then
    return nil                 -- Returning nil will terminate the loop
  else
    local newX = x + 1         -- The next value to use within the loop is the current value of x plus 1
                               -- This value will be used as the value of 'x' the next time this function is called.
    return newX
  end
end

The generic for loop expects a tuple containing the transformation function, the state value and the initial values of the loop variables. This tuple can be included directly after the in keyword.

-- Will display the numbers from 1 to 10
for x in seq, nil, 0 do
  print(x)
end

In most cases, however, this tuple will be returned by a function. This allows for the use of iterator factories which, when called, returns a new iterator that can be used with a generic for loop:

function countTo10()
  local function seq(state, x)
    if (x >= 10) then
      return nil
    else
      local newX = x + 1
      return newX
    end
  end
 
  return seq, nil, 0
end
 
for x in countTo10() do
  print(x)
end

Since Lua supports closures and functions as first-class objects, the iterator factory can also take arguments which can be used within the transformation function.

function countTo(limit)
  local function seq(state, x)
    if (x >= limit) then
      return nil
    else
      local newX = x + 1
      return newX
    end
  end
 
  return seq, nil, 0
end
 
for x in countTo(10) do       -- Print the numbers from 1 to 10
  print(x)
end
 
for y in countTo(20) do       -- Print the numbers from 1 to 20
  print(y)
end