Haskell solution to "12 balls and scale" problem

Balls.hs

import Control.Exception (assert)
import Control.Monad.Writer
import Data.List (intercalate, nub)
import Prelude hiding (Left, Right)

slice :: (Int, Int) -> [a] -> [a]
slice (l, r) x = take (r - l) $ drop l $ x

data Ball = Ball Int deriving (Eq, Show)
data Bias = Heavier | Lighter deriving (Eq, Show)
data Side = Left | Right deriving (Eq, Show)
data Result = LeftHeavier | RightHeavier | Balanced

is :: Side -> Bias -> Result -> Bool
(Left `is` Heavier) LeftHeavier = True
(Left `is` Heavier) RightHeavier = False
(Left `is` Lighter) LeftHeavier = False
(Left `is` Lighter) RightHeavier = True
(Right `is` Heavier) LeftHeavier = False
(Right `is` Heavier) RightHeavier = True
(Right `is` Lighter) LeftHeavier = True
(Right `is` Lighter) RightHeavier = False
(_ `is` _) Balanced = error "balanced"

sideOf :: Side -> Result -> Bias
Left `sideOf` LeftHeavier = Heavier
Left `sideOf` RightHeavier = Lighter
Right `sideOf` LeftHeavier = Lighter
Right `sideOf` RightHeavier = Heavier

whichSideIs :: Bias -> Result -> Side
whichSideIs Heavier LeftHeavier = Left
whichSideIs Heavier RightHeavier = Right
whichSideIs Lighter LeftHeavier = Right
whichSideIs Lighter RightHeavier = Left

actuallyWeigh :: (Ball, Bias) -> [Ball] -> [Ball] -> Result
actuallyWeigh (ball, bias) left right = let
	inLeft = ball `elem` left
	inRight = ball `elem` right
	in case (inLeft, inRight, bias) of
		(False, False, _) -> Balanced
		(False, True, Heavier) -> RightHeavier
		(False, True, Lighter) -> LeftHeavier
		(True, False, Heavier) -> LeftHeavier
		(True, False, Lighter) -> RightHeavier

class Monad w => WeighMonad w where
	weigh :: [Ball] -> [Ball] -> w Result
	illegal :: w a

solve12 :: WeighMonad w => [Ball] -> w (Ball, Bias)
solve12 b = do
	x <- weigh (slice (0,4) b) (slice (4,8) b)
	case x of
		Balanced -> solve4 (slice (8,12) b) (slice (0,8) b)
		_ -> solve4And4 (slice (0,4) b) x (slice (4,8) b) (slice (8,12) b)

solve4 :: WeighMonad w => [Ball] -> [Ball] -> w (Ball, Bias)
solve4 suspect neutral = do
	x <- weigh [suspect!!0, suspect!!1] [suspect!!2, head neutral]
	case x of
		Balanced -> solve1 (suspect!!3) neutral
		_ -> do
			y <- weigh [suspect!!0] [suspect!!1]
			case y of
				Balanced -> return (suspect!!2, Right `sideOf` x)
				_ -> case whichSideIs (Left `sideOf` x) y of
					Left -> return (suspect!!0, Left `sideOf` x)
					Right -> return (suspect!!1, Left `sideOf` x)

solve1 :: WeighMonad w => Ball -> [Ball] -> w (Ball, Bias)
solve1 ball neutral = do
	x <- weigh [ball] [head neutral]
	case x of
		Balanced -> illegal
		_ -> return (ball, Left `sideOf` x)

solve4And4 :: WeighMonad w => [Ball] -> Result -> [Ball] -> [Ball] -> w (Ball, Bias)
solve4And4 left x right neutral = do
	y <- weigh [left!!0, left!!1, right!!0, right!!1] [right!!2, neutral!!0, neutral!!1, neutral!!2]
	case y of
		Balanced -> do
			z <- weigh [left!!2] [left!!3]
			case z of
				Balanced -> return (right!!3, Right `sideOf` x)
				_ -> case whichSideIs (Left `sideOf` x) z of
					Left -> return (left!!2, Left `sideOf` x)
					Right -> return (left!!3, Left `sideOf` x)
		_ | (Left `is` (Left `sideOf` x)) y -> do
			z <- weigh [left!!0] [left!!1]
			case z of
				Balanced -> return (right!!2, Right `sideOf` x)
				_ -> case whichSideIs (Left `sideOf` x) z of
					Left -> return (left!!0, Left `sideOf` x)
					Right -> return (left!!1, Left `sideOf` x)
		_ | (Right `is` (Left `sideOf` x)) y -> do
			z <- weigh [right!!0] [right!!1]
			case z of
				Balanced -> illegal
				_ -> case whichSideIs (Left `sideOf` y) z of
					Left -> return (right!!0, Left `sideOf` y)
					Right -> return (right!!1, Left `sideOf` y)

balls = [Ball i | i <- [1..12]]

{- Mode 1: Print exhaustive list of what happens in case of each ball being
heavier or lighter -}

data WeighAndTrack a = WeighAndTrack { runWeighAndTrack :: (Ball, Bias) -> Writer [([Ball], [Ball])] a }

instance Monad WeighAndTrack where
	return x = WeighAndTrack (const (return x))
	a >>= b = WeighAndTrack $ \ secret -> do
		a' <- runWeighAndTrack a secret
		b' <- runWeighAndTrack (b a') secret
		return b'

instance WeighMonad WeighAndTrack where
	weigh left right =
		assert (length left == length right) $
		assert (nub (left ++ right) == left ++ right) $
		WeighAndTrack $ \ (ball, bias) -> do
			tell [(left, right)]
			return (actuallyWeigh (ball, bias) left right)
	illegal = error "impossible"

mode1 = forM_ [(ball, bias) | ball <- balls, bias <- [Lighter, Heavier]] $ \ (ball, bias) -> do
	putStrLn ("If " ++ formatBall ball ++ " is " ++ formatBias bias ++ ":")
	let ((ball', bias'), weighings) = runWriter (runWeighAndTrack (solve12 balls) (ball, bias))
	forM weighings $ \ (left, right) -> do
		let r = actuallyWeigh (ball, bias) left right
		putStrLn ("    Weigh " ++ formatBalls left ++ " against " ++ formatBalls right ++ " -> " ++ formatResult r)
	putStrLn ("    Answer: " ++ formatBall ball' ++ " is " ++ formatBias bias')
	putStrLn ("    " ++ if (ball, bias) == (ball', bias') then "Correct" else "Wrong")

{- Mode 2: Print decision tree -}

data DecisionTree a = Weighing [Ball] [Ball] (DecisionTree a) (DecisionTree a) (DecisionTree a) | Conclusion a | Impossible

instance Monad DecisionTree where
	return x = Conclusion x
	a >>= b = case a of
		Weighing left right ifLeftHeavier ifRightHeavier ifBalanced ->
			Weighing left right (ifLeftHeavier >>= b) (ifRightHeavier >>= b) (ifBalanced >>= b)
		Conclusion a -> b a

instance WeighMonad DecisionTree where
	weigh left right = Weighing left right (Conclusion LeftHeavier) (Conclusion RightHeavier) (Conclusion Balanced)
	illegal = Impossible

printDecisionTree :: Int -> DecisionTree (Ball, Bias) -> IO ()
printDecisionTree indent (Conclusion (ball, bias)) = do
	printIndent indent $ "conclude " ++ formatBall ball ++ " is " ++ formatBias bias
printDecisionTree indent (Weighing left right ifLeftHeavier ifRightHeavier ifBalanced) = do
	printIndent indent $ "weigh " ++ formatBalls left ++ " against " ++ formatBalls right
	printIndent indent $ "  if left heavier:"
	printDecisionTree (indent + 1) ifLeftHeavier
	printIndent indent $ "  if right heavier:"
	printDecisionTree (indent + 1) ifRightHeavier
	printIndent indent $ "  if balanced:"
	printDecisionTree (indent + 1) ifBalanced
printDecisionTree indent Impossible = do
	printIndent indent $ "impossible"

printIndent :: Int -> String -> IO ()
printIndent indent string = putStrLn (replicate (indent * 4) ' ' ++ string)

mode2 = printDecisionTree 0 (solve12 balls)

{- Mode 3: Interactively solve problem -}

instance WeighMonad IO where
	weigh left right = do
		putStrLn $ "Which is heavier of " ++ formatBalls left ++ " or " ++ formatBalls right ++ "?"
		let tryAgain = do
			putStr "Please enter 'left', 'right', or 'balanced':"
			answer <- getLine
			case answer of
				"left" -> return LeftHeavier
				"right" -> return RightHeavier
				"balanced" -> return Balanced
				_ -> do
					putStrLn "Huh?"
					tryAgain
		tryAgain
	illegal = do
		fail "Impossible."

mode3 = do
	maybeSolution <- catch (liftM Just $ solve12 balls) (const $ return Nothing)
	case maybeSolution of
		Just (ball, bias) -> putStrLn $ "Answer: " ++ formatBall ball ++ " is " ++ formatBias bias
		Nothing -> putStrLn "That's impossible."

main = do
	putStrLn "Modes:"
	putStrLn "1. Print how each possible heavier or lighter ball would be detected"
	putStrLn "2. Print decision tree"
	putStrLn "3. Interactively solve a problem"
	putStrLn "4. Quit"
	mode <- getLine
	case mode of
		"1" -> do
			mode1
			main
		"2" -> do
			mode2
			main
		"3" -> do
			mode3
			main
		"4" -> do
			return ()
		_ -> do
			putStrLn "I don't understand that mode."
			main

formatBall (Ball i) = show i
formatBias Heavier = "heavier"
formatBias Lighter = "lighter"
formatBalls bs = intercalate ", " (map formatBall bs)
formatResult LeftHeavier = "left is heavier"
formatResult RightHeavier = "right is heavier"
formatResult Balanced = "balanced"