Iterators Tutorial

wiki

Algorithm

We can write our own iterators to be called by the for statement. The following Lua pseudo-code describes how the for statement works with iterators:

do
  local iterator, state, var1 = explist
  local var2, ... , varn
  while true do
    var1, ..., varn = iterator(state, var1)
    if var1 == nil then break end
      -- for block code
  end
end

Simple example

The following iterator will return a sequence of squared values. When we return no value (i.e. n>=state), Lua returns nil which terminates iteration. Notice the iterator is returning the next value in the sequence for a given iteration. We use the state to hold the number of iterations we wish to perform.

> function square(state,n) if n<state then n=n+1 return n,n*n end end

Here is the for statement calling the iterator:

> for i,n in square,5,0 do print(i,n) end
1       1
2       4
3       9
4       16
5       25

We could wrap the above example up (like pairs()) and provide a squares(nbvals) iterator constructor function. E.g.,

> function squares(nbvals) return square,nbvals,0 end  -- iterator,state,initial value

Now we can call it like pairs():

> for i,n in squares(5) do print(i,n) end
1       1
2       4
3       9
4       16
5       25

Complicated example

The following iterator is similar to ipairs, but allows for multiple tables to be iterated over.

function ipairs(...)
  local t = {...}
  local tmp = {...}
  -- if nothing to iterate over just return a dummy iterator
  if #tmp==0 then
    return function() end, nil, nil
  end
  local function mult_ipairs_it(t, i)
    i = i+1
    for j=1,#t do
      local val = t[j][i]
      if val == nil then return val end
      tmp[j] = val
    end
    return i, unpack(tmp)
  end
  return mult_ipairs_it, t, 0
end

local t1 = {'a', 'b', 'c', 'd', 'e'}
local t2 = {'A', 'B', 'C', 'D', 'E', 'F'}

for k,v1,v2 in ipairs(t1, t2) do
  print(k,v1,v2)
end

Number Theory Example

This iterator produces the prime divisors of its argument. It relies on the obvious fact that the smallest positive divisor, greater than one, of a nonzero integer must be a prime number.

primdiv = function (n)
   assert(n ~= 0)
   if n < 0 then n = -n end
   local f = function (s,v) -- s not used
                   local p
                   if n == 1 then return end
                   while n%v > 0 and v*v < n do
                    if v == 2 then v = 3
                    else v = v + 2 end
                   end -- while
                   if n%v == 0 then
                    n = n/v; return v
                   end
                   if v*v > n then
                    p = n
                    n = 1; return p
                   end
                 end -- function
     return f,nil,2
    end -- function

 for p in primdiv(84) do io.write(p," ") end --> 2 2 3 7

RiciLake says: I couldn't resist rewriting that:

 function primdiv(n)
   assert(n ~= 0)
   if n < 0 then n = -n end
   local function f(_, v)
     if n > 1 then
       while n%v > 0 do
         v = v + (v == 2 and 1 or 2)
         if v*v > n then v = n end
       end -- while
       n = n / v
       return v
     end -- if
   end -- function
   return f,nil,2
 end -- function primdiv

 for p in primdiv(84) do io.write(p," ") end --> 2 2 3 7

GavinWraith says: That is much nicer. I do regard, however, the use of function primdiv(n) as tantamount to methadone for first-order language junkies.