-----------------------------------------------------------------------------
-- |
-- Module    : Documentation.SBV.Examples.Puzzles.NQueens
-- Copyright : (c) Levent Erkok
-- License   : BSD3
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- Solves the NQueens puzzle: <http://en.wikipedia.org/wiki/Eight_queens_puzzle>
-----------------------------------------------------------------------------

{-# OPTIONS_GHC -Wall -Werror #-}

module Documentation.SBV.Examples.Puzzles.NQueens where

import Data.SBV

-- | A solution is a sequence of row-numbers where queens should be placed
type Solution = [SWord8]

-- | Checks that a given solution of @n@-queens is valid, i.e., no queen
-- captures any other.
isValid :: Int -> Solution -> SBool
isValid :: Int -> Solution -> SBool
isValid Int
n Solution
s = (SWord8 -> SBool) -> Solution -> SBool
forall a. (a -> SBool) -> [a] -> SBool
sAll SWord8 -> SBool
forall a. (OrdSymbolic a, Num a) => a -> SBool
rangeFine Solution
s SBool -> SBool -> SBool
.&& Solution -> SBool
forall a. EqSymbolic a => [a] -> SBool
distinct Solution
s SBool -> SBool -> SBool
.&& ((Int, Int) -> SBool) -> [(Int, Int)] -> SBool
forall a. (a -> SBool) -> [a] -> SBool
sAll (Int, Int) -> SBool
checkDiag [(Int, Int)]
ijs
  where rangeFine :: a -> SBool
rangeFine a
x = a
x a -> a -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.>= a
1 SBool -> SBool -> SBool
.&& a
x a -> a -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.<= Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
        ijs :: [(Int, Int)]
ijs = [(Int
i, Int
j) | Int
i <- [Int
1..Int
n], Int
j <- [Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1..Int
n]]
        checkDiag :: (Int, Int) -> SBool
checkDiag (Int
i, Int
j) = SWord8
diffR SWord8 -> SWord8 -> SBool
forall a. EqSymbolic a => a -> a -> SBool
./= SWord8
diffC
           where qi :: SWord8
qi = Solution
s Solution -> Int -> SWord8
forall a. [a] -> Int -> a
!! (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                 qj :: SWord8
qj = Solution
s Solution -> Int -> SWord8
forall a. [a] -> Int -> a
!! (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                 diffR :: SWord8
diffR = SBool -> SWord8 -> SWord8 -> SWord8
forall a. Mergeable a => SBool -> a -> a -> a
ite (SWord8
qi SWord8 -> SWord8 -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.>= SWord8
qj) (SWord8
qiSWord8 -> SWord8 -> SWord8
forall a. Num a => a -> a -> a
-SWord8
qj) (SWord8
qjSWord8 -> SWord8 -> SWord8
forall a. Num a => a -> a -> a
-SWord8
qi)
                 diffC :: SWord8
diffC = Int -> SWord8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i)

-- | Given @n@, it solves the @n-queens@ puzzle, printing all possible solutions.
nQueens :: Int -> IO ()
nQueens :: Int -> IO ()
nQueens Int
n
 | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"n must be non-negative, received: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n
 | Bool
True  = do String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Finding all " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"-queens solutions.."
              AllSatResult
res <- SymbolicT IO SBool -> IO AllSatResult
forall a. Provable a => a -> IO AllSatResult
allSat (SymbolicT IO SBool -> IO AllSatResult)
-> SymbolicT IO SBool -> IO AllSatResult
forall a b. (a -> b) -> a -> b
$ Int -> Solution -> SBool
isValid Int
n (Solution -> SBool) -> SymbolicT IO Solution -> SymbolicT IO SBool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Int -> SymbolicT IO Solution
forall a. SymVal a => Int -> Symbolic [SBV a]
mkExistVars Int
n
              Int
cnt <- ([(Bool, [Word8])] -> [(Bool, [Word8])])
-> (Int -> (Bool, [Word8]) -> IO ()) -> AllSatResult -> IO Int
forall a.
SatModel a =>
([(Bool, a)] -> [(Bool, a)])
-> (Int -> (Bool, a) -> IO ()) -> AllSatResult -> IO Int
displayModels [(Bool, [Word8])] -> [(Bool, [Word8])]
forall a. a -> a
id Int -> (Bool, [Word8]) -> IO ()
forall a a. Show a => a -> (a, [Word8]) -> IO ()
disp AllSatResult
res
              String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Found: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
cnt String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" solution(s)."
   where disp :: a -> (a, [Word8]) -> IO ()
disp a
i (a
_, [Word8]
s) = do String -> IO ()
putStr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Solution #" String -> String -> String
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
i String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": "
                            [Word8] -> IO ()
dispSolution [Word8]
s
         dispSolution :: [Word8] -> IO ()
         dispSolution :: [Word8] -> IO ()
dispSolution [Word8]
model
           | Int
lmod Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n = String -> IO ()
forall a. HasCallStack => String -> a
error (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Impossible! Backend solver returned " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
lmod String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" values, was expecting: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n
           | Bool
True      = do String -> IO ()
putStr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ [Word8] -> String
forall a. Show a => a -> String
show [Word8]
model
                            String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
" (Valid: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SBool -> String
forall a. Show a => a -> String
show (Int -> Solution -> SBool
isValid Int
n ((Word8 -> SWord8) -> [Word8] -> Solution
forall a b. (a -> b) -> [a] -> [b]
map Word8 -> SWord8
forall a. SymVal a => a -> SBV a
literal [Word8]
model)) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
           where lmod :: Int
lmod  = [Word8] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word8]
model