Lua - writing iterator similar to ipairs, but selects indices

Question

I'd like to write an iterator that behaves exactly like ipairs, except which takes a second argument. The second argument would be a table of the indices that ipairs should loop over.

I'm wondering if my current approach is inefficient, and how I could improve it with closures.

I'm also open to other methods of accomplishing the same thing. But I like iterators because they're easy to use and debug.

I'll be making references to and using some of the terminology from Programming in Lua (PiL), especially the chapter on closures (chapter 7 in the link).

---

So I'd like to have this,

ary = {10,20,30,40}

for i,v in selpairs(ary, {1,3}) do
    ary[i] = v+5
    print(string.format("ary[%d] is now = %g", i, ary[i]))
end

which would output this:

ary[1] is now = 15
ary[3] is now = 35

My current approach is this : (in order: iterator, factory, then generic for)

iter = function (t, s)
    s = s + 1   
    local i = t.sel[s]
    local v = t.ary[i]

    if v then 
        return s, i, v
    end
end

function selpairs (ary, sel)
    local t = {}
    t.ary = ary
    t.sel = sel

    return iter, t, 0
end

ary = {10,20,30,40}
for _,i,v in selpairs(ary, {1,3}) do
    ary[i] = v+5
    print(string.format("ary[%d] is now = %g", i, ary[i]))
end

-- same output as before

It works. sel is the array of 'selected' indices. ary is the array you want to perform the loop on. Inside iter, s indexes sel, and i indexes ary.

But there are a few glaring problems.

• I must always discard the first returned argument s (_ in the for loop). I never need s, but it has to be returned as the first argument since it is the "control variable".
• The "invariant state" is actually two invariant states (ary and sel) packed into a single table. Pil says that this is more expensive, and recommends using closures. (Hence my writing this question).

---

The rest can of this can be ignored. I'm just providing more context for what I'm wanting to use selpairs for.

I'm mostly concerned with the second problem. I'm writing this for a library I'm making for generating music. Doing simple stuff like ary[i] = v+5 won't really be a problem. But when I do stuff like accessing object properties and checking bounds, then I get concerned that the 'invariant state as a table' approach may be creating unnecessary overhead. Should I be concerned about this?

If anything, I'd like to know how to write this with closures just for the knowledge.

Of course, I've tried using closures, but I'm failing to understand the scope of "locals in enclosing functions" and how it relates to a for loop calling an iterator.

---

As for the first problem, I imagine I could make the control variable a table of s, i, and v. And at the return in iter, unpack the table in the desired order. But I'm guessing that this is inefficient too.

---

Eventually, I'd like to write an iterator which does this, except nested into itself. My main data structure is arrays of arrays, so I'd hope to make something like this:

ary_of_arys = {
    {10,  20, 30,  40},
    {5,   6,  7,   8},
    {0.9, 1,  1.1, 1.2},
}

for aoa,i,v in selpairs_inarrays(ary_of_arys, {1,3}, {2,3,4}) do
    ary_of_arys[aoa][i] = v+5
end

And this too, could use the table approach, but it'd be nice to know how to take advantage of closures.

I've actually done something similar: A function that basically does the same thing by taking a function as it's fourth and final argument. It works just fine, but would this be less inefficient than an iterator?

Answer 1

And here's a co-routine based possibility:

function selpairs(t,selected)
  return coroutine.wrap(function()
                          for _,k in ipairs(selected) do
                            coroutine.yield(k,t[k])
                          end
                        end)
end

Answer 2

Not sure if I understand what you want to achive but why not simply write

local values = {"a", "b", "c", "d"}
for i,key in ipairs {3,4,1} do

  print(values[key])

end

and so forth, instead of implementing all that interator stuff? I mean your use case is rather simple. It can be easily extended to more dimensions.

Answer 3

You can hide "control variable" in an upvalue:

local function selpairs(ary, sel)
   local s = 0
   return
      function()
         s = s + 1
         local i = sel[s]
         local v = ary[i]
         if v then
            return i, v
         end
      end
end

Usage:

local ary = {10,20,30,40}

for i, v in selpairs(ary, {1,3}) do
   ary[i] = v+5
   print(string.format("ary[%d] is now = %g", i, ary[i]))
end

Nested usage:

local ary_of_arys = {
   {10,  20, 30,  40},
   {5,   6,  7,   8},
   {0.9, 1,  1.1, 1.2},
}
local outer_indices = {1,3}
local inner_indices = {2,3,4}

for aoa, ary in selpairs(ary_of_arys, outer_indices) do
   for i, v in selpairs(ary, inner_indices) do
      ary[i] = v+5  -- This is the same as  ary_of_arys[aoa][i] = v+5
   end
end