Lua - Custom Iterators

Iterators in Lua are functions which helps in traversing or iterating a collection. We can write our own custom iterator easily. Consider the following code to write a iterator which can be used in a for loop to generate a sequence of square of numbers.

Example - a Custom stateless Iterator

Following is an example of a custom stateless iterator to get a sequence of square of numbers.

main.lua

-- define an iterator to get square of an element
-- in case index >= state, lua will return Nil and terminates the iteration
function square(state,index)
   -- if index is less than maximum number
   if index < state then
      -- increment the index    
      index = index + 1
      return index, index*index 
   end	  
end

-- define an iterator
function squares(value)
   -- iterator, state, initial value
   return square, value, 0
end
   
-- call the iterator to get sequence of 5 numbers and their squares
for index, value in squares(5)
do
   print(index, value)
end

Output

When we run the above program, we will get the following output−

1	1
2	4
3	9
4	16
5	25

Example - a Custom Stateful Iterator

Following is an example of a custom stateful iterator which accepts an array and return squared value.

main.lua

-- Initialize an array
array = {1, 2, 3, 4, 5, 6}

-- return a stateful iterator to get values of collection passed
function elementIterator (collection)

   local index = 0
   local count = #collection
	
   -- The closure function is returned
	
   return function ()
      index = index + 1
		
      if index <= count
      then
	     -- return the current element of the iterator
	     local value = collection[index];
         return index, value * value
      end		
   end	
end

-- loop through the iterator
for index, value in elementIterator(array)  
do
   -- print the element
   print(index, value)
end

Output

When we run the above program, we will get the following output−

1	1
2	4
3	9
4	16
5	25
6	36