Constraint Programming in Haskell

Having had a bit of free time in the evenings lately, I've finally got around to starting one of the things on my list of cool Haskell projects I'd like to get around to one day. Since my time working on Mercury and HAL, I've been interested in logic programming and constraint logic programming and I've been wanting to find out how close I can get to something that looks like a constraint logic programming system in Haskell.

The Haskell list monad already allows a style of programming that feels somewhat like non-deterministic logic programming, with backtracking to do a depth-first left-to-right search. I'm going to build a monad on top of the list monad to do finite domain constraint programming. I'm aiming for a nice clean interface and a simple implementation, so this is not going to be as efficient as it could be, but that's ok for a proof of concept. If I get time later, I'll go back and try to make it more efficient.

In a finite domain solver the variables represent a finite set of (usually integer) values to which the programmer can add various constraints. They are particularly useful for scheduling and timetabling problems. They are also good for solving puzzles such as Sudoku, which I will detail in a later post.

An example of the kind of constraint program we're aiming to be able to write is

runTest = runFD test

test = do
    x <- newVar [0..3]
    y <- newVar [0..3]
    (( x .<. y) `mplus` (x `same` y))
    x `hasValue` 2
    labelling [x, y]

Here we create two finite domain variables x and y, each with an initial domain of {0, 1, 2, 3}. We then add a constraint that says either x is less than y or x is the same as y. We then further constrain x to have the value 2. Finally, the call to labelling will search for all possible solutions that satisfy the given constraints. When we evaluate runTest we get the result [[2,3],[2,2]] which represents the two possible solutions {x = 2, y = 3} and {x = 2, y = 2}.

Finite domain solvers also allow constraints to contain arithmetic expressions involving constraint variables and integers. To keep things simple, we'll leave them out for now. The interface we are going to implement is shown below.

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes #-}
module FD (
    -- Types
    FD,           -- Monad for finite domain constraint solver
    FDVar,        -- Finite domain solver variable

    -- Functions
    runFD,        -- Run the monad and return a list of solutions.
    newVar,       -- Create a new FDVar
    newVars,      -- Create multiple FDVars
    hasValue,     -- Constrain a FDVar to a specific value
    same,         -- Constrain two FDVars to be the same
    different,    -- Constrain two FDVars to be different
    allDifferent, -- Constrain a list of FDVars to be different
    (.<.),        -- Constrain one FDVar to be less than another
    labelling     -- Backtracking search for all solutions
    ) where

Modules we need to import:

import Prelude hiding (lookup)
import Control.Monad.State.Lazy
import Control.Monad.Trans
import qualified Data.Map as Map
import Data.Map ((!), Map)
import qualified Data.IntSet as IntSet
import Data.IntSet (IntSet)

Below we define the types for the solver.

-- The FD monad
newtype FD s a = FD { unFD :: StateT (FDState s) [] a }
    deriving (Monad, MonadPlus, MonadState (FDState s))

-- FD variables
newtype FDVar s = FDVar { unFDVar :: Int } deriving (Ord, Eq)

type VarSupply s = FDVar s
data VarInfo s = VarInfo
     { delayedConstraints :: FD s (), values :: IntSet }
type VarMap s = Map (FDVar s) (VarInfo s)
data FDState s = FDState
     { varSupply :: VarSupply s, varMap :: VarMap s }

The type FD s a is our constraint solver monad. It contains a list monad to provide the search capability. This is wrapped in a StateT monad transformer which threads through the constraint store FDState s. The type variable s is a phantom type. We will see later how this can be used to prevent any of the implementation detail "leaking out" of the monad.

Our constraint store FDState s contains a supply of fresh constraint variables and also keeps track of the information we need to know about existing variables. For each existing variable we record its set of possible values (its domain) and a set of constraints on it. Whenever the domain of a variable changes, we need to execute its constraints to check that they are still satisfied. This, in turn, may further constrain the domain of other variables. This is known as constraint propagation.

-- Run the FD monad and produce a lazy list of possible solutions.
runFD :: (forall s . FD s a) -> [a]
runFD fd = evalStateT (unFD fd) initState

initState :: FDState s
initState = FDState { varSupply = FDVar 0, varMap = Map.empty }

The function runFD runs a constraint solver, starting with an initially empty constraint store, and return a list of all possible solutions. The type (forall s . FD s a) -> [a] ensures that any values of a type containing the phantom type variable s can't "escape" from the monad. This means that we can't take a constraint variable from one monad and use it inside another one, thus ensuring, through the type system, that the monad is used safely.

-- Get a new FDVar
newVar :: [Int] -> FD s (FDVar s)
newVar domain= do
    v <- nextVar
    v `isOneOf` domain
    return v
    where
        nextVar :: FD s (FDVar s)
        nextVar = do
            s <- get
            let v = varSupply s
            put $ s { varSupply = FDVar (unFDVar v + 1) }
            return v
        isOneOf :: FDVar s -> [Int] -> FD s ()
        x `isOneOf` domain=
            modify $ \s ->
                let vm = varMap s
                    vi = VarInfo {
                        delayedConstraints = return (),
                        values = IntSet.fromList domain}
                in
                s { varMap = Map.insert x vi vm }

newVars :: Int -> [Int] -> FD s [FDVar s]
newVars n domain = replicateM n (newVar domain)

The function newVar domain creates a new constraint variable constrained to values in domain. The function newVars n domain is a convenient way of creating multiple variables with the same domain.

Some helper functions which are not exported, but are used when we define the constraint functions:

-- Lookup the current domain of a variable.
lookup :: FDVar s -> FD s IntSet
lookup x = do
    s <- get
    return . values $ varMap s ! x

-- Update the domain of a variable and fire all delayed constraints
-- associated with that variable.
update :: FDVar s -> IntSet -> FD s ()
update x i = do
    s <- get
    let vm = varMap s
    let vi = vm ! x
    put $ s { varMap = Map.insert x (vi { values = i}) vm }
    delayedConstraints vi

-- Add a new constraint for a variable to the constraint store.
addConstraint :: FDVar s -> FD s () -> FD s ()
addConstraint x constraint = do
    s <- get
    let vm = varMap s
    let vi = vm ! x
    let cs = delayedConstraints vi
    put $ s { varMap =
        Map.insert x (vi { delayedConstraints = cs >> constraint }) vm }
 
-- Useful helper function for adding binary constraints between FDVars.
type BinaryConstraint s = FDVar s -> FDVar s -> FD s ()
addBinaryConstraint :: BinaryConstraint s -> BinaryConstraint s
addBinaryConstraint f x y = do
    let constraint  = f x y
    constraint
    addConstraint x constraint
    addConstraint y constraint

The function lookup returns the current domain for a variable; update updates the domain for a variable and propagates the change into all constraints on that variable; addConstraint inserts a constraint into the constraint store; addBinaryConstraint tests a constraint on two variable and then adds it to the constraint store for each variable.

Now we can define the actual constraint functions:

-- Constrain a variable to a particular value.
hasValue :: FDVar s -> Int -> FD s ()
var `hasValue` val = do
    vals <- lookup var
    guard $ val `IntSet.member` vals
    let i = IntSet.singleton val
    when (i /= vals) $ update var i

In hasValue we lookup the current domain of the variable and test that the value to be set is within the domain. If the domain has changed, we update the constraint store and propagate the change. The other constraints are defined similarly:

-- Constrain two variables to have the same value.
same :: FDVar s -> FDVar s -> FD s ()
same = addBinaryConstraint $ \x y -> do
    xv <- lookup x
    yv <- lookup y
    let i = IntSet.intersection xv yv
    guard $ not $ IntSet.null i
    when (i /= xv) $ update x i
    when (i /= yv) $ update y i

-- Constrain two variables to have different values.
different :: FDVar s -> FDVar s -> FD s ()
different = addBinaryConstraint $ \x y -> do
    xv <- lookup x
    yv <- lookup y
    guard $ IntSet.size xv > 1 || IntSet.size yv > 1 || xv /= yv

-- Constrain a list of variables to all have different values.
allDifferent :: [FDVar s] -> FD s ()
allDifferent (x:xs) = do
    mapM_ (different x) xs
    allDifferent xs
allDifferent _ = return ()

-- Constrain one variable to have a value less than the value of another
-- variable.
(.<.) :: FDVar s -> FDVar s -> FD s ()
(.<.) = addBinaryConstraint $ \x y -> do
    xv <- lookup x
    yv <- lookup y
    let xv' = IntSet.filter (< IntSet.findMax yv) xv
    let yv' = IntSet.filter (> IntSet.findMin xv) yv
    guard $ not $ IntSet.null xv'
    guard $ not $ IntSet.null yv'
    when (xv /= xv') $ update x xv'
    when (yv /= yv') $ update y yv'

Finally, in the labelling function we make use of the underlying list monad to search for all solutions for the given set of variables.

-- Label variables using a depth-first left-to-right search.
labelling :: [FDVar s] -> FD s [Int]
labelling = mapM label where
    label var = do
        vals <- lookup var
        val <- FD . lift $ IntSet.toList vals
        var `hasValue` val
        return val

In a later posts I plan to show how to use this finite domain solver monad to write a solver for Sudoku puzzles, and extend the monad to support arithmetic expressions.