Haskell markov text generator

Question

I'm new to Haskell, and here is my first not-totally-trivial program. It's the first program I tend to write in any language -- a Markov text generator. I'm wondering what I can change to make it more idiomatic, or what language features I could make better use of.

import Data.List
import System.Environment
import qualified Data.Map as Map
import Control.Monad.State
import System.Random

type MarkovMap = Map.Map String String
type MarkovState = (MarkovMap, StdGen, String)

transition :: State MarkovState Char
transition = do 
    (m, gen, sofar) <- get
    let options = m Map.! sofar
        (index, newGen) = randomR (0, length options - 1) gen
        next = options !! index
    put (m, newGen, tail sofar ++ [next])
    return next

generateText :: State MarkovState Char -> State MarkovState String
generateText s = do
    x <- s
    xs <- generateText s
    return (x:xs)

getWords :: MarkovMap -> Int -> [String]
getWords m n =
    let keys        = filter ((==) ' ' . last) $ Map.keys m 
        (r, gen)    = randomR (0, length keys - 1) $ mkStdGen 137
        startState  = keys !! r 
        markovChars = evalState (generateText transition) (m, gen, startState)
    in  take n . words $ markovChars

printWords :: [String] -> IO ()
printWords ws = mapM_ putStrLn $ makeLines ws
    where makeLines [] = []
          makeLines ws = unwords (take 10 ws) : makeLines (drop 10 ws)

main :: IO ()
main = do
    (n:nwords:fileName:[]) <- getArgs
    contents <- readFile fileName
    let chainMap = chain (read n :: Int) . unwords . words $ contents
    printWords $ getWords chainMap (read nwords :: Int)

chain :: Int -> String -> Map.Map String String
chain n xs = 
    let from = map (take n) . tails $ xs ++ " " ++ xs
        to   = drop n . map (:[]) $ xs ++ " " ++ take n xs
    in Map.fromListWith (++) $ zip from to

Example usage:

Keep track of the last 3 characters, take 100 words from tobeornottobe.txt

> runhaskell markov.hs 3 100 tobeornottobe.txt

delay, the law's count with who would bear things all;
and that make and sweary life; fortal shocks the hue
of returns tural consience of somethis againsolution. To die, thers
the he law's the have, or nobler retus resolence to
troublesh is noblesh is sicklied of regards of some will,
and arrows of? To beary from who would by a
life; for inst give spurns, and, but to sleep: perchan
flesh is heir thing afterprises us for no mortal shocks
turns of action devoutly to dreathe sleep: perchance thance the
ills we hue of time, to suffled of great pith

Answer 1

First of all, it seems like you never change the MarkovMap in your States. So, why don't you take it out of MarkovState and change the type of, say, transition, to MarkovMap -> State MarkovState Char?

Try to use more combinators and less pattern-matching. For example, the generate function is something like sequence . repeat.

Answer 2

The code is already quite Haskell-ish. I might have used unfoldr on plain functions instead of 'sequence' on the state monad, but that's just a matter of taste. Nevertheless, your implementation could be more efficient, as it looks up an entry of the MarkovMap for every character that is emitted. The Haskell way would be to exploit laziness, so that at most one lookup per map entry is performed. Define the map as

type MarkovMap = Map.Map String (StdGen -> String)

and you can move the map lookup into the map entries by "tying the knot". It's a bit tricky if you try to re-implement this from scratch, but you can get there by small-step refactorings of your code.