Functional programming contest

jars.hs

-- http://users.livejournal.com/_darkus_/641529.html
-- Usage: runghc jars.hs 5 8 2

import Control.Category
import Control.Applicative
import Control.Monad.Logic
import Control.Monad.Writer
import Control.Monad.Reader
import Control.Monad.State
import Data.List
import Data.Ord
import qualified Data.Map as Map
import System.Environment

data JarsState = JarsState !Int !Int
    deriving (Eq, Ord, Show)

type M a =
    StateT Int -- current step number
    (WriterT [String] -- descriptions of steps
    (LogicT -- for backtracking
    (StateT (Map.Map JarsState Int) -- seen states
    (Reader JarsState)))) -- capacities
    a

makeTransition :: String -> JarsState -> M JarsState
makeTransition desc js = do
    seen <- lift get
    currentStep <- get
    let seenStep = Map.lookup js seen
    if maybe False (currentStep >=) seenStep
        then mzero
        else
            do
            lift $ put $! Map.insert js currentStep seen
            put $! currentStep + 1
            tell [desc]
            return js

makeTransitions :: JarsState -> M JarsState
makeTransitions (JarsState n1 n2) = do
    JarsState c1 c2 <- ask

    msum
        [ makeTransition "Fill jar 1" $ JarsState c1 n2
        , makeTransition "Fill jar 2" $ JarsState n1 c2

        , makeTransition "Empty jar 1" $ JarsState 0 n2
        , makeTransition "Empty jar 2" $ JarsState n1 0

        , makeTransition "Pour from jar 1 to jar 2" $
            let q = min n1 (c2 - n2)
            in JarsState (n1 - q) (n2 + q)
        , makeTransition "Pour from jar 2 to jar 1" $
            let q = min n2 (c1 - n1)
            in JarsState (n1 + q) (n2 - q)
        ]

solve :: Int -> Int -> Int -> Maybe [String]
solve c1 c2 need =
    let (options, _) =
            flip runStateT 1 >>>
            runWriterT >>>
            observeAllT >>>
            flip runStateT Map.empty >>>
            flip runReader (JarsState c1 c2) $
            loop (JarsState 0 0) 
    in
        if null options
            then Nothing
            else Just $ snd $ minimumBy (comparing fst)
                [ (n, log) | ((_, n), log) <- options ]
    where
    loop js@(JarsState n1 n2) = msum
        [ makeTransitions js >>= loop
        , if n1 == need
            then tell ["Jar 1 now has the needed amount"]
            else mzero
        , if n2 == need
            then tell ["Jar 2 now has the needed amount"]
            else mzero
        ]

main = do
    [c1, c2, need] <- map read <$> getArgs
    maybe (putStrLn "No solution") (mapM_ putStrLn) $ solve c1 c2 need