Solving Musihikui-zan (arithmetical restorations) using SBV (SMT front-end for Haskell)

Mushikui.hs

{-# LANGUAGE DeriveGeneric, FlexibleInstances, GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses, NoMonomorphismRestriction             #-}
{-# LANGUAGE OverloadedStrings, PatternGuards, RecordWildCards            #-}
{-# LANGUAGE TypeFamilies, ViewPatterns                                   #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Main where
import Control.Lens  (ifoldlM, iforM_, imap, imapM, ix, (^?))
import Data.Char     (digitToInt, intToDigit, isDigit)
import Data.Foldable (foldlM, foldrM)
import Data.Maybe    (catMaybes, mapMaybe)
import Data.SBV      (Predicate, SBV, SBool, SWord8, SymWord, Symbolic, Word8)
import Data.SBV      (allSat, bAll, bAnd, constrain, exists, extractModels)
import Data.SBV      (false, isVacuous, literal, sDivMod, true, (&&&), (./=))
import Data.SBV      ((.<=), (.==))
import Data.String   (IsString (..))

type Cell =  SWord8

data Setting = Setting { factorA  :: [Cell]
                       , factorB  :: [Cell]
                       , summands :: [[Cell]]
                       , answer   :: [Cell]
                       }
             deriving (Show)


data Board = Board { facA :: [Maybe Word8]
                   , facB :: [Maybe Word8]
                   , sums :: [[Maybe Word8]]
                   , resl :: [Maybe Word8]
                   }

instance IsString [Maybe Word8] where
  fromString = map step . filter (`elem` (['0'..'9'] ++ "_"))
    where
      step n | isDigit n = Just $ toEnum $ digitToInt n
             | otherwise = Nothing

instance Show Board where
  show = init . prettyBoard

quantify :: (Num a, SymWord a) => String -> Maybe a -> Symbolic (SBV a)
quantify _ (Just n) = return (literal n)
quantify s Nothing = do
  x <- exists s
  constrain $ 0 .<= x &&& x .<= 9
  return x

prettyCell :: Maybe Word8 -> Char
prettyCell Nothing  = '_'
prettyCell (Just i) = intToDigit $ fromEnum i

prettyLine :: [Maybe Word8] -> String
prettyLine = map prettyCell

prettyBoard :: Board -> String
prettyBoard Board{..} =
  let width = maximum (map length $ facA : facB : resl : imap (\i -> (replicate i Nothing ++)) sums) + 3
      pad n s = replicate (n - length s) ' ' ++ s
  in unlines $ map (pad width) $
     [prettyLine facA ++ " "
     ,"x " ++ pad (width - 2) (prettyLine facB ++ " ")
     ,replicate width '~'
     ] ++ imap (\i -> (++ replicate (i+1) ' ') . prettyLine) sums
       ++ [replicate width '~'
          ,prettyLine resl ++ " "]

buildSetting :: Board -> Symbolic Setting
buildSetting Board{..} = do
  factorA  <- imapM (\i -> quantify $ "A_" ++ show (i+1)) $ reverse facA
  factorB  <- imapM (\i -> quantify $ "B_" ++ show (i+1)) $ reverse facB
  summands <- imapM (\i -> imapM (\j -> quantify $ "sum_" ++ show (i+1) ++ '_' : show (j+1))) $
              map reverse sums
  answer   <- imapM (\i -> quantify $ "M_" ++ show (i+1)) $ reverse resl
  constrain $ bAll (\a -> last a ./= 0) (factorA : factorB : answer : summands)
  return Setting{..}

sumsUp :: Setting -> Symbolic ()
sumsUp Setting{..} = do
  c <- ifoldlM step 0 answer
  constrain (c .== 0)
  where
    step i c a = do
      let ixs   = catMaybes [summands ^? ix x . ix (i - x) | x <- [0..i]]
          adds  = sum ixs + c
          (upper, lower) = adds `sDivMod` 10
      constrain $ lower .== a
      return upper

problem :: Board -> Predicate
problem ip = problem' =<< buildSetting ip

problem' :: Setting -> Predicate
problem' bd = do
  sumsUp bd
  multLines bd
  return true

multLines :: Setting -> Symbolic ()
multLines Setting{..} = do
  iforM_ factorB $ \i b -> do
    let targ = summands !! i
        cnt = length factorA
        step j co a = do
          let (q, r) = (b * a + co) `sDivMod` 10
          case targ ^? ix j of
            Nothing -> constrain false >> return 0
            Just el -> constrain (r .== el) >> return q
    c <- ifoldlM step 0 factorA
    e <- foldlM (\co a -> do let (q,r) = co `sDivMod` 10
                             constrain (a .== r)
                             return q)
         c (drop cnt targ)
    constrain $ e .== 0

readBoard :: Board -> [Word8] -> Maybe Board
readBoard (Board factorA factorB summands answer) cws0 = do
  (facA, restA) <- matching factorA cws0
  (facB, restB) <- matching factorB restA
  (sums, restS) <- foldlM matchings ([], restB) summands
  (resl, _cws)   <- matching answer restS
  return Board{..}
  where
    matchings (xs,rs) ys = do
      (as, rs') <- matching ys rs
      return (xs ++ [as], rs')
    matching xs cws = foldrM match ([], cws) xs
    match c (cs, rs)
      | Just l <- c = return (Just l:cs,rs)
      | r : rs' <- rs = return (Just r : cs, rs')
      | otherwise = Nothing

main :: IO ()
main = do
  mapM_ print =<< getAllSolutions question

getAllSolutions :: Board -> IO [Board]
getAllSolutions puz = do
  let q = problem puz
  unsol <- isVacuous q
  if unsol
    then return []
    else mapMaybe (readBoard puz) . extractModels <$> allSat q

question :: Board
question = Board  "  __"
                  "x 2_"
                 ["  __",
                  "  __"]
                  "__1_"


lonely1 :: Board
lonely1 = Board "  __"
                "x __"
               ["  1_",
                " __ "]
                "____"

lonely5 :: Board
lonely5 = Board "  _5"
                "x __"
               ["  __",
                " __ "]
                "____"

lonely7_1 :: Board
lonely7_1 = Board  "  __"
                   "x 7_"
                  [" ___",
                   " __ "]
                   " ___"

lonely7_2 :: Board
lonely7_2 = Board  "  __"
                   "x __"
                  ["  __",
                   " _7 "]
                   "____"

lonely8_1 :: Board
lonely8_1 = Board  "  __"
                   "x 8_"
                  [" ___",
                   " __ "]
                   "____"

lonely8_2 :: Board
lonely8_2 = Board  "  __"
                   "x _8"
                  ["  __",
                   "___ "]
                   "____"

lonely9_1 :: Board
lonely9_1 = Board  "  9_"
                   "x __"
                  ["  __",
                   " __ "]
                   " ___"

lonely9_2 :: Board
lonely9_2 = Board  "  __"
                   "x _9"
                  ["  __",
                   " __ "]
                   "____"

lonely9_3 :: Board
lonely9_3 = Board  "  __"
                   "x __"
                  ["  __",
                   "___ "]
                   "9___"

lonely7_1_ans :: Board
lonely7_1_ans = Board "12" "79" ["108", "84"] "948"

overlaps :: Setting -> Board -> SBool
overlaps Setting{..} Board{..} =
  bAnd $ zipWith eql facA (reverse factorA)
      ++ zipWith eql facB (reverse factorB)
      ++ zipWith ((bAnd .) . zipWith eql) sums (map reverse summands)
      ++ zipWith eql resl answer
  where
    eql ma b = maybe true (\a -> literal a .== b) ma