Haskell Quiz/Probable Iterations/Solution Dolio

This quiz was pretty simple. The list monad makes generation of the test cases simple, and the writer monad is handy for capturing potential output for each line. I used a DList for the writer accumulator to avoid repeated copying.

module Main where

import Data.DList

import Control.Monad.Writer.Lazy

import System.Environment
import System.Exit
import Control.Monad

import Text.Printf

die = [1..6]

check :: ([Int] -> Bool) -> (Int, [Int]) -> Writer (DList String) Bool
check p (line, roll) = do tell $ if b then singleton hit else singleton noHit ; return b
 where
 b = p roll
 noHit = printf "%12d  %s" line (show roll)
 hit   = noHit ++ "  <=="

sample :: Int -> Int -> (Int, (Int, DList String))
sample i j = (length l, runWriter . liftM length . filterM (check p) $ zip [1..] l)
 where
 p l = length (filter (==5) l) >= j
 l = replicateM i die

chop :: [a] -> [a]
chop [] = []
chop (x:xs) = x : chop (drop 49999 xs)

main = do (v,s,i,j) <- processArgs
          let (total, (selected, out)) = sample i j
          if v
           then mapM_ putStrLn $ toList out
           else when s . mapM_ putStrLn . chop $ toList out
          putStrLn ""
          putStr "Number of desirable outcomes is "
          print selected
          putStr "Number of possible outcomes is "
          print total
          putStrLn ""
          putStr "Probability is "
          print $ fromIntegral selected / fromIntegral total

processArgs = do l <- getArgs
                 case l of
                      [i,j]        -> return (False, False, read i, read j)
                      ["-v", i, j] -> return (True, False, read i, read j)
                      ["-s", i, j] -> return (False, True, read i, read j)
                      _            -> do putStrLn "Unrecognized arguments."
                                         exitFailure