lua iterators

Iterators Tutorial

lua-users home wiki

for statement iterators were covered in the ForTutorial. These are some notes on writing your own custom iterators.

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:

-- Equivalent to "for var1, ···, varn in explist do block end"
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

The state and current key value are passed into the iterator. The iterator returns the new value of the key and any other values, e.g. the value generated by the iterator. If nil is returned then the for loop is terminated.

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.

Misc examples

ipairs in Lua

function ipairs(t)
  local function ipairs_it(t, i)
    i = i+1
    local v = t[i]
    if v ~= nil then
      return i,v
    else
      return nil
    end
  end
  return ipairs_it, t, 0
end

Reversed ipairs in Lua

function ripairs(t)
  local max = 1
  while t[max] ~= nil do
    max = max + 1
  end
  local function ripairs_it(t, i)
    i = i-1
    local v = t[i]
    if v ~= nil then
      return i,v
    else
      return nil
    end
  end
  return ripairs_it, t, max
end

-- from the end backwards
function ripairs(t)
  local function ripairs_it(t,i)
    i=i-1
    local v=t[i]
    if v==nil then return v end
    return i,v
  end
  return ripairs_it, t, #t+1
end

-- traversing the whole 'array'
function ripairs(t)
  idx={}
  for k,v in pairs(t) do
    if type(k)=="number" then idx[#idx+1]=k end
  end
  table.sort(idx)
  local function ripairs_it(t,_)
    if #idx==0 then return nil end
    k=idx[#idx]
    idx[#idx]=nil
    return k,t[k]
  end
  return ripairs_it, t, nil
end

t1 = {'a', 'b', nil, 'd', 'e', nil}
for k,v in ripairs(t1) do print(k,v) end
--> 5 e
--> 4 d
--> 2 b
--> 1 a

RecentChanges · preferences edit · history Last edited December 9, 2009 11:03 am GMT (diff)