Haskell Quiz/Verbal Arithmetic/Solution Dolio

module Main(main) where
 
import Control.Monad
 
import Data.Char
import Data.List
import Data.Maybe
 
import Data.Map (Map)
import qualified Data.Map as M
 
import Text.ParserCombinators.Parsec
 
import CSP
 
data Equation = Equation { summands :: [String], result :: String } deriving Show
 
-- A quick parser for the equations in use. I didn't bother trying to cover all
-- eventualities, but it will parse 'word1 + word2 + word3 ... = result'
add = char '+' >> spaces >> return (++)
equ = char '=' >> spaces
vnum = many1 lower >>= \a -> spaces >> return a
 
equation = do l <- chainl1 (liftM return vnum) add
              equ
              r <- vnum
              return (Equation l r)
 
-- Constraints on individual digits. Digits in a column, plus the carry-in, mod 10
-- must equal the result digit. Returns the carry-out
digitcsp :: [Char] -> Int -> Char -> CSP Int Int
digitcsp vs carry result = do l <- mapM byName vs
                              let (d, m) = divMod (carry + sum l) 10
                              byName result >>= guard . (m ==)
                              return d
 
-- Constraint to make sure the left-most digit on any term is not zero. We don't want
-- 'two = 012', for instance.
sigfig :: Equation -> CSP Int ()
sigfig (Equation ss r) = chk r >> foldM_ (const chk) () ss
 where
 chk (h:_) = byName h >>= guard . (0 /=)
 
-- The overall constraint for solving equations
eqncsp :: Equation -> CSP Int ()
eqncsp e@(Equation s r@(hr:_)) = go (columns s) (reverse r) 0 >> sigfig e
 where
 go (vs:vss) (r:rs) carry = digitcsp vs carry r >>= go vss rs
 go []       [r]    carry = byName r >>= guard . (carry ==)
 go []       []     0     = return ()
 go _        _      _     = mzero
 
-- turns a list of terms into a list of letters in each column, to be summed.
columns :: [String] -> [[Char]]
columns = transpose . map reverse
 
-- Given an equation, and mappings for all the letters, constructs the numerals
showEqn :: Equation -> Map Char Int -> String
showEqn (Equation ss r) m = intercalate " + " (map showWord ss) ++ " = " ++ showWord r
 where
 showWord = map (\d -> intToDigit (fromJust (M.lookup d m)))
 
main = mapM_ (display . fmap solve . parse equation "foo") . lines =<< getContents
 where
 display (Left _) = putStrLn "Bad parse."
 display (Right ss) = mapM_ putStrLn ss
 
solve :: Equation -> [String]
solve e = map (showEqn e) ms
 where ms = execCSP (eqncsp e) [0..9]

{-# OPTIONS_GHC -fglasgow-exts #-}
 
module CSP where
 
import Data.Map (Map)
import qualified Data.Map as M
 
import Control.Arrow
import Control.Monad
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Trans
 
newtype CSP b a = CSP { unCSP :: StateT (Map Char b) (SupplyT b []) a }
                    deriving (Monad, MonadPlus, MonadState (Map Char b),
                                MonadSupply b)
 
newtype SupplyT s m a = SupplyT { unSupplyT :: StateT [s] m a }
                        deriving (Monad, MonadPlus, MonadState [s])
 
class MonadPlus m => MonadSupply s m where
    getSupply :: m s
 
instance MonadPlus m => MonadSupply s (SupplyT s m) where
    getSupply = do l <- get
                   (v, l') <- choose (pick l)
                   put l'
                   return v
     where
     pick (h:t) = (h, t) : map (second (h:)) (pick t)
     pick []    = mzero
 
instance MonadSupply s m => MonadSupply s (StateT r m) where
    getSupply = lift getSupply
 
runSupplyT :: Monad m => SupplyT s m a -> [s] -> m a
runSupplyT st s = flip evalStateT s . unSupplyT $ st
 
runCSP :: CSP b a -> [b] -> [(a, Map Char b)]
runCSP c l = flip runSupplyT l . flip runStateT M.empty . unCSP $ c
 
evalCSP :: CSP b a -> [b] -> [a]
evalCSP c l = flip runSupplyT l . flip evalStateT M.empty . unCSP $ c
 
execCSP :: CSP b a -> [b] -> [Map Char b]
execCSP c l = flip runSupplyT l . flip execStateT M.empty . unCSP $ c
 
byName :: Char -> CSP b b
byName k = gets (M.lookup k) >>= maybe c return
 where c = getSupply >>= assign k
 
assign :: Char -> a -> CSP a a
assign k a = modify (M.insert k a) >> return a
 
choose :: MonadPlus m => [a] -> m a
choose = msum . map return