Countdown problem solution in Lua

countdown.lua

--local strict = require"strict"

local function permutations(t)
  local fn = coroutine.yield
  local function permgen(t, n)
    for i = 1, n do
      t[n], t[i] = t[i], t[n]
      if n == 1 then
        fn(t)
      else
        permgen(t, n - 1)
      end
      t[n], t[i] = t[i], t[n]
    end
  end
  local co = coroutine.create(function () permgen(t, #t) end)
  return function()
    local code, res = coroutine.resume(co)
    return res
  end
end

local function genParam(name, from, to, skip, sep)
  local t = { }
  for i = from, to do
    if i ~= skip then
      t[#t + 1] = name.."["..i.."]"
    end
  end
  return "{ "..table.concat(t, sep or ", ").." }"
end

-- generates (loop unrolls) and compiles a function that is calculating operators for n parameters
local function genFn(n, env)
  local t = { }
  t[#t + 1] = "return function(a, exp, fn)"
  for i = 1, n do
    t[#t + 1] = "  calc("..genParam("a", 1, n, i)..", "..genParam("exp", 1, n, i)..", fn)"
  end
  for i = 1, n-1 do
    t[#t + 1] = "  calc("..genParam("a", 1, i)..", "..genParam("exp", 1, i)..", function(b1, bexp1)"
    t[#t + 1] = "    calc("..genParam("a", i + 1, n)..", "..genParam("exp", i + 1, n)..", function(b2, bexp2)"
    t[#t + 1] = "      calc({ b1, b2 }, { bexp1, bexp2 }, fn)"
    t[#t + 1] = "    end)"
    t[#t + 1] = "  end)"
  end
  t[#t + 1] = "end"
  local src = table.concat(t, "\n")
  if setfenv and loadstring then
    return setfenv(assert(loadstring(src)), env)()
  end
  return load(src, nil, "t", env)()
end

-- setup array and environment for generated functions and a dispatcher function that specializes on table size
local calcFn = { }
local calc, calcEnv

calc = function(a, exp, fn)
  if not calcFn[#a] then
    calcFn[#a] = genFn(#a, calcEnv)
  end
  calcFn[#a](a, exp, fn)
end

-- manually specialized template for array size of 1 (identity)
calcFn[1] = function(a, exp, fn)
  fn(a[1], exp[1])
end

-- manually specialized template for array size of 2. `exp` is generated in reverse polish notation
calcFn[2] = function(a, exp, fn)
  local a1, a2 = a[1], a[2]
  local exp1 = exp[1] or tostring(a1)
  local exp2 = exp[2] or tostring(a2)
  if a1 <= a2 then
    fn(a1 + a2, exp1.." "..exp2.." +")
    if a1 > 1 and a2 > 1 then
      fn(a1 * a2, exp1.." "..exp2.." *")
    end
  end
  if a1 > a2 then
    fn(a1 - a2, exp1.." "..exp2.." -")
  end
  if a2 > 1 and a1 % a2 == 0 then
    fn(a1 / a2, exp1.." "..exp2.." /")
  end
end

local prec = { ["+"] = 1, ["-"] = 1, ["*"] = 2, ["/"] = 2 }
local function toinfix(rpn)  
  local t = { }
  for tok in rpn:gmatch("%S+") do
    if tonumber(tok) then
      t[#t + 1] = { val = tok, prec = 10 }
    else
      local b = table.remove(t)
      local a = table.remove(t)
      local p = (prec[tok] or 0)
      a = a.prec < p and "("..a.val..")" or a.val
      b = b.prec < (p + ((tok == "-" or tok == "/") and 1 or 0)) and "("..b.val..")" or b.val
      t[#t + 1] = { val = a..tok..b, prec = p }
    end
  end
  return table.remove(t).val
end

local function countdown(t, target, fn)
  local ret, total, found = 0, 0, { }
  calcEnv = setmetatable({ calc = calc, cache = { } }, { __index = _G })
  for t in permutations(t) do
    calc(t, { }, function(result, exp)
      total = total + 1
      if result == target then
        local inf = toinfix(exp)
        if not found[inf] then
          found[inf] = exp
          fn(result, inf, exp)
          ret = ret + 1
        end
      end
    end)
  end
  return ret, total, found
end

if #arg > 0 then
  local t = { }
  for i = 1, #arg-1 do
    t[i] = tonumber(arg[i])
  end
  print(countdown(t, tonumber(arg[#arg]), print))
  if decoda_output then
    io.read()
  end
else
  print("usage: "..arg[0].." <n1> <n2> ... <target>")
end