Haskell Quiz/Yahtzee/Solution Bobstopper

This is a bit longer than the ruby solution but is a bit more featureful.

import Control.Arrow (first, second, (&&&))
import Control.Monad
import Control.Monad.State

import Data.List
import Data.Maybe
import qualified Data.Map as Map

import System.Exit
import System.Random


-- Possible values of a 6 sided dice roll
data DiceRoll = One | Two | Three | Four | Five | Six
              deriving (Eq, Ord, Enum, Bounded, Show)

value = (1+) . fromEnum


-- Some Random Utilities

instance Random DiceRoll where
    randomR (min,max) g = first toEnum $ randomR (fromEnum min,fromEnum max) g
    random = randomR (minBound,maxBound)


randomIOs :: Random a => IO [a]
randomIOs = fmap randoms newStdGen


rollDice = randomIO :: IO DiceRoll

-- Generates n random dice
rollNDice :: Int -> IO [DiceRoll]
rollNDice n = (take n) `fmap` randomIOs


-- reRollDice ns dice returns a new set of dice with dices at positions
-- given by ns replaced with newly generated dice
reRollDice :: [Int] -> [DiceRoll] -> IO [DiceRoll]
reRollDice [] dice = return dice
reRollDice (n:ns) dice = do
  next <- fmap (replace dice n) $ rollDice
  reRollDice ns next

replace :: [a] -> Int -> a -> [a]
replace xs i x = (take i xs) ++ [x] ++ (drop (i+1) xs)


-- Scoring logic

-- Yahtzee categories which can be scored against
data Category = Ones | Twos | Threes | Fours | Fives | Sixes
              | ThreeOfAKind | FourOfAKind | FullHouse
              | SmallStraight | LargeStraight | Chance | Yahtzee
                deriving (Show, Read, Ord, Eq)


-- Small Mapping just for counting
type Count = Map.Map DiceRoll Int
initCount = Map.fromAscList [ (One, 0), (Two, 0), (Three, 0)
                            , (Four, 0), (Five, 0), (Six, 0)]


-- Count how many of each DiceRoll there is in a hand
countDice :: [DiceRoll] -> Count
countDice = foldr (Map.adjust (+1)) initCount


-- Count how many of a specific DiceRoll there is in a hand
numberMatching :: DiceRoll -> [DiceRoll] -> Int
numberMatching = (length .) . filter . (==)


-- Test if a list is ascending incrementally
isStraight :: (Bounded a, Enum a, Eq a) => [a] -> Bool
isStraight [] = True
isStraight xs = and $ zipWith isSucc xs (tail xs)
    where isSucc x y | x == maxBound = False
                     | otherwise     = y == succ x


-- Determine a hand's score against a specific category
score :: [DiceRoll] -> Category -> Int
score dice Ones = numberMatching One dice
score dice Twos = (numberMatching Two dice) * 2
score dice Threes = (numberMatching Three dice) * 3
score dice Fours = (numberMatching Four dice) * 4
score dice Fives = (numberMatching Five dice) * 5
score dice Sixes = (numberMatching Six dice) * 6
score dice ThreeOfAKind | Map.null $ Map.filter (>=3) $ countDice dice = 0
                        | otherwise = sum (map value dice)
score dice FourOfAKind | Map.null $ Map.filter (>=4) $ countDice dice = 0
                       | otherwise = sum (map value dice)
score dice Yahtzee | Map.null $ Map.filter (==5) $ countDice dice = 0
                   | otherwise = 50
score dice FullHouse | Map.size gt2 == 1 && head (Map.elems gt2) == 5 = 25
                     | Map.size gt2 == 2 
                       && (not $ Map.null $ Map.filter (==3) gt2)
                         = 25
                     | otherwise = 0
    where gt2 = Map.filter (>=2) $ countDice dice
score dice SmallStraight | (isStraight $ tail $ sort dice)
                           || (isStraight $ take 4 $ sort dice) = 30
                         | otherwise = 0
score dice LargeStraight | isStraight $ sort dice = 40
                         | otherwise = 0
score dice Chance = sum $ map value dice


-- The Game Logic

type PlayerName = String
type PlayerTotal = (PlayerName, Int)
type ScoreCard = Map.Map Category (Maybe Int)
maxReRolls = 2


-- An initial score with all categories unscored
clearScore = Map.fromAscList [ (Ones, Nothing), (Twos, Nothing)
                             , (Threes, Nothing), (Fours, Nothing)
                             , (Fives, Nothing), (Sixes, Nothing)
                             , (ThreeOfAKind, Nothing), (FourOfAKind, Nothing)
                             , (FullHouse, Nothing), (SmallStraight, Nothing)
                             , (LargeStraight, Nothing), (Chance, Nothing)
                             , (Yahtzee, Nothing) ]


data GameState = GameState { playerQueue :: [ PlayerName ]
                           , players :: Map.Map PlayerName ScoreCard
                           , reRolls :: Int }


-- Generate a new game using the player names given
newGame :: [PlayerName] -> GameState
newGame [] = error "Need to define at least one player"
newGame ps = GameState { playerQueue = cycle ps
                       , players = Map.fromList $ zip ps (repeat clearScore)
                       , reRolls = maxReRolls }


-- Change a specific player's score by scoring the given hand against
-- the given category
adjustScore :: PlayerName -> Category -> [DiceRoll] -> GameState -> GameState
adjustScore n c ds g = g { players = 
                           Map.adjust (adjustCategory c (score ds c))
                              n (players g) }


-- Adjust the score for a specific category in a scorecard
adjustCategory :: Category -> Int -> ScoreCard -> ScoreCard
adjustCategory Yahtzee x s = Map.adjust (maybe (Just 50) (Just . (+x)))
                             Yahtzee s
adjustCategory c x s = Map.adjust (maybe (Just x) 
                                   (error "Category already used"))
                       c s


currentPlayer :: GameState -> PlayerName
currentPlayer = head . playerQueue


-- Get the current scorecard of the given player
playerScore :: PlayerName -> GameState -> ScoreCard
playerScore n g = Map.findWithDefault err n (players g)
    where err = error $ "Couldn't find Player " ++ n ++ " in game!"


-- Get the current scorecard for the current player
currentPlayerScore :: GameState -> ScoreCard
currentPlayerScore g = playerScore (currentPlayer g) g


-- Uses up one of the available rerolls
useReRoll :: GameState -> GameState
useReRoll g = g { reRolls = (reRolls g)-1 }


-- Reset the game for the next player
newRound :: GameState -> GameState
newRound g = g { reRolls = maxReRolls
               , playerQueue = tail $ playerQueue g }


-- Returns the number of rounds in a game
gameLength :: GameState -> Int
gameLength = (13*) . Map.size . players


-- Look up the score of the given category from a scorecard
category :: Category -> ScoreCard -> Maybe Int
category c s = Map.findWithDefault err c s
    where err = error $ "Couldn't find Category " ++ (show c) ++ 
                " in player's score!"


-- Determines whether it is ok to place a score against the category
-- This function should always be used prior to scoring a category
categoryOK :: Category -> ScoreCard -> Bool
categoryOK Yahtzee _ = True
categoryOK c s = isNothing $ category c s


-- Calculates the "upper" total of the scorecard
scoreUpper :: ScoreCard -> Int
scoreUpper s = sum $ mapMaybe ((flip category) s) 
               [Ones,Twos,Threes,Fours,Fives,Sixes]


-- Calculates the "lower" total of the scorecard
scoreLower :: ScoreCard -> Int
scoreLower s = sum $ mapMaybe ((flip category) s) 
               [ ThreeOfAKind, FourOfAKind, FullHouse
               , SmallStraight, LargeStraight, Yahtzee, Chance ]


-- Calculates the grand total of the scorecard
scoreTotal :: ScoreCard -> Int
scoreTotal = (uncurry (+)) . (scoreUpper &&& scoreLower)


-- Data structure to account for draws as well as outright winners
data GameOutcome = Winner PlayerTotal | Draw [PlayerTotal]

-- Determines the winner between two players
winner :: PlayerTotal -> GameOutcome -> GameOutcome
winner y (Winner x) | (snd x) > (snd y)      = Winner x
                    | (snd x) < (snd y)      = Winner y
                    | otherwise              = Draw [x,y]
winner y (Draw xs@(x:_)) | (snd x) > (snd y) = Draw xs
                         | (snd x) < (snd y) = Winner y
                         | otherwise         = Draw $ y:xs


-- Score printing utilities

-- Prints the player's score
printScore :: PlayerName -> Yahtzee ()
printScore p = do
  score <- gets (playerScore p)
  liftIO $ putStrLn $ show score


-- Prints the player's final score totals
printFinalScore :: PlayerName -> ScoreCard -> IO ()
printFinalScore n s = do
  putStrLn $ "*** Final Score for " ++ n ++ " ***"
  putStrLn $ "Lower Score: " ++ (show $ scoreLower s)
  putStrLn $ "Upper Score: " ++ (show $ scoreUpper s)
  newline
  putStrLn $ "TOTAL Score: " ++ (show $ scoreTotal s)


-- Print the overall winner(s) and their score(s)
printWinner :: GameOutcome -> IO ()
printWinner (Winner (n,s)) = 
    putStrLn $ n++ " is the winner with " ++ (show s) ++ "!"
printWinner (Draw ws) = do 
  putStrLn "The game is a Draw!"
  mapM_ (\(n,s) -> putStrLn $ n++ "with " ++ (show s)) ws


-- The Interface

-- Valid user input types
data Input = Accept Category | ShowScore | ReRoll [Int] | ReRollAll 
           | Quit | Help | CategoryHelp
             deriving Read


-- The game monad
type Yahtzee a = StateT GameState IO a


newline = putStrLn ""

-- Persists in asking the given question until a valid answer is given.
-- Valid answers are determined and returned by the given function
ask :: String -> (String -> Maybe a) -> IO a
ask question f = do
  putStr (question ++ "> ") 
  x <- fmap f getLine
  case x of
    (Just answer) -> return answer
    Nothing -> putStrLn "Invalid input. Try again" >> ask question f


-- Attempts to read a string. Returns Nothing if it fails
maybeRead :: Read a => String -> Maybe a
maybeRead s = case r of
                [] -> Nothing
                ((x,_):_) -> Just x
    where r = reads s


-- The Game start and completion
main :: IO ()
main = do
  playerNames <- getPlayerNames
  let game = newGame playerNames
  gameOver <- execStateT (replicateM_ (gameLength game) playerUp) game
  let scores = Map.toList $ players gameOver
  mapM_ (uncurry printFinalScore) scores
  let topScore = foldr winner (Winner ("Nobody",0)) 
                 (map (second scoreTotal) scores)
  printWinner topScore


-- Ask for each player names until a blank line is reached
getPlayerNames :: IO [PlayerName]
getPlayerNames = do
  name <- ask "Enter player name (blank line to finish)" Just
  if null name 
     then return []
     else fmap (name:) getPlayerNames


-- Start a round for the next player
playerUp :: Yahtzee ()
playerUp = do
  player <- gets currentPlayer
  liftIO $ newline
  liftIO $ putStrLn "Next Round"
  liftIO $ newline
  printScore player
  liftIO $ newline
  liftIO $ putStrLn $ player ++ ", you're up!"
  (liftIO $ rollNDice 5) >>= userAction
  modify newRound
  

-- Interact with the user during the round
userAction :: [DiceRoll] -> Yahtzee ()
userAction dice = do
  liftIO $ putStrLn $ "Your roll is " ++ (show dice)
  input <- liftIO $ ask "Input Action (type Help for help)" maybeRead
  processInput input dice


-- Process a user's input
processInput :: Input -> [DiceRoll] -> Yahtzee ()
processInput (Accept c) ds = do
  player <- gets currentPlayer
  playerScore <- gets $ playerScore player
  if categoryOK c playerScore
     then do modify $ adjustScore player c ds
             printScore player
     else do (liftIO $ putStrLn ("Category " ++ (show c) ++ " already used"))
             userAction ds
processInput ShowScore ds = gets currentPlayer >>= printScore >> userAction ds
processInput (ReRoll is) ds = maybeReRoll ds $ 
                              modify useReRoll >> (liftIO $ reRollDice is ds)
processInput ReRollAll ds = maybeReRoll ds $
                            modify useReRoll >> (liftIO $ rollNDice 5)
processInput Quit _ = liftIO $ exitWith ExitSuccess
processInput Help ds = do
  liftIO $ putStrLn "* Command List *"
  liftIO $ putStrLn "Accept category: Accept this roll and score against category"
  liftIO $ putStrLn "Score: Print your current score"
  liftIO $ putStrLn "ReRoll indexes: reroll the dice specified by the Haskell list of Ints in indexes"
  liftIO $ putStrLn "ReRollAll: discard all dice and roll a fresh set"
  liftIO $ putStrLn "Quit: Exit the game"
  liftIO $ putStrLn "Help: Display this help"
  liftIO $ putStrLn "CategoryHelp: Display available categories"
  liftIO $ newline
  userAction ds
processInput CategoryHelp ds = do
  liftIO $ putStrLn "* Category List *"
  liftIO $ putStrLn "Ones | Twos | Threes | Fours | Fives | Sixes"
  liftIO $ putStrLn "| ThreeOfAKind | FourOfAKind | FullHouse"
  liftIO $ putStrLn "| SmallStraight | LargeStraight | Yahtzee | Chance"
  userAction ds


-- takes a current diceroll and a command to reroll.
-- Performs the reroll action only if the player has rerolls left.
-- Otherwise tells the player no rerolls are left
maybeReRoll :: [DiceRoll] -> Yahtzee [DiceRoll] -> Yahtzee ()
maybeReRoll dice cmd = do
  nextRound <- gets ((<=0) . reRolls)
  if nextRound
     then (liftIO $ putStrLn "No rolls left. Accept a category") >> 
          userAction dice
     else cmd >>= userAction