Return one particular element first, when iterating a lua table

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

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

StackExchange Code Review Q#61183, answer score: 2