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:

-- 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