Wednesday, December 28, 2011

My first use of the state monad

While learning Haskell, I was looking for a concise implementation of a function which "reshapes" a list into a matrix.  Given the number of rows r, the number of columns c, and a list vs, the function should take r*c values from the list and create a r by c matrix out of them.  Here's the type:
toMatrix :: Int -> Int -> [a] -> [[a]]

First solution

First I wrote  a simpler function that would split a list into chunks of a given size, like this:
chunksOf :: Int -> [a] -> [[a]]
chunksOf _ [] = []
chunksOf c vs = h : (chunksOf c t)
    where (h, t) = splitAt c vs
Using the above, toMatrix could be implemented this way:
toMatrix r c = chunksOf c . take (r*c)
I had a feeling that a function like chunksOf should be already present somewhere in the standard library, so I asked Hoogle, but to no avail. There was chunksOf in Data.Text, but it operated on Text only (I retroactively named my function after the one in Data.Text).  However, Hoogle returned replicateM as well…

Second solution

… and I realized I could use it with the state monad to implement toMatrix.  The state could contain the list of values yet to be consumed, and the action to be replicated could be chopping off c values from the list:
splitOnce :: Int -> State [a] [a]
splitOnce c = do
    s <- get
    let (h, t) = splitAt c s
    put t
    return h
After a while I realized that the same function could be written in a much more concise form:
splitOnce' :: Int -> State [a] [a]
splitOnce' = state . splitAt
The solution was, therefore:
toMatrix r = evalState . replicateM r . state . splitAt

Summary

The second version is good enough for me and as a bonus it helped me understand the state monad. Note that the two implementations of toMatrix are not equivalent, as they handle lists shorter than r*c in different ways.  Future work: find a concise and preferably point-free implementation of chunksOf.

Update (2013-01-08)

This answer on StackOverflow contains a very nice implementation of chunksOf.