NQueens in lua

NQueens.lua

local dump = require 'pl.pretty'.dump

local function printboard(board)
  if not board then return end
  local n = #board
  io.stdout:write ("\n", (" -"):rep(n), "\n")
  for i = 1, n do
    io.stdout:write "|"
    for j = 1, n do
      io.stdout:write(board[j][i] and "Q" or " ", "|")
    end
    io.stdout:write ("\n", (" -"):rep(n), "\n")
  end
end

local checkcount = 0
local function attacked(board, c, r)
  checkcount = checkcount + 1
  local n = #board
  for i = r - 1, 1, -1 do
    if board[c][i] then return true end
    local slope = r - i
    if board[c - slope] and board[c - slope][i] then return true end
    if board[c + slope] and board[c + slope][i] then return true end
  end
  return false
end

local function NQueens(board, Q)
  if Q > #board then return false end
  local function Solve_NQueens(board, Q, r)
    local n = #board

    if Q == 0 or r > n then return board end
    for c = 1, n do
      board[c][r] = true
      if not attacked(board, c, r)
         and Solve_NQueens(board, Q - 1, r + 1) then
         return board end
      board[c][r] = false
    end
    return false
  end

  return Solve_NQueens(board, Q, 1)
end

local function all_NQueens(board, Q)
  local function Solve_NQueens(board, Q, r)
    local n = #board

    if Q == 0 or r > n then 
      coroutine.yield(board)
      return nil
    end
    for c = 1, n do
      board[c][r] = true
      if not attacked(board, c, r) then
        Solve_NQueens(board, Q - 1, r + 1)
      end
      board[c][r] = false
    end
    return nil
  end

  return coroutine.wrap(function()
                          return Q <= #board 
                             and Solve_NQueens(board, Q, 1)
                              or nil
                        end)
end

local board =
{
  {}, {}, {}, {},
  -- {}, {}, {}, {},
}

-- board[3][1] = true
local unittest =
{
  testattack_clear = function()
    assert(not attacked(board, 4, 4))
  end,

  testattack_vertical = function()
    assert(attacked(board, 3, 4))
    assert(attacked(board, 3, 2))
  end,

  testattack_diag = function()
    assert(attacked(board, 2, 2))
    assert(attacked(board, 4, 2))
    assert(attacked(board, 1, 3))
  end,
}

for _, test in pairs(unittest) do
  -- test()
end

-- board[3][1] = nil
-- NQueens(board, 8)
-- printboard(board)
-- print("attacked called:", checkcount .. " times")

for s in all_NQueens(board, 4) do
  printboard(s)
end