Variation and feedback on Michael Bernstein's Markov chain implementation

Chain.hs

module Chain where

import Control.Applicative
import Data.List (tails, unfoldr)
import System.Random (RandomGen, newStdGen, randomR)
import qualified Data.Map.Strict as M

import Data.Sequence as S (Seq, (><), (|>))
import qualified Data.Sequence as S

-- A map of prefixes to possible successors.
type PrefixMap k = M.Map (Seq k) (Seq k)

-- Build a prefix map of width n for the given sequence.
build :: Ord k => Int -> [k] -> PrefixMap k
build n ks = M.fromListWith (><) prefixes
  where prefixes = [ (S.fromList prefix, S.singleton k) | (prefix, k:_) <- splitAt n `map` tails ks ]

-- Generate a Markov chain from the given transition map and initial state.
-- The state should be a fixed-width prefix.
generate :: (Ord k, RandomGen g) => PrefixMap k -> Seq k -> g -> [k]
generate transitions initState rng = unfoldr step (initState, rng)
  where
    -- Choose a random successor k for prefix state.
    step (state, g) | Just candidates <- M.lookup state transitions,
                      Just (k, g') <- pick candidates g,
                      state' <- S.drop 1 state |> k     -- new prefix
                    = Just (k, (state', g'))            -- output k
                    -- If state has no successor, stop.
                    | otherwise = Nothing

-- Like generate, but use the global random generator in IO.
generateIO :: Ord k => PrefixMap k -> Seq k -> IO [k]
generateIO t i = generate t i <$> newStdGen

-- Helper: Pick a random element from a list.
pick :: RandomGen g => Seq k -> g -> Maybe (k, g)
pick xs _ | S.null xs = Nothing
pick xs g = Just (xs `S.index` i, g')
  where (i, g') = randomR (0, S.length xs - 1) g


main :: IO ()
main = do
    print' =<< unwords <$> demonstrate 1 (words "a man a plan a canal panama") (words "a")
    print' =<< demonstrate 1 " abracadabra " "a"
    print' =<< demonstrate 2 "--a--b--c--d--e++A++B++C++D++E--" "--"
  where
    print' = putStrLn . take 70
    -- Generate a Markov chain for the given prefix width, input, and seed.
    demonstrate n ks s = (s ++) <$> generateIO (build n ks) (S.fromList s)

This version generalizes the element type to any Ord k, simplifies build, and lets generate produce an infinite lazy output stream.

There are two variations:

1. List version
2. Seq version (same, except for using Data.Sequence for the transition map instead of lists)

Sample output:

a man a plan a man a man a plan a canal panama
adadabrabracabrabrabra ada abradabracada aca abrabrabrabrabradabracabr
--a--d--b--e++E--d--d--a--b--b--a--b--e++D++B++A++C++C++E--e++C++C++A+