Simple Constraint Solver in Lua

--
-- Simple Finite Domain, Constraint Propagation Solver
--
-- Inspired by 'Fast Procedural Level Population with Playability Constraints'
-- by Ian Horswill and Leif Foged.
--
-- paper: http://www.aaai.org/ocs/index.php/AIIDE/AIIDE12/paper/view/5466/5691
-- code:  http://code.google.com/p/constraint-thingy/
--
-- I haven't implemented any of the more advanced parts of the paper like
-- interval domains of the DAG based pathing stuff.
--
-- It's not very efficient as it uses tables (Lua's built in hash map) for
-- representing domains instead of bit arrays but it is simple.
--


-- variable = {
--     name = <string>,
--     values = { [<value>] = true }*
-- }

-- constraint = function ( solver, variable )

-- solver = {
--     variables = { [variable] = true }+
--     varmap = { [<string>] = variable }+
--     constraints = { [variable] = { constraint }+ }+
--     stack = { { variables } }*
-- }

function printf( ... ) print(string.format(...)) end

function set( value )
    return { [value] = true }
end

function set_empty( set )
    return next(set) == nil
end

function set_singular( set )
    local k = next(set)

    return k ~= nil and next(set, k) == nil
end

function set_count( set )
    local result = 0

    for k, _ in pairs(set) do
        result = result + 1
    end

    return result
end

function set_copy( set )
    local result = {}

    for k, _ in pairs(set) do
        result[k] = true
    end

    return result
end

function set_intersection( set1, set2 )
    local result = {}

    for k, _ in pairs(set1) do
        if set2[k] then
            result[k] = true
        end
    end

    return result
end

function set_subtract( set1, set2 )
    local result = {}

    for k, _ in pairs(set1) do
        if not set2[k] then
            result[k] = true
        end
    end

    return result
end

function set_equal( set1, set2 )
    for k, _ in pairs(set1) do
        if not set2[k] then
            return false
        end
    end

    for k,_ in pairs(set2) do
        if not set1[k] then
            return false
        end
    end

    return true
end

function set_toarray( set )
    local result = {}

    for k, _ in pairs(set) do
        result[#result+1] = k
    end

    return result
end

function set_fromarray( array )
    local result = {}

    for _, v in ipairs(array) do
        result[v] = true
    end

    return result
end

function set_tostring( set )
    local parts = {}

    for k, _ in pairs(set) do
        parts[#parts+1] = k
    end

    return string.format("{%s}", table.concat(parts, ' '))
end

function print_stack( solver )
    for index, frame in ipairs(solver.stack) do
        printf('%d ->', index)

        for variable, values in pairs(frame) do
            printf('  %s = %s', variable.name, set_tostring(values))
        end
    end
end

function newframe( solver )
    local stack = solver.stack
    local frame = {}
    stack[#stack+1] = frame
end

function save( solver, var )
    local stack = solver.stack
    local frame = stack[#stack]

    frame[var] = var.values
end

function unwind( solver )
    local stack = solver.stack
    local frame = stack[#stack]

    for variable, values in pairs(frame) do
        variable.values = values
    end

    stack[#stack] = nil
end

function solution( solver )
    local result = {}

    for variable, _ in pairs(solver.variables) do
        assert(set_singular(variable.values))

        result[variable.name] = next(variable.values)
    end

    return result
end

function narrow( solver, var, set )
    -- printf('  narrow %s to %s', vartos(var), set_tostring(set))
    local newvalues = set_intersection(var.values, set)

    if set_empty(newvalues) then
        return false
    end

    if not set_equal(var.values, newvalues) then
        save(solver, var)

        var.values = newvalues

        local constraints = solver.constraints[var]

        -- printf('  propagate %s over %d constraints', vartos(var), set_count(constraints))

        for constraint, _ in pairs(constraints) do
            if not constraint(solver, var) then
                return false
            end
        end
    end

    return true
end

function shuffle( array )
    for i = 1, #array-1 do
        local j = math.random(i, #array)
        array[i], array[j] = array[j], array[i]
    end

    return array
end

function randpairs( tbl )
    local kvs = {}

    for k, v in pairs(tbl) do
        kvs[#kvs+1] = { k, v }
    end

    kvs = shuffle(kvs)

    local index = 0

    return
        function ()
            index = index + 1
            local kv = kvs[index]

            if kv then
                return kv[1], kv[2]
            else
                return nil
            end
        end
end

function solve( solver )
    local varsarray = set_toarray(solver.variables)

    -- We want to try and solve the 'smallest' variables first. This also
    -- allows a quick check for any empty variables.
    table.sort(varsarray,
        function ( lhs, rhs )
            return set_count(lhs.values) < set_count(rhs.values)
        end)

    -- If there's any empty variables we're buggered.
    if set_empty(varsarray[1].values) then
        error('empty')
    end

    if solver.random then
        shuffle(varsarray)
    end

    local iterator = (solver.random) and randpairs or pairs

    local function aux( solver, index )
        if index > #varsarray then
            coroutine.yield(solution(solver))

            return
        end

        local variable = varsarray[index]

        -- If the variable is singular (i.e. has a value) we just carry on.
        if set_singular(variable.values) then
            aux(solver, index+1)
        else
            -- local failed = true

            for value, _ in iterator(variable.values) do
                newframe(solver)

                -- printf(' try %s = %s', variable.name, value)

                local success = narrow(solver, variable, set(value))

                if success then
                    -- failed = false
                    -- print(debug.traceback())
                    aux(solver, index+1)
                end

                unwind(solver)
            end

            -- if not failed then
                return
            -- else
                -- error('unsatisfiable')
            -- end
        end
    end

    aux(solver, 1)
end

function printTable( tbl )
    for k, v in pairs(tbl) do
        print(k, v)
    end
end

function environment( varnames, varmap )
    local result = {}

    local repeated = {}

    for _, varname in ipairs(varnames) do
        local variable = varmap[varname]
        assert(variable)
        assert(not repeated[variable], variable.name)
        result[#result+1] = variable
        repeated[variable] = true
    end

    return result
end

function vartos( var, values )
    values = values or var.values
    return string.format("%s=%s", var.name, set_tostring(values))
end

local numcfails = 0
local numfwfails = 0

local numcardsame = 0
local numcarddiff = 0

local _constraints = {
    NotEqual =
        function ( vars )
            assert(#vars == 2)

            local var1, var2 = vars[1], vars[2]

            return
                function ( solver, var )
                    assert(var == var1 or var == var2)
                    local other = (var == var1) and var2 or var1

                    if set_singular(var.values) then
                        local newvalues = set_subtract(other.values, var.values)
                        local oldvalues = other.values

                        -- printf('  NotEqual %s %s => %s', vartos(var), vartos(other, oldvalues), set_tostring(newvalues))

                        local result = narrow(solver, other, newvalues)

                        -- printf('   %s', result and 'succeeded' or 'failed')

                        return result
                    end

                    return true
                end
        end,
    Equal =
        function ( vars )
            assert(#vars == 2)

            local var1, var2 = vars[1], vars[2]

            return
                function ( solver, var )
                    assert(var == var1 or var == var2)
                    local other = (var == var1) and var2 or var1

                    if set_singular(var.values) then
                        local newvalues = set_copy(var.values)
                        local oldvalues = other.values

                        -- printf('  Equal %s %s => %s', vartos(var), vartos(other, oldvalues), set_tostring(newvalues))

                        local result = narrow(solver, other, newvalues)

                        -- printf('   %s', result and 'succeeded' or 'failed')

                        return result
                    end

                    return true
                end
        end,
    Cardinality =
        function ( vars, params )
            local min = params.min
            local max = params.max
            local value = params.value

            assert(0 <= min)
            assert(min <= max)
            assert(value)
            assert(min <= #vars)

            local variables = set_fromarray(varss)

            return
                function ( solver, var )
                    assert(variables[var])

                    local possible = {}
                    local definite = {}
                    local numpossible = 0
                    local numdefinite = 0

                    for variable, _ in pairs(variables) do
                        if variable.values[value] then
                            possible[variable] = true
                            numpossible = numpossible + 1

                            if set_singular(variable.values) then
                                definite[variable] = true
                                numdefinite = numdefinite + 1
                            end
                        end
                    end

                    if numpossible < min or numdefinite > max then
                        numcfails = numcfails + 1
                        return false
                    end

                    -- We have enough possibles, but only just, so try and make
                    -- them definite to be sure.
                    if numpossible == min then
                        for variable, _ in pairs(possible) do
                            local result = narrow(solver, variable, set(value))

                            if not result then
                                numcfails = numcfails + 1
                                return false
                            end
                        end

                        return true
                    end

                    -- We've got the maximum number of definites so make the
                    -- possibles that are not also definite into impossibles.
                    if numdefinite == max then
                        local valueset = set(value)

                        for variable, _ in pairs(possible) do
                            if not definite[variable] then
                                local newvalues = set_subtract(variable.values, valueset)

                                -- printf('  Cardinality %s => %s', vartos(variable), set_tostring(newvalues))

                                local result = narrow(solver, variable, newvalues)

                                if not result then
                                    numcfails = numcfails + 1
                                    return false
                                end
                            end
                        end

                        return true
                    end

                    return true
                end
        end,
    Sum =
        function ( vars, params )
            assert(#vars > 0)

            local total = params.total
            assert(total)

            return
                function ( solver, var )
                    local count = 0

                    -- print(#vars)

                    for _, variable in ipairs(vars) do
                        if not set_singular(variable.values) then
                            return true
                        else
                            -- TODO: Need a 'get unique value' function.
                            count = count + next(variable.values)
                        end
                    end

                    -- printf('  Sum count:%d total:%d', count, total)

                    return count == total
                end
        end,
    FourWayConnected =
        function ( vars, params )
            local height = params.height
            local width = params.width
            local passable = params.passable

            assert(#vars == width * height)

            -- Build up the 'neighbour' relation.
            local neighbours = {}

            local at =
                function ( x, y )
                    local result = x + (y-1)*width
                    assert(1 <= result and result <= width * height)
                    return result
                end

            for y = 1, height do
                for x = 1, width do
                    local peers = {}

                    -- North
                    if y > 1 then
                        peers[#peers+1] = vars[at(x, y-1)]
                    end

                    -- East
                    if x < width then
                        peers[#peers+1] = vars[at(x+1, y)]
                    end

                    -- South
                    if y < height then
                        peers[#peers+1] = vars[at(x, y+1)]
                    end

                    -- West
                    if x > 1 then
                        peers[#peers+1] = vars[at(x-1, y)]
                    end

                    neighbours[vars[at(x, y)]] = peers
                end
            end

            return
                function ( solver, var )
                    for _, variable in ipairs(vars) do
                        if not set_singular(variable.values) then
                            return true
                        end
                    end

                    -- First create a set of passable variables.
                    local passables = {}
                    local numpassables = 0

                    for _, variable in ipairs(vars) do
                        if passable[next(variable.values)] then
                            passables[variable] = true
                            numpassables = numpassables + 1
                        end
                    end

                    if set_count(passables) == 0 then
                        return true
                    else
                        local bfs = function( variable )
                            local result = { [variable] = true }
                            local frontier = { [variable] = true }
                            local count = 1

                            while next(frontier) do
                                local newFrontier = {}

                                for variable, _  in pairs(frontier) do
                                    for _, other in ipairs(neighbours[variable]) do
                                        if not result[other] and not frontier[other] and passables[other] then
                                            count = count + 1
                                            result[other] = true
                                            newFrontier[other] = true
                                        end
                                    end
                                end

                                frontier = newFrontier
                            end

                            return count
                        end

                        local count = bfs(next(passables))

                        if count ~= numpassables then
                            numfwfails = numfwfails + 1

                            if numcfails % 1000 == 0 or numfwfails % 1000 == 0 then
                                -- printf('#card:%d #fourw:%d', numcfails, numfwfails)
                            end
                        end

                        return count == numpassables
                    end
                end
        end,
}

-- tbl = {
--     vars = { [<string>] = { <string> }+ },
--     constraints = { { <string>, vars = { <string> }+ } }*
-- }


function newSolver( tbl )
    local dump = tbl.dump == true

    local vars = tbl.vars

    local variables = {}
    local varmap = {}

    for name, values in pairs(vars) do
        assert(#values >= 1)

        local variable = { name = name, values = set_fromarray(values) }
        variables[variable] = true
        varmap[name] = variable

        if dump then
            print(vartos(variable))
        end
    end

    local constraints = {}

    for variable, _ in pairs(variables) do
        constraints[variable] = {}
    end

    for index, v in ipairs(tbl.constraints) do
        local name = v[1]
        local vars = environment(v.vars, varmap)
        local params = v.params

        local constraint = _constraints[name](vars, params)

        for _, variable, _ in ipairs(vars) do
            if dump then
                printf('constraint #%d %s', index, variable.name)
            end

            constraints[variable][constraint] = true
        end
    end

    local random = tbl.random and true or false

    local result = {
        variables = variables,
        varmap = varmap,
        constraints = constraints,
        stack = {},
        random = random,
    }

    local coro = coroutine.create(function () solve(result) end)

    return
        function ()
            local status, result = coroutine.resume(coro)

            if status then
                return result
            else
                error(result)
                return nil
            end
        end
end

function solution_tostring( solution )
    local sorted = {}

    for varname, value in pairs(solution) do
        sorted[#sorted+1] = { varname, value }
    end

    table.sort(sorted,
        function ( lhs, rhs )
            return lhs[1] < rhs[1]
        end)

    local parts = {}

    for _, data in pairs(sorted) do
        parts[#parts+1] = string.format("%s=%s", data[1], data[2])
    end

    return string.format("{ %s }", table.concat(parts, ' '))
end

function enumerate( solver )
    local count = 0

    for solution in solver do
        count = count + 1
        printf('#%d %s', count, solution_tostring(solution))
    end

    printf('%d solutions', count)
    print()

    return count
end

function enumerate2( solver )
    local count = 0

    for solution in solver do
        count = count + 1

        if count % 1000 == 0 then
            printf('#%d %s', count, solution_tostring(solution))
        end
    end

    printf('%d solutions', count)
    print()
end

local solver = newSolver {
    vars = {
        a = { 'x', 'y', 'z' },
        b = { 'x', 'y', 'z' },
        c = { 'x', 'y', 'z' },
    },
    constraints = {},
}

enumerate(solver)

local solver = newSolver {
    dump = true,
    vars = {
        a = { 'x', 'y', 'z' },
        b = { 'x', 'y', 'z' },
        c = { 'x', 'y', 'z' },
    },
    constraints = {
        { 'NotEqual', vars = { 'a', 'b' } },
        { 'NotEqual', vars = { 'a', 'c' } },
        { 'NotEqual', vars = { 'b', 'c' } },
    },
}

assert(enumerate(solver) == 6)

local solver = newSolver {
    vars = {
        a = { 'x', 'y', 'z' },
        b = { 'x', 'y', 'z' },
        c = { 'x', 'y', 'z' },
    },
    constraints = {
        { 'Equal', vars = { 'a', 'b' } },
        { 'Equal', vars = { 'a', 'c' } },
        { 'Equal', vars = { 'b', 'c' } },
    },
}

assert(enumerate(solver) == 3)

function genNotEquals( varnames )
    local result = {}

    for i = 1, #varnames do
        for j = i+1, #varnames do
            result[#result+1] = { 'NotEqual', vars = { varnames[i], varnames[j] } }
        end
    end

    return result
end

local solver = newSolver {
    vars = {
        a = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
        b = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
        c = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
        d = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
        e = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
        f = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
        g = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
        h = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
        i = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
        j = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
        k = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
        l = { '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',  },
    },
    constraints = genNotEquals { 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l' } ,
}

-- This should have 12! = 479,001,600 solutions.
-- enumerate2(solver)

local solver = newSolver {
    vars = {
        a = { '0', '1', },
        b = { '0', '1', },
        c = { '0', '1', },
        d = { '0', '1', },
    },
    constraints = {
        { 'Cardinality', vars = { 'a', 'b', 'c', 'd' }, params = { min = 2, max = 3, value = '1' } },
    },
}

assert(enumerate(solver) == 10)

local solver = newSolver {
    vars = {
        a = { '0', '1', },
        b = { '0', '1', },
        c = { '0', '1', },
        d = { '0', '1', },
        e = { '0', '1', },
        f = { '0', '1', },
        g = { '0', '1', },
        h = { '0', '1', },
        i = { '0', '1', },
        j = { '0', '1', },
        k = { '0', '1', },
        l = { '0', '1', },
        m = { '0', '1', },
        n = { '0', '1', },
        o = { '0', '1', },
        r = { '0', '1', },
        s = { '0', '1', },
        t = { '0', '1', },
        u = { '0', '1', },
        v = { '0', '1', },
        w = { '0', '1', },
        x = { '0', '1', },
        y = { '0', '1', },
        z = { '0', '1', },
    },
    constraints = {
        { 'Cardinality', vars = { 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z' }, params = { min = 13, max = 13, value = '1' } },
    },
}

-- enumerate2(solver)

local solver = newSolver {
    vars = {
        a = { 0, 1, 2, 3, 4, 5, },
        b = { 0, 1, 2, 3, 4, 5, },
        c = { 0, 1, 2, 3, 4, 5, },
        d = { 0, 1, 2, 3, 4, 5, },
    },
    constraints = {
        { 'Sum', vars = { 'a', 'b', 'c', 'd' }, params = { total = 2 } },
    },
}

enumerate(solver)

local solver = newSolver {
    vars = {
        a11 = { '#', '.' },
        a21 = { '#', '.' },
        a31 = { '#', '.' },
        a12 = { '#', '.' },
        a22 = { '#', '.' },
        a32 = { '#', '.' },
        a13 = { '#', '.' },
        a23 = { '#', '.' },
        a33 = { '#', '.' },
    },
    constraints = {
        { 'FourWayConnected', vars = { 'a11', 'a21', 'a31', 'a12', 'a22', 'a32', 'a13', 'a23', 'a33' }, params = { width = 3, height = 3, passable = { ['.'] = true } } },
    },
}

-- enumerate3(solver)

function enumeratemap( width, height, minf, maxf )
    assert(width >= 3)
    assert(height >= 3)
    assert(0 <= minf and minf <= 1)
    assert(0 <= maxf and maxf <= 1)
    assert(minf <= maxf)

    local min = math.floor((width*height)*minf)
    local max = math.floor((width*height)*maxf)

    local vars = {}
    local varnames = {}
    local domain = { '#', '.' }
    local passable = { ['.'] = true }

    for y = 1, height do
        for x = 1, width do
            local var = string.format("%d_%d", x, y)
            vars[var] = domain
            varnames[#varnames+1] = var
        end
    end

    local solver = newSolver {
        random = true,
        vars = vars,
        constraints = {
            {
                'FourWayConnected',
                vars = varnames,
                params = {
                    width = width,
                    height = height,
                    passable = passable
                },
            },
            {
                'Cardinality',
                vars = varnames,
                params = {
                    min = min,
                    max = max,
                    value = '.'
                }
            }
        },
    }

    local count = 0

    for solution in solver do
        count = count + 1

        printf('#%d', count)

        for y = 1, height do
            local line = {}
            for x = 1, width do
                local var = string.format("%d_%d", x, y)
                line[#line+1] = solution[var]
            end
            print(table.concat(line))
        end

        print()
    end

    printf('%d solutions', count)
    print()

    return count
end

enumeratemap(10, 9, 0.2, 0.5)