Search code examples
luaiterator

Lua, custom iterator - proper way to define?


I'm working with a lot of data files written in Lua. Most of them is written in this manner, a "phone book" as an example:

data = {
    -- First Level - country
    USA = {
        -- Second level - city
        Denver = {
            -- Third level - actual entries
            {name = 'John', number = '12345'},
            -- more entries
        },
        Washington = {
            {name = 'Ann', number = '54321'},
            -- more entries
        },
        -- more cities with entries
    },
    -- more countries with cities and entries
}

So the fact that first level is 'Country' and second is 'City' is implicit, but it makes the data more compact.

Now, when actually searching for some data, I'd like to iterate over this data as entries including this leveled, implicit information.

-- Coroutine yielding entries including level data
function corIter(data)
    for country,l1 in pairs(data) do
        for city,l2 in pairs(l1) do
            for _,entry in pairs(l2) do
                -- Copy the entry
                local out = {}
                for k,v in pairs(entry) do
                    out[k] = v
                end
                -- Add level properties
                out.country = country
                out.city = city
                coroutine.yield(out)
            end
        end
    end
end

-- Iterate over the entries
local cor = coroutine.create(corIter)
local _, entry = coroutine.resume(cor, data)
while entry do
    -- Handle the entry, has 'name', 'number', 'country' and 'city' keys
    table.print(entry) -- (custom print function I use)

    -- Get the next one
    _, entry = coroutine.resume(cor)
end  

But I think this approach may be bad, since it keeps a whole thread alive just to iterate over a damn table in a specific way.

Is there any other "obvious" solution to this? The key there is performance and ease of use. I don't exactly need a general solution (for arbitrary numbers of "levels" inside data table), but this all in all fells like a hack.


Solution

  • local function phones(d)
       local cn, c, tn, t, i
       return
          function()
             local a
             repeat
                if tn then
                   a, i = t[i], i+1
                   if not a then
                      i, tn, t = 1, next(c, tn)
                   end
                else
                   cn, c = next(d, cn)
                   i, tn, t = 1, next(c or {})
                end
             until a or not cn
             return cn, tn, a
          end
    end
    
    for country, town, abonent in phones(data) do
       print(country, town, abonent.name, abonent.number)
    end