A Haskell Sudoku Solver using Finite Domain Constraints

In a previous post we looked at how to write a finite domain constraint solver in Haskell. Now we're going to look at how to use this solver to solve Sudoku puzzles. First we give the module header and define a useful type:

module Sudoku (Puzzle, printSudoku, displayPuzzle, sudoku) where

import Control.Monad
import Data.List (transpose)
import FD

type Puzzle = [Int]

We represent both unsolved and solved puzzles as a list of 81 Ints. In an unsolved puzzle, we use 0 to represent a blank square. The numbers 1 to 9 represent squares with known values. Here is an example of an unsolved puzzle that I copied from the local newspaper:

test :: Puzzle
test = [
    0, 0, 0, 0, 8, 0, 0, 0, 0,
    0, 0, 0, 1, 0, 6, 5, 0, 7,
    4, 0, 2, 7, 0, 0, 0, 0, 0,
    0, 8, 0, 3, 0, 0, 1, 0, 0,
    0, 0, 3, 0, 0, 0, 8, 0, 0,
    0, 0, 5, 0, 0, 9, 0, 7, 0,
    0, 5, 0, 0, 0, 8, 0, 0, 6,
    3, 0, 1, 2, 0, 4, 0, 0, 0,
    0, 0, 6, 0, 1, 0, 0, 0, 0 ]

The function displayPuzzle displays a puzzle for us in rows and columns. The function printSudoku will solve a puzzle by calling sudoku, which we will define below. It then prints each solution.

displayPuzzle :: Puzzle -> String
displayPuzzle = unlines . map show . chunk 9

printSudoku :: Puzzle -> IO ()
printSudoku = putStr . unlines . map displayPuzzle . sudoku

chunk :: Int -> [a] -> [[a]]
chunk _ [] = []
chunk n xs = ys : chunk n zs where
    (ys, zs) = splitAt n xs

We now present the code to actually solve the puzzle:

sudoku :: Puzzle -> [Puzzle]
sudoku puzzle = runFD $ do
    vars <- newVars 81 [1..9]
    zipWithM_ (\x n -> when (n > 0) (x `hasValue` n)) vars puzzle
    mapM_ allDifferent (rows vars)
    mapM_ allDifferent (columns vars)
    mapM_ allDifferent (boxes vars)
    labelling vars

rows, columns, boxes :: [a] -> [[a]]
rows = chunk 9
columns = transpose . rows
boxes = concat . map (map concat . transpose) . chunk 3 . chunk 3 . chunk 3

We start by initialising 81 new solver variables, one for each square in the puzzle. Next, we constrain each known square in the puzzle to its given value. The next three lines create the Sudoku constraints: each number 1 to 9 may occur only once in each row, column and 3x3 box. The functions for grouping the variables into rows, columns and boxes are given below the main function. Finally, we call labelling to search for solutions.

Let's see how this performs with our test puzzle, running on a 2.4GHz Intel Core 2 Duo with 4GB RAM:

> ghc --make -O2 test
[1 of 3] Compiling FD               ( FD.hs, FD.o )
[2 of 3] Compiling Sudoku           ( Sudoku.hs, Sudoku.o )
[3 of 3] Compiling Main             ( test.hs, test.o )
Linking test ...
> time ./test
[5,6,7,4,8,3,2,9,1]
[9,3,8,1,2,6,5,4,7]
[4,1,2,7,9,5,3,6,8]
[6,8,9,3,7,2,1,5,4]
[7,4,3,6,5,1,8,2,9]
[1,2,5,8,4,9,6,7,3]
[2,5,4,9,3,8,7,1,6]
[3,7,1,2,6,4,9,8,5]
[8,9,6,5,1,7,4,3,2]


real 0m6.518s
user 0m6.480s
sys 0m0.032s

So it takes around 6.5 seconds to find all solutions to this puzzle (of which there is exactly one, as expected). Can we do any better than that? Yes we can, actually. Recall our earlier definition of the function different which constrains two variables to have different values, and is used by the allDifferent function which features prominently in our Sudoku solving code:

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

Notice how this function doesn't try to constrain the domains of the variables or do any constraint propagation. That's because, in the general case, we don't have enough information to do this, we just have to store the constraint and retest it each time the domains are changed. This means that in our Sudoku solver, the labelling step will effectively try each value 1 to 9 in turn for each blank square and then test whether the allDifferent constraints still hold. This is almost a brute force algorithm.

However, we have overlooked one case where we can further constrain the variables: if the domain of one variable is a singleton set, then we can remove that value from the domain of the other variable. This should drastically reduce the number of tests we need to do during labelling. Here's the modified function:

different = addBinaryConstraint $ \x y -> do
    xv <- lookup x
    yv <- lookup y
    guard $ IntSet.size xv > 1 || IntSet.size yv > 1 || xv /= yv
    when (IntSet.size xv == 1 && xv `IntSet.isProperSubsetOf` yv) $
        update y (yv `IntSet.difference` xv)
    when (IntSet.size yv == 1 && yv `IntSet.isProperSubsetOf` xv) $
        update x (xv `IntSet.difference` yv)

Here are the times with the modified funtion:

> time ./test
[5,6,7,4,8,3,2,9,1]
[9,3,8,1,2,6,5,4,7]
[4,1,2,7,9,5,3,6,8]
[6,8,9,3,7,2,1,5,4]
[7,4,3,6,5,1,8,2,9]
[1,2,5,8,4,9,6,7,3]
[2,5,4,9,3,8,7,1,6]
[3,7,1,2,6,4,9,8,5]
[8,9,6,5,1,7,4,3,2]


real 0m0.180s
user 0m0.160s
sys 0m0.020s

That's a 36-fold improvement. Not bad for adding four lines of code!

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.