Return one particular element first, when iterating a lua table

Question

I wish to iterate through a Lua table, but want to get one particular element first always (its key is known beforehand).

Since the pairs function does not guarantee the order of a table's elements, I wrote a custom iterator function (I can insert every element into a indexed table and use the table.sort function, but I wanted to try writing a custom iterator).

local function 
pairsWithPriority(t,priorityKey)
 local i = 0 
 local function 
 closure(_,lastReturnedKey)
 local k,v 
 if i == 0 then --the first element should always be t[prioritykey]
 k,v = priorityKey,t[priorityKey]
 elseif i == 1 then --since we have returned the first element we reset the next pointer 
 k,v = next(t, nil) 
 else 
 k,v = next(t,lastReturnedKey)
 end 
 if i > 0 then 
 if k == priorityKey then --the first element is encountered AFTER it has been manually returned, so discard
 k,v = next(t, k)
 end
 end 
 i = i + 1 
 return k,v 
 end 
 return closure
 end 

When I use this iterator function:

local t = {a=1,b=2,c=6,d=4,e=4}
for i,v in pairsWithPriority(t,"c") do 
 print(i,v)
end

t[c] is always returned first.

This iterator works, but I'd like to know if there is a better way of doing this in a cleaner and more efficient way

Answer 1

You are performing the check i==0 and i==1 every time after i has incremented beyond those values. You can reduce this check by creating a lookup table.

lookup = {
 [0] = function()
 return priorityKey, t[priorityKey]
 end,
 [1] = function()
 return next(t, nil)
 end
}

The same can be done for you later checking i == priorityKey. I modified the lookup functions instead to deal with that case:

local function pairsWithPriority (t,priorityKey)
 local i, lookup = 0, {
 [0] = function()
 return priorityKey, t[priorityKey], true
 end,
 [1] = function()
 return next( t, nil )
 end,
 }
 local function closure( _, lastReturnedKey )
 local k, v, b
 if i > 1 then
 k, v = next( t, lastReturnedKey )
 else
 k, v, b = lookup[i]()
 if b then i = i + 1 return k, v end
 end
 if k == priorityKey then --the first element is encountered AFTER it has been manually returned, so discard
 k, v = next( t, k )
 end
 i = i + 1
 return k, v
 end 
 return closure
end