SBV Sudoku Solver

Problem96.hs

module Problem96
    ( answer
    ) where

import           Data.Foldable (for_)
import           Data.List (transpose)
import qualified Data.Map.Strict as M
import           Data.SBV
import           Problem96Input (grids)

toMap :: [[a]] -> M.Map (Int, Int) a
toMap vs = M.fromList $ do
  (y, row)   <- zip [0..] vs
  (x, value) <- zip [0..] row
  pure ((x, y), value)

unknownValue (x, y) = do
  v <- sInteger $ show x ++ "," ++ show y
  constrain $ 0 .< v .&& v .< 10
  pure v

rows :: M.Map (Int, Int) a -> [[a]]
rows m =
  map (\y -> map (\x -> m M.! (x, y)) [0..8]) [0..8]

columns :: M.Map (Int, Int) a -> [[a]]
columns = transpose . rows

blocks :: M.Map (Int, Int) a -> [[a]]
blocks m =
  concatMap (\h -> map (\w -> aux (w, h)) [0..2]) [0..2]

  where
    aux (w, h) =
      map (\(dw, dh) -> m M.! (w*3+dw, h*3+dh)) [
        (0,0), (1, 0), (2, 0),
        (0,1), (1, 1), (2, 1),
        (0,2), (1, 2), (2, 2)
      ]

solveGrid grid = sat $ do
  let knowns = toMap grid
  solveSpace <- mapM sequence $ map (\y -> map (\x -> unknownValue (x, y)) [0..8]) [0..8]
  let solutionLookup = toMap solveSpace

  -- Known values come from input
  for_ (filter ((/= 0) . snd) $ M.toList knowns) $ \(pt, value) ->
    constrain $ solutionLookup M.! pt .== literal value

  for_ (rows solutionLookup) $ constrain . distinct

  for_ (columns solutionLookup) $ constrain . distinct

  for_ (blocks solutionLookup) $ constrain . distinct

upperDigits :: [[Integer]] -> IO Integer
upperDigits grid = do
  sol <- solveGrid grid
  case (getModelValue "0,0" sol :: Maybe Integer, getModelValue "1,0" sol, getModelValue "2,0" sol) of
    (Just a, Just b, Just c) -> pure $ a*100+b*10+c
    _ ->
      error "No result!"

puzzle :: Integer -> [[[Integer]]] -> IO Integer
puzzle acc [] = pure acc
puzzle acc (g:gs) = do
  v <- upperDigits g
  puzzle (acc+v) gs

answer :: IO ()
answer =
  puzzle 0 grids >>= print