Ersatzを使って数独を解く

ErsatzはSATソルバの入力をモナドの力で生成してくれるライブラリ。適当にexistsとassertを並べればうまくやってくれるらしい。ということで、数独ソルバを実装してみた。素晴らしい…

sudoku.hs

import Prelude hiding ((||), (&&), not, and, or)
import Data.Char
import Data.Array
import Data.Ix
import Data.Word
import Control.Monad
import Control.Monad.State
import Ersatz -- (ersatz-0.2)

fieldRange = ((1,1),(9,9))

unique :: Equatable a => [a] -> Bit
unique = go [] where
    go ss (x:xs) = and (map (/==x) ss) && go (x : ss) xs
    go _ [] = true

sudoku :: (MonadState s m, HasSAT s) => Array (Int, Int) Char
    -> Array (Int, Int) Bit4 -> m (Array (Int, Int) Bit4)
sudoku problem answer = do
    forM_ (range fieldRange) $ \ix -> assert $ if isDigit (problem ! ix)
        -- 既に埋められているマスを確定
        then answer ! ix === encode (toEnum $ digitToInt $ problem ! ix)
        -- 空いているマスには1から9までの数字が入る
        else let Bit4 b3 b2 b1 b0 = answer ! ix in or
                [ b3 && not b2 && not b1 -- 8,9
                , not b3 && (b2 || b1 || b0) -- excepting 0
                ]
    -- 行が一意
    forM_ [1..9] $ \r -> assert $ unique [answer ! (c, r) | c <- [1..9]]
    -- 列が一意
    forM_ [1..9] $ \c -> assert $ unique [answer ! (c, r) | r <- [1..9]]
    -- ブロック内が一意
    forM_ (range ((0,0),(2,2))) $ \(sc, sr) ->
        assert $ unique [answer ! (sc * 3 + c, sc * 3 + r) | c <- [1..3], r <- [1..3]]
    return answer

showField :: (Result, Maybe (Array (Int, Int) Word8)) -> IO ()
showField (_, Just answer) = forM_ [1..9] $ \r -> do
    forM_ [1..9] $ \c -> putChar $ intToDigit $ fromEnum $ answer ! (c, r)
    putChar '\n'
showField _ = fail "Impossible"

main :: IO ()
main = do
    problem <- liftM (array fieldRange . concat)
        $ forM [1..9] $ \r -> do
            line <- getLine
            return [((c, r), ch) | (c, ch) <- zip [1..9] line]
    let sat = replicateM 81 exists >>= sudoku problem . listArray fieldRange
    minisatPath "cryptominisat" `solveWith` sat >>= showField

$ runhaskell sudoku.hs
-9---3---
-7----8--
---5--2--
5--------
-----7-91
2---4----
----6-4--
6--2-----
-------7-
192783645
375691824
864539217
548912763
426857391
217346958
739165482
681274539
953428176