Permutations in Lua. An iterator example

Here's some code I just wrote when messing with an algorithm to generate permutations. The code is just perfect to be used in the form of an iterator, and although I was reading the example in Perl (I read it in Higher Order Perl, an amazing book no matter what's your programming language of choice), I thought Lua would be a good candidate for that.

Along the way, I wrote a few utility functions you can see in there. Mostly tests on function composition and mapping over iterators.

Here's the code. If you need to generate permutations, I found this algorithm (which I don't know the name, but let's call it 'odometer counting') very easy to implement and understand. At least easier than Randal's way of doing it.

local inspect = require'inspect'

function map(f, t)
  local r = {}
  if type(t) == 'table' then
    for _, x in ipairs(t) do
      r[#r+1] = f(x)
    end
  else
    for x in t do
      r[#r+1] = f(x)
    end
  end
  return r
end

function count()
  local c = 0
  return function()
    c = c + 1
    return c
  end
end

-- for x in count() do
--   print(x)
--   if x > 10 then break end
-- end

function permutations(...)
  local function inc(t, pos)
    if t[pos][3] == t[pos][2] then
      if pos == 1 then return nil end
      t[pos][3] = 1
      return inc(t, pos-1)
    else
      t[pos][3] = t[pos][3] + 1
      return true
    end

  end

  local sets = {...}
  local state = map(function(x)
                      return {x, #x , 1}
                    end , sets)
  state[#state][3] = 0

  local curr = #state

  return function()
    while true do
      if inc(state, curr) then
        return map(function(s)
                     return s[1][s[3]] end,
                   state)
      else
        return nil
      end
    end
  end
end

function compose(f,g)
  return function (...)
    return f(g(unpack(arg)))
  end
end

pinspect = compose(print, inspect)

map(pinspect, permutations({1,2,3}, {5,6,7}))

local c = 0
for i in permutations({1,2,3,5,6} , {3,4,5,6,7,6}) do
  c = c+1
  if c == 10 then break end
  pinspect(i)
end