-----------------------------------------------------------------------------
-- |
-- Module    : Documentation.SBV.Examples.Puzzles.Tower
-- Copyright : (c) Levent Erkok
-- License   : BSD3
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- Solves the tower puzzle, <http://www.chiark.greenend.org.uk/%7Esgtatham/puzzles/js/towers.html>.
-----------------------------------------------------------------------------

{-# OPTIONS_GHC -Wall -Werror #-}

module Documentation.SBV.Examples.Puzzles.Tower where

import Control.Monad
import Data.Array hiding (inRange)
import Data.SBV
import Data.SBV.Control

-------------------------------------------------------------------
-- * Modeling Towers
-------------------------------------------------------------------

-- | Count of visible towers as an array.
type Count a = Array Integer a

-- | The grid itself. The indexes are tuples, first coordinate increases as you go from
-- left to right, and the second increases as you go from top to bottom.
type Grid a = Array (Integer, Integer) a

-- | The problem has 4 counts, from top, left, bottom, and right. And the grid itself.
type Problem a = (Count a, Count a, Count a, Count a, Grid a)

-- | Example problem. Encodes:
--
-- @
--     - - 3 - - 4
--   -     2       5
--   -   2         -
--   4             -
--   2             -
--   -             2
--   3             -
--     - - 3 4 - -
-- @
problem :: Problem (Maybe Integer)
problem :: Problem (Maybe Integer)
problem = (Array Integer (Maybe Integer)
top, Array Integer (Maybe Integer)
left, Array Integer (Maybe Integer)
bot, Array Integer (Maybe Integer)
right, Array (Integer, Integer) (Maybe Integer)
grid)
  where build :: (i, i) -> [(i, a)] -> Array i (Maybe a)
build (i, i)
ix [(i, a)]
es = (Maybe a -> Maybe a -> Maybe a)
-> Maybe a -> (i, i) -> [(i, Maybe a)] -> Array i (Maybe a)
forall i e a.
Ix i =>
(e -> a -> e) -> e -> (i, i) -> [(i, a)] -> Array i e
accumArray (\Maybe a
_ Maybe a
a -> Maybe a
a) Maybe a
forall a. Maybe a
Nothing (i, i)
ix [(i
i, a -> Maybe a
forall a. a -> Maybe a
Just a
v) | (i
i, a
v) <- [(i, a)]
es]

        top :: Array Integer (Maybe Integer)
top   = (Integer, Integer)
-> [(Integer, Integer)] -> Array Integer (Maybe Integer)
forall {i} {a}. Ix i => (i, i) -> [(i, a)] -> Array i (Maybe a)
build  (Integer
1, Integer
6)          [(Integer
3, Integer
3), (Integer
6, Integer
4)]
        left :: Array Integer (Maybe Integer)
left  = (Integer, Integer)
-> [(Integer, Integer)] -> Array Integer (Maybe Integer)
forall {i} {a}. Ix i => (i, i) -> [(i, a)] -> Array i (Maybe a)
build  (Integer
1, Integer
6)          [(Integer
3, Integer
4), (Integer
4, Integer
2), (Integer
6, Integer
3)]
        bot :: Array Integer (Maybe Integer)
bot   = (Integer, Integer)
-> [(Integer, Integer)] -> Array Integer (Maybe Integer)
forall {i} {a}. Ix i => (i, i) -> [(i, a)] -> Array i (Maybe a)
build  (Integer
1, Integer
6)          [(Integer
3, Integer
3), (Integer
4, Integer
4)]
        right :: Array Integer (Maybe Integer)
right = (Integer, Integer)
-> [(Integer, Integer)] -> Array Integer (Maybe Integer)
forall {i} {a}. Ix i => (i, i) -> [(i, a)] -> Array i (Maybe a)
build  (Integer
1, Integer
6)          [(Integer
1, Integer
5), (Integer
5, Integer
2)]
        grid :: Array (Integer, Integer) (Maybe Integer)
grid  = ((Integer, Integer), (Integer, Integer))
-> [((Integer, Integer), Integer)]
-> Array (Integer, Integer) (Maybe Integer)
forall {i} {a}. Ix i => (i, i) -> [(i, a)] -> Array i (Maybe a)
build ((Integer
1, Integer
1), (Integer
6, Integer
6)) [((Integer
3, Integer
1), Integer
2), ((Integer
2, Integer
2), Integer
2)]

-- | Given a concrete partial board, turn it into a symbolic board, by filling in the
-- empty cells with symbolic variables.
symProblem :: Problem (Maybe Integer) -> Symbolic (Problem SInteger)
symProblem :: Problem (Maybe Integer) -> Symbolic (Problem SInteger)
symProblem (Array Integer (Maybe Integer)
t, Array Integer (Maybe Integer)
l, Array Integer (Maybe Integer)
b, Array Integer (Maybe Integer)
r, Array (Integer, Integer) (Maybe Integer)
g) = (,,,,) (Count SInteger
 -> Count SInteger
 -> Count SInteger
 -> Count SInteger
 -> Grid SInteger
 -> Problem SInteger)
-> SymbolicT IO (Count SInteger)
-> SymbolicT
     IO
     (Count SInteger
      -> Count SInteger
      -> Count SInteger
      -> Grid SInteger
      -> Problem SInteger)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array Integer (Maybe Integer) -> SymbolicT IO (Count SInteger)
forall (f :: * -> *).
Traversable f =>
f (Maybe Integer) -> Symbolic (f SInteger)
fill Array Integer (Maybe Integer)
t SymbolicT
  IO
  (Count SInteger
   -> Count SInteger
   -> Count SInteger
   -> Grid SInteger
   -> Problem SInteger)
-> SymbolicT IO (Count SInteger)
-> SymbolicT
     IO
     (Count SInteger
      -> Count SInteger -> Grid SInteger -> Problem SInteger)
forall a b.
SymbolicT IO (a -> b) -> SymbolicT IO a -> SymbolicT IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Array Integer (Maybe Integer) -> SymbolicT IO (Count SInteger)
forall (f :: * -> *).
Traversable f =>
f (Maybe Integer) -> Symbolic (f SInteger)
fill Array Integer (Maybe Integer)
l SymbolicT
  IO
  (Count SInteger
   -> Count SInteger -> Grid SInteger -> Problem SInteger)
-> SymbolicT IO (Count SInteger)
-> SymbolicT
     IO (Count SInteger -> Grid SInteger -> Problem SInteger)
forall a b.
SymbolicT IO (a -> b) -> SymbolicT IO a -> SymbolicT IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Array Integer (Maybe Integer) -> SymbolicT IO (Count SInteger)
forall (f :: * -> *).
Traversable f =>
f (Maybe Integer) -> Symbolic (f SInteger)
fill Array Integer (Maybe Integer)
b SymbolicT IO (Count SInteger -> Grid SInteger -> Problem SInteger)
-> SymbolicT IO (Count SInteger)
-> SymbolicT IO (Grid SInteger -> Problem SInteger)
forall a b.
SymbolicT IO (a -> b) -> SymbolicT IO a -> SymbolicT IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Array Integer (Maybe Integer) -> SymbolicT IO (Count SInteger)
forall (f :: * -> *).
Traversable f =>
f (Maybe Integer) -> Symbolic (f SInteger)
fill Array Integer (Maybe Integer)
r SymbolicT IO (Grid SInteger -> Problem SInteger)
-> SymbolicT IO (Grid SInteger) -> Symbolic (Problem SInteger)
forall a b.
SymbolicT IO (a -> b) -> SymbolicT IO a -> SymbolicT IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Array (Integer, Integer) (Maybe Integer)
-> SymbolicT IO (Grid SInteger)
forall (f :: * -> *).
Traversable f =>
f (Maybe Integer) -> Symbolic (f SInteger)
fill Array (Integer, Integer) (Maybe Integer)
g
 where fill :: Traversable f => f (Maybe Integer) -> Symbolic (f SInteger)
       fill :: forall (f :: * -> *).
Traversable f =>
f (Maybe Integer) -> Symbolic (f SInteger)
fill = (Maybe Integer -> SymbolicT IO SInteger)
-> f (Maybe Integer) -> SymbolicT IO (f SInteger)
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) -> f a -> m (f b)
mapM (SymbolicT IO SInteger
-> (Integer -> SymbolicT IO SInteger)
-> Maybe Integer
-> SymbolicT IO SInteger
forall b a. b -> (a -> b) -> Maybe a -> b
maybe SymbolicT IO SInteger
forall a. SymVal a => Symbolic (SBV a)
free_ (SInteger -> SymbolicT IO SInteger
forall a. a -> SymbolicT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SInteger -> SymbolicT IO SInteger)
-> (Integer -> SInteger) -> Integer -> SymbolicT IO SInteger
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> SInteger
forall a. SymVal a => a -> SBV a
literal))

-------------------------------------------------------------------
-- * Counting visible towers
-------------------------------------------------------------------

-- | Given a list of tower heights, count the number of visible ones in the given order.
-- We simply keep track of the tallest we have seen so far, and increment the count for
-- each tower we see if it's taller than the tallest seen so far.
visible :: [SInteger] -> SInteger
visible :: [SInteger] -> SInteger
visible = SInteger -> SInteger -> [SInteger] -> SInteger
forall {a} {t}.
(OrdSymbolic a, Num t, Mergeable t) =>
a -> t -> [a] -> t
go SInteger
0 SInteger
0
 where go :: a -> t -> [a] -> t
go a
_            t
visibleSofar []     = t
visibleSofar
       go a
tallestSofar t
visibleSofar (a
x:[a]
xs) = a -> t -> [a] -> t
go (a
tallestSofar a -> a -> a
forall a. OrdSymbolic a => a -> a -> a
`smax` a
x)
                                                (SBool -> t -> t -> t
forall a. Mergeable a => SBool -> a -> a -> a
ite (a
x a -> a -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.> a
tallestSofar) (t
1 t -> t -> t
forall a. Num a => a -> a -> a
+ t
visibleSofar) t
visibleSofar)
                                                [a]
xs

-------------------------------------------------------------------
-- * Building constraints
-------------------------------------------------------------------

-- | Build the constraints for a given problem. We scan the elements and add the required
-- visibility counts for each row and column, viewed both in the correct order and in the backwards order.
tower :: Problem SInteger -> Symbolic ()
tower :: Problem SInteger -> Symbolic ()
tower (Count SInteger
top, Count SInteger
left, Count SInteger
bot, Count SInteger
right, Grid SInteger
grid) = do
  let (Integer
minX, Integer
maxX) = Count SInteger -> (Integer, Integer)
forall i e. Array i e -> (i, i)
bounds Count SInteger
top
      (Integer
minY, Integer
maxY) = Count SInteger -> (Integer, Integer)
forall i e. Array i e -> (i, i)
bounds Count SInteger
left

  -- Constraints from top and bottom
  [Integer] -> (Integer -> Symbolic ()) -> Symbolic ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Integer
minX .. Integer
maxX] ((Integer -> Symbolic ()) -> Symbolic ())
-> (Integer -> Symbolic ()) -> Symbolic ()
forall a b. (a -> b) -> a -> b
$ \Integer
x -> do
      let reqT :: SInteger
reqT = Count SInteger
top Count SInteger -> Integer -> SInteger
forall i e. Ix i => Array i e -> i -> e
! Integer
x
          reqB :: SInteger
reqB = Count SInteger
bot Count SInteger -> Integer -> SInteger
forall i e. Ix i => Array i e -> i -> e
! Integer
x
          elts :: [SInteger]
elts = [Grid SInteger
grid Grid SInteger -> (Integer, Integer) -> SInteger
forall i e. Ix i => Array i e -> i -> e
! (Integer
x, Integer
y) | Integer
y <- [Integer
minY .. Integer
maxY]]
      (SInteger -> Symbolic ()) -> [SInteger] -> Symbolic ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\SInteger
e -> SBool -> Symbolic ()
forall a. QuantifiedBool a => a -> Symbolic ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain (SInteger -> (SInteger, SInteger) -> SBool
forall a. OrdSymbolic a => a -> (a, a) -> SBool
inRange SInteger
e (Integer -> SInteger
forall a. SymVal a => a -> SBV a
literal Integer
1, Integer -> SInteger
forall a. SymVal a => a -> SBV a
literal Integer
maxY))) [SInteger]
elts
      SBool -> Symbolic ()
forall a. QuantifiedBool a => a -> Symbolic ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain (SBool -> Symbolic ()) -> SBool -> Symbolic ()
forall a b. (a -> b) -> a -> b
$ [SInteger] -> SBool
forall a. EqSymbolic a => [a] -> SBool
distinct [SInteger]
elts
      SBool -> Symbolic ()
forall a. QuantifiedBool a => a -> Symbolic ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain (SBool -> Symbolic ()) -> SBool -> Symbolic ()
forall a b. (a -> b) -> a -> b
$ SInteger
reqT SInteger -> SInteger -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== [SInteger] -> SInteger
visible [SInteger]
elts
      SBool -> Symbolic ()
forall a. QuantifiedBool a => a -> Symbolic ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain (SBool -> Symbolic ()) -> SBool -> Symbolic ()
forall a b. (a -> b) -> a -> b
$ SInteger
reqB SInteger -> SInteger -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== [SInteger] -> SInteger
visible ([SInteger] -> [SInteger]
forall a. [a] -> [a]
reverse [SInteger]
elts)

  -- Constraints from left and right
  [Integer] -> (Integer -> Symbolic ()) -> Symbolic ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Integer
minY .. Integer
maxY] ((Integer -> Symbolic ()) -> Symbolic ())
-> (Integer -> Symbolic ()) -> Symbolic ()
forall a b. (a -> b) -> a -> b
$ \Integer
y -> do
      let reqL :: SInteger
reqL = Count SInteger
left  Count SInteger -> Integer -> SInteger
forall i e. Ix i => Array i e -> i -> e
! Integer
y
          reqR :: SInteger
reqR = Count SInteger
right Count SInteger -> Integer -> SInteger
forall i e. Ix i => Array i e -> i -> e
! Integer
y
          elts :: [SInteger]
elts = [Grid SInteger
grid Grid SInteger -> (Integer, Integer) -> SInteger
forall i e. Ix i => Array i e -> i -> e
! (Integer
x, Integer
y) | Integer
x <- [Integer
minX .. Integer
maxX]]
      (SInteger -> Symbolic ()) -> [SInteger] -> Symbolic ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\SInteger
e -> SBool -> Symbolic ()
forall a. QuantifiedBool a => a -> Symbolic ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain (SInteger -> (SInteger, SInteger) -> SBool
forall a. OrdSymbolic a => a -> (a, a) -> SBool
inRange SInteger
e (Integer -> SInteger
forall a. SymVal a => a -> SBV a
literal Integer
1, Integer -> SInteger
forall a. SymVal a => a -> SBV a
literal Integer
maxX))) [SInteger]
elts
      SBool -> Symbolic ()
forall a. QuantifiedBool a => a -> Symbolic ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain (SBool -> Symbolic ()) -> SBool -> Symbolic ()
forall a b. (a -> b) -> a -> b
$ [SInteger] -> SBool
forall a. EqSymbolic a => [a] -> SBool
distinct [SInteger]
elts
      SBool -> Symbolic ()
forall a. QuantifiedBool a => a -> Symbolic ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain (SBool -> Symbolic ()) -> SBool -> Symbolic ()
forall a b. (a -> b) -> a -> b
$ SInteger
reqL SInteger -> SInteger -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== [SInteger] -> SInteger
visible [SInteger]
elts
      SBool -> Symbolic ()
forall a. QuantifiedBool a => a -> Symbolic ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain (SBool -> Symbolic ()) -> SBool -> Symbolic ()
forall a b. (a -> b) -> a -> b
$ SInteger
reqR SInteger -> SInteger -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== [SInteger] -> SInteger
visible ([SInteger] -> [SInteger]
forall a. [a] -> [a]
reverse [SInteger]
elts)

-------------------------------------------------------------------
-- * Example run
-------------------------------------------------------------------

-- | Solve the puzzle descibed above. We get:
--
-- >>> example
--   1 2 3 2 2 4
-- 1 6 5 2 4 3 1 5
-- 3 3 2 5 6 1 4 2
-- 4 2 4 1 5 6 3 2
-- 2 5 3 6 1 4 2 3
-- 2 1 6 4 3 2 5 2
-- 3 4 1 3 2 5 6 1
--   3 2 3 4 2 1
example :: IO ()
example :: IO ()
example = Symbolic () -> IO ()
forall a. Symbolic a -> IO a
runSMT (Symbolic () -> IO ()) -> Symbolic () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        sp <- Problem (Maybe Integer) -> Symbolic (Problem SInteger)
symProblem Problem (Maybe Integer)
problem
        tower sp
        query $ do cs <- checkSat
                   case cs of
                     CheckSatResult
Unsat -> IO () -> Query ()
forall a. IO a -> Query a
io (IO () -> Query ()) -> IO () -> Query ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStrLn String
"Unsolvable"
                     CheckSatResult
Sat   -> Problem SInteger -> Query ()
display Problem SInteger
sp
                     CheckSatResult
_     -> String -> Query ()
forall a. HasCallStack => String -> a
error (String -> Query ()) -> String -> Query ()
forall a b. (a -> b) -> a -> b
$ String
"Unexpected result: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ CheckSatResult -> String
forall a. Show a => a -> String
show CheckSatResult
cs
 where display :: Problem SInteger -> Query ()
       display :: Problem SInteger -> Query ()
display (Count SInteger
top, Count SInteger
left, Count SInteger
bot, Count SInteger
right, Grid SInteger
grid) = do
          let (Integer
minX, Integer
maxX) = Count SInteger -> (Integer, Integer)
forall i e. Array i e -> (i, i)
bounds Count SInteger
top
              (Integer
minY, Integer
maxY) = Count SInteger -> (Integer, Integer)
forall i e. Array i e -> (i, i)
bounds Count SInteger
left

          -- Display top row
          IO () -> Query ()
forall a. IO a -> Query a
io (IO () -> Query ()) -> IO () -> Query ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStr String
"  "
          topVals <- [Integer] -> (Integer -> QueryT IO Integer) -> QueryT IO [Integer]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Integer
minX .. Integer
maxX] ((Integer -> QueryT IO Integer) -> QueryT IO [Integer])
-> (Integer -> QueryT IO Integer) -> QueryT IO [Integer]
forall a b. (a -> b) -> a -> b
$ \Integer
x -> SInteger -> QueryT IO Integer
forall a. SymVal a => SBV a -> Query a
getValue (Count SInteger
top Count SInteger -> Integer -> SInteger
forall i e. Ix i => Array i e -> i -> e
! Integer
x)
          io $ putStrLn $ unwords (map show topVals)

          -- Display each row, sandwiched between left/right
          forM_ [minY .. maxY] $ \Integer
y -> do
             lv <- SInteger -> QueryT IO Integer
forall a. SymVal a => SBV a -> Query a
getValue (Count SInteger
left  Count SInteger -> Integer -> SInteger
forall i e. Ix i => Array i e -> i -> e
! Integer
y)
             rv <- getValue (right ! y)
             row <- forM [minX .. maxX] $ \Integer
x -> SInteger -> QueryT IO Integer
forall a. SymVal a => SBV a -> Query a
getValue (Grid SInteger
grid Grid SInteger -> (Integer, Integer) -> SInteger
forall i e. Ix i => Array i e -> i -> e
! (Integer
x, Integer
y))
             io $ putStrLn $ unwords (map show (lv : row ++ [rv]))

          -- Finish with bottom row
          io $ putStr "  "
          botVals <- forM [minX .. maxX] $ \Integer
x -> SInteger -> QueryT IO Integer
forall a. SymVal a => SBV a -> Query a
getValue (Count SInteger
bot Count SInteger -> Integer -> SInteger
forall i e. Ix i => Array i e -> i -> e
! Integer
x)
          io $ putStrLn $ unwords (map show botVals)