Tower example

[LeventErkok]

Solution to the tower puzzle https://www.chiark.greenend.org.uk/%7Esgtatham/puzzles/js/towers.html

Would make a nice example to add in the next release

{-# OPTIONS_GHC -Wall -Werror #-}

module Tower where

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

type Count   a = Array Integer            a
type Grid    a = Array (Integer, Integer) a
type Problem a = (Count a, Count a, Count a, Count a, Grid a)

-- Left to right: x increases
-- Top to bottom: y increases
problem :: Problem (Maybe Integer)
problem = (top, left, bot, right, grid)
  where build ix es = accumArray (\_ a -> a) Nothing ix [(i, Just v) | (i, v) <- es]

        top   = build  (1, 6)          [(3, 3), (6, 4)]
        left  = build  (1, 6)          [(3, 4), (4, 2), (6, 3)]
        bot   = build  (1, 6)          [(3, 3), (4, 4)]
        right = build  (1, 6)          [(1, 5), (5, 2)]
        grid  = build ((1, 1), (6, 6)) [((3, 1), 2), ((2, 2), 2)]

symProblem :: Problem (Maybe Integer) -> Symbolic (Problem SInteger)
symProblem (t, l, b, r, g) = (,,,,) <$> fill t <*> fill l <*> fill b <*> fill r <*> fill g
 where fill :: Traversable f => f (Maybe Integer) -> Symbolic (f SInteger)
       fill = mapM (maybe free_ (pure . literal))

visible :: [SInteger] -> SInteger
visible = go 0 0
 where go _            visibleSofar []     = visibleSofar
       go tallestSofar visibleSofar (x:xs) = go (tallestSofar `smax` x)
                                                (ite (x .> tallestSofar) (1 + visibleSofar) visibleSofar)
                                                xs

tower :: Problem SInteger -> Symbolic ()
tower (top, left, bot, right, grid) = do
  let (minX, maxX) = bounds top
      (minY, maxY) = bounds left

  forM_ [minX .. maxX] $ \x -> do
      let reqT = top ! x
          reqB = bot ! x
          elts = [grid ! (x, y) | y <- [minY .. maxY]]
      mapM_ (\e -> constrain (inRange e (literal 1, literal maxY))) elts
      constrain $ distinct elts
      constrain $ reqT .== visible elts
      constrain $ reqB .== visible (reverse elts)

  forM_ [minY .. maxY] $ \y -> do
      let reqL = left  ! y
          reqR = right ! y
          elts = [grid ! (x, y) | x <- [minX .. maxX]]
      mapM_ (\e -> constrain (inRange e (literal 1, literal maxX))) elts
      constrain $ distinct elts
      constrain $ reqL .== visible elts
      constrain $ reqR .== visible (reverse elts)

  pure ()

main :: IO ()
main = runSMT $ do
        sp <- symProblem problem
        tower sp
        query $ do cs <- checkSat
                   case cs of
                     Unsat -> io $ putStrLn "Unsolvable"
                     Sat   -> walk sp
                     _     -> error $ "Unexpected result: " ++ show cs

 where walk :: Problem SInteger -> Query ()
       walk (top, left, bot, right, grid) = do
          let (minX, maxX) = bounds top
              (minY, maxY) = bounds left

              sh x = do vx <- getValue x
                        io $ putStr $ show vx ++ " "

          io $ putStr "  "
          forM_ [minX .. maxX] $ \x -> sh (top ! x)
          io $ putStrLn ""

          forM_ [minY .. maxY] $ \y -> do
             sh (left ! y)
             forM_ [minX .. maxX] $ \x -> do
               sh (grid ! (x, y))
             sh (right ! y)
             io $ putStrLn ""

          io $ putStr "  "
          forM_ [minX .. maxX] $ \x -> sh (bot ! x)
          io $ putStrLn ""