-----------------------------------------------------------------------------
-- |
-- Module    : Documentation.SBV.Examples.Puzzles.Sudoku
-- Copyright : (c) Levent Erkok
-- License   : BSD3
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- The Sudoku solver, quintessential SMT solver example!
-----------------------------------------------------------------------------

{-# LANGUAGE CPP              #-}
{-# LANGUAGE FlexibleContexts #-}

{-# OPTIONS_GHC -Wall -Werror #-}

module Documentation.SBV.Examples.Puzzles.Sudoku where

#if MIN_VERSION_base(4,18,0)
import Control.Monad (when, zipWithM_)
#endif

import Control.Monad.State.Lazy

import Data.List     (transpose)

import Data.SBV
import Data.SBV.Control

-------------------------------------------------------------------
-- * Modeling Sudoku
-------------------------------------------------------------------
-- | A row is a sequence of digits that we represent symbolic integers
type Row = [SInteger]

-- | A Sudoku board is a sequence of 9 rows
type Board = [Row]

-- | Given a series of elements, make sure they are all different
-- and they all are numbers between 1 and 9
check :: [SInteger] -> SBool
check :: [SInteger] -> SBool
check [SInteger]
grp = [SBool] -> SBool
sAnd ([SBool] -> SBool) -> [SBool] -> SBool
forall a b. (a -> b) -> a -> b
$ [SInteger] -> SBool
forall a. EqSymbolic a => [a] -> SBool
distinct [SInteger]
grp SBool -> [SBool] -> [SBool]
forall a. a -> [a] -> [a]
: (SInteger -> SBool) -> [SInteger] -> [SBool]
forall a b. (a -> b) -> [a] -> [b]
map SInteger -> SBool
forall {a}. (OrdSymbolic a, Num a) => a -> SBool
rangeFine [SInteger]
grp
  where rangeFine :: a -> SBool
rangeFine a
x = a
x a -> (a, a) -> SBool
forall a. OrdSymbolic a => a -> (a, a) -> SBool
`inRange` (a
1, a
9)

-- | Given a full Sudoku board, check that it is valid
valid :: Board -> SBool
valid :: Board -> SBool
valid Board
rows = [SBool] -> SBool
sAnd ([SBool] -> SBool) -> [SBool] -> SBool
forall a b. (a -> b) -> a -> b
$ Bool -> SBool
forall a. SymVal a => a -> SBV a
literal Bool
sizesOK SBool -> [SBool] -> [SBool]
forall a. a -> [a] -> [a]
: ([SInteger] -> SBool) -> Board -> [SBool]
forall a b. (a -> b) -> [a] -> [b]
map [SInteger] -> SBool
check (Board
rows Board -> Board -> Board
forall a. [a] -> [a] -> [a]
++ Board
columns Board -> Board -> Board
forall a. [a] -> [a] -> [a]
++ Board
squares)
  where sizesOK :: Bool
sizesOK = Board -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Board
rows Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
9 Bool -> Bool -> Bool
&& ([SInteger] -> Bool) -> Board -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\[SInteger]
r -> [SInteger] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SInteger]
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
9) Board
rows

        columns :: Board
columns = Board -> Board
forall a. [[a]] -> [[a]]
transpose Board
rows
        regions :: [Board]
regions = [Board] -> [Board]
forall a. [[a]] -> [[a]]
transpose [Int -> [SInteger] -> Board
forall a. Int -> [a] -> [[a]]
chunk Int
3 [SInteger]
row | [SInteger]
row <- Board
rows]
        squares :: Board
squares = [Board -> [SInteger]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat Board
sq | Board
sq <- Int -> Board -> [Board]
forall a. Int -> [a] -> [[a]]
chunk Int
3 ([Board] -> Board
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Board]
regions)]

        chunk :: Int -> [a] -> [[a]]
        chunk :: forall a. Int -> [a] -> [[a]]
chunk Int
_ [] = []
        chunk Int
i [a]
xs = let ([a]
f, [a]
r) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
i [a]
xs in [a]
f [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chunk Int
i [a]
r

-- | A puzzle is simply a list of rows. Put 0 to indicate blanks.
type Puzzle = [[Integer]]

-------------------------------------------------------------------
-- * Solving Sudoku puzzles
-------------------------------------------------------------------

-- | Fill a given board, replacing 0's with appropriate elements to solve the puzzle
fillBoard :: Puzzle -> IO Puzzle
fillBoard :: Puzzle -> IO Puzzle
fillBoard Puzzle
board = Symbolic Puzzle -> IO Puzzle
forall a. Symbolic a -> IO a
runSMT (Symbolic Puzzle -> IO Puzzle) -> Symbolic Puzzle -> IO Puzzle
forall a b. (a -> b) -> a -> b
$ do
     let emptyCellCount :: Int
emptyCellCount = [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Integer] -> Int) -> [Integer] -> Int
forall a b. (a -> b) -> a -> b
$ ([Integer] -> [Integer]) -> Puzzle -> [Integer]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Integer -> Bool) -> [Integer] -> [Integer]
forall a. (a -> Bool) -> [a] -> [a]
filter (Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0)) Puzzle
board
     subst <- Int -> Symbolic [SInteger]
forall a. SymVal a => Int -> Symbolic [SBV a]
mkFreeVars Int
emptyCellCount
     constrain $ valid (fill literal subst)

     query $ do cs <- checkSat
                case cs of
                  CheckSatResult
Sat   -> do vals <- (SInteger -> QueryT IO Integer)
-> [SInteger] -> QueryT IO [Integer]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SInteger -> QueryT IO Integer
forall a. SymVal a => SBV a -> Query a
getValue [SInteger]
subst
                              pure $ fill id vals
                  CheckSatResult
Unsat -> [Char] -> Query Puzzle
forall a. HasCallStack => [Char] -> a
error [Char]
"Unsolvable puzzle!"
                  CheckSatResult
_     -> [Char] -> Query Puzzle
forall a. HasCallStack => [Char] -> a
error ([Char] -> Query Puzzle) -> [Char] -> Query Puzzle
forall a b. (a -> b) -> a -> b
$ [Char]
"Solver said: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ CheckSatResult -> [Char]
forall a. Show a => a -> [Char]
show CheckSatResult
cs

 where fill :: (Integer -> b) -> [b] -> [[b]]
fill Integer -> b
xform = State [b] [[b]] -> [b] -> [[b]]
forall s a. State s a -> s -> a
evalState (([Integer] -> StateT [b] Identity [b]) -> Puzzle -> State [b] [[b]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((Integer -> StateT [b] Identity b)
-> [Integer] -> StateT [b] Identity [b]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Integer -> StateT [b] Identity b
forall {m :: * -> *}. MonadState [b] m => Integer -> m b
replace) Puzzle
board)
         where replace :: Integer -> m b
replace Integer
0 = do supply <- m [b]
forall s (m :: * -> *). MonadState s m => m s
get
                              case supply of
                                []     -> [Char] -> m b
forall a. HasCallStack => [Char] -> a
error [Char]
"Run out of supplies while filling in the board!"
                                (b
s:[b]
ss) -> [b] -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put [b]
ss m () -> m b -> m b
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
s
               replace Integer
n = b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> m b) -> b -> m b
forall a b. (a -> b) -> a -> b
$ Integer -> b
xform Integer
n

-- | Solve a given puzzle and print the results
sudoku :: Puzzle -> IO ()
sudoku :: Puzzle -> IO ()
sudoku Puzzle
board = Puzzle -> IO Puzzle
fillBoard Puzzle
board IO Puzzle -> (Puzzle -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Puzzle -> IO ()
displayBoard
 where displayBoard :: Puzzle -> IO ()
       displayBoard :: Puzzle -> IO ()
displayBoard Puzzle
puzzle = do
            let sh :: a -> a -> [Char]
sh       a
i a
r = a -> [Char]
forall a. Show a => a -> [Char]
show a
r [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ if a
i a -> [a] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [a
3, a
6] then [Char]
" " else [Char]
""
                printRow :: a -> [b] -> IO ()
printRow a
i [b]
r = do [Char] -> IO ()
putStrLn ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char]
"    " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [[Char]] -> [Char]
unwords ((Int -> b -> [Char]) -> [Int] -> [b] -> [[Char]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> b -> [Char]
forall {a} {a}. (Show a, Eq a, Num a) => a -> a -> [Char]
sh [(Int
1::Int)..] [b]
r)
                                  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (a
i a -> [a] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [a
3, a
6]) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char] -> IO ()
putStrLn [Char]
""
            (Int -> [Integer] -> IO ()) -> [Int] -> Puzzle -> IO ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ Int -> [Integer] -> IO ()
forall {b} {a}. (Show b, Eq a, Num a) => a -> [b] -> IO ()
printRow [(Int
1::Int)..] Puzzle
puzzle

            let isValid :: SBool
isValid = Board -> SBool
valid (([Integer] -> [SInteger]) -> Puzzle -> Board
forall a b. (a -> b) -> [a] -> [b]
map ((Integer -> SInteger) -> [Integer] -> [SInteger]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> SInteger
forall a. SymVal a => a -> SBV a
literal) Puzzle
puzzle)
            case SBool -> Maybe Bool
forall a. SymVal a => SBV a -> Maybe a
unliteral SBool
isValid of
               Just Bool
True  -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
               Just Bool
False -> [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid solution generated!"
               Maybe Bool
Nothing    -> [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Impossible happened, got a symbolic result for valid."

-------------------------------------------------------------------
-- * Example boards
-------------------------------------------------------------------

-- | A random puzzle, found on the internet..
puzzle1 :: Puzzle
puzzle1 :: Puzzle
puzzle1 = [ [Integer
0, Integer
6, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
1, Integer
0]
          , [Integer
0, Integer
0, Integer
0,   Integer
6, Integer
5, Integer
1,   Integer
0, Integer
0, Integer
0]
          , [Integer
1, Integer
0, Integer
7,   Integer
0, Integer
0, Integer
0,   Integer
6, Integer
0, Integer
2]

          , [Integer
6, Integer
2, Integer
0,   Integer
3, Integer
0, Integer
5,   Integer
0, Integer
9, Integer
4]
          , [Integer
0, Integer
0, Integer
3,   Integer
0, Integer
0, Integer
0,   Integer
2, Integer
0, Integer
0]
          , [Integer
4, Integer
8, Integer
0,   Integer
9, Integer
0, Integer
7,   Integer
0, Integer
3, Integer
6]

          , [Integer
9, Integer
0, Integer
6,   Integer
0, Integer
0, Integer
0,   Integer
4, Integer
0, Integer
8]
          , [Integer
0, Integer
0, Integer
0,   Integer
7, Integer
9, Integer
4,   Integer
0, Integer
0, Integer
0]
          , [Integer
0, Integer
5, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
7, Integer
0] ]

-- | Another random puzzle, found on the internet..
puzzle2 :: Puzzle
puzzle2 :: Puzzle
puzzle2 = [ [Integer
1, Integer
0, Integer
3,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
8, Integer
0]
          , [Integer
0, Integer
0, Integer
6,   Integer
0, Integer
4, Integer
8,   Integer
0, Integer
0, Integer
0]
          , [Integer
0, Integer
4, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
0]

          , [Integer
2, Integer
0, Integer
0,   Integer
0, Integer
9, Integer
6,   Integer
1, Integer
0, Integer
0]
          , [Integer
0, Integer
9, Integer
0,   Integer
8, Integer
0, Integer
1,   Integer
0, Integer
4, Integer
0]
          , [Integer
0, Integer
0, Integer
4,   Integer
3, Integer
2, Integer
0,   Integer
0, Integer
0, Integer
8]

          , [Integer
0, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
7, Integer
0]
          , [Integer
0, Integer
0, Integer
0,   Integer
1, Integer
5, Integer
0,   Integer
4, Integer
0, Integer
0]
          , [Integer
0, Integer
6, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
2, Integer
0, Integer
3] ]

-- | Another random puzzle, found on the internet..
puzzle3 :: Puzzle
puzzle3 :: Puzzle
puzzle3 = [ [Integer
6, Integer
0, Integer
0,   Integer
0, Integer
1, Integer
0,   Integer
5, Integer
0, Integer
0]
          , [Integer
8, Integer
0, Integer
3,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
0]
          , [Integer
0, Integer
0, Integer
0,   Integer
0, Integer
6, Integer
0,   Integer
0, Integer
2, Integer
0]

          , [Integer
0, Integer
3, Integer
0,   Integer
1, Integer
0, Integer
8,   Integer
0, Integer
9, Integer
0]
          , [Integer
1, Integer
0, Integer
0,   Integer
0, Integer
9, Integer
0,   Integer
0, Integer
0, Integer
4]
          , [Integer
0, Integer
5, Integer
0,   Integer
2, Integer
0, Integer
3,   Integer
0, Integer
1, Integer
0]

          , [Integer
0, Integer
7, Integer
0,   Integer
0, Integer
3, Integer
0,   Integer
0, Integer
0, Integer
0]
          , [Integer
0, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
3, Integer
0, Integer
6]
          , [Integer
0, Integer
0, Integer
4,   Integer
0, Integer
5, Integer
0,   Integer
0, Integer
0, Integer
9] ]

-- | According to the web, this is the toughest.
-- sudoku puzzle ever.. It even has a name: Al Escargot:
-- <http://zonkedyak.blogspot.com/2006/11/worlds-hardest-sudoku-puzzle-al.html>
puzzle4 :: Puzzle
puzzle4 :: Puzzle
puzzle4 = [ [Integer
1, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
7,   Integer
0, Integer
9, Integer
0]
          , [Integer
0, Integer
3, Integer
0,   Integer
0, Integer
2, Integer
0,   Integer
0, Integer
0, Integer
8]
          , [Integer
0, Integer
0, Integer
9,   Integer
6, Integer
0, Integer
0,   Integer
5, Integer
0, Integer
0]

          , [Integer
0, Integer
0, Integer
5,   Integer
3, Integer
0, Integer
0,   Integer
9, Integer
0, Integer
0]
          , [Integer
0, Integer
1, Integer
0,   Integer
0, Integer
8, Integer
0,   Integer
0, Integer
0, Integer
2]
          , [Integer
6, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
4,   Integer
0, Integer
0, Integer
0]

          , [Integer
3, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
1, Integer
0]
          , [Integer
0, Integer
4, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
7]
          , [Integer
0, Integer
0, Integer
7,   Integer
0, Integer
0, Integer
0,   Integer
3, Integer
0, Integer
0] ]

-- | This one has been called diabolical, apparently
puzzle5 :: Puzzle
puzzle5 :: Puzzle
puzzle5 = [ [ Integer
0, Integer
9, Integer
0,   Integer
7, Integer
0, Integer
0,   Integer
8, Integer
6, Integer
0]
          , [ Integer
0, Integer
3, Integer
1,   Integer
0, Integer
0, Integer
5,   Integer
0, Integer
2, Integer
0]
          , [ Integer
8, Integer
0, Integer
6,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
0]

          , [ Integer
0, Integer
0, Integer
7,   Integer
0, Integer
5, Integer
0,   Integer
0, Integer
0, Integer
6]
          , [ Integer
0, Integer
0, Integer
0,   Integer
3, Integer
0, Integer
7,   Integer
0, Integer
0, Integer
0]
          , [ Integer
5, Integer
0, Integer
0,   Integer
0, Integer
1, Integer
0,   Integer
7, Integer
0, Integer
0]

          , [ Integer
0, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
1, Integer
0, Integer
9]
          , [ Integer
0, Integer
2, Integer
0,   Integer
6, Integer
0, Integer
0,   Integer
3, Integer
5, Integer
0]
          , [ Integer
0, Integer
5, Integer
4,   Integer
0, Integer
0, Integer
8,   Integer
0, Integer
7, Integer
0] ]

-- | Another example
puzzle6 :: Puzzle
puzzle6 :: Puzzle
puzzle6 = [ [Integer
0, Integer
0, Integer
0,   Integer
0, Integer
6, Integer
0,   Integer
0, Integer
8, Integer
0]
          , [Integer
0, Integer
2, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
0]
          , [Integer
0, Integer
0, Integer
1,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
0]

          , [Integer
0, Integer
7, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
1, Integer
0, Integer
2]
          , [Integer
5, Integer
0, Integer
0,   Integer
0, Integer
3, Integer
0,   Integer
0, Integer
0, Integer
0]
          , [Integer
0, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
4, Integer
0, Integer
0]

          , [Integer
0, Integer
0, Integer
4,   Integer
2, Integer
0, Integer
1,   Integer
0, Integer
0, Integer
0]
          , [Integer
3, Integer
0, Integer
0,   Integer
7, Integer
0, Integer
0,   Integer
6, Integer
0, Integer
0]
          , [Integer
0, Integer
0, Integer
0,   Integer
0, Integer
0, Integer
0,   Integer
0, Integer
5, Integer
0] ]

-- | Solve them all, this takes a fraction of a second to run for each case
allPuzzles :: IO ()
allPuzzles :: IO ()
allPuzzles = (Puzzle -> IO ()) -> [Puzzle] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Puzzle -> IO ()
sudoku [Puzzle
puzzle1, Puzzle
puzzle2, Puzzle
puzzle3, Puzzle
puzzle4, Puzzle
puzzle5, Puzzle
puzzle6]