ぷよぷよ in Lens

Lenshttp://hackage.haskell.org/package/lensの練習にちょうどいいなあと思って。
30分はかかってない気がします。

import           Control.Lens
import           Control.Monad
import           Control.Monad.State
import           Data.List
import qualified Data.Map            as M

paint :: M.Map (Int, Int) Char -> (Int, Int) -> Int
paint mm pos = evalState (go pos) mm where
  go (i, j) = do
    b <- use $ at (i, j)
    if b == M.lookup pos mm then do
      at (i, j) ?= 'x'
      vs <- forM [(1, 0), (-1, 0), (0, 1), (0, -1)] $ \(di, dj) ->
        go (i + di, j + dj)
      return $ 1 + sum vs
      else do
      return 0

erase :: [String] -> [String]
erase bd =
  [ [ if paint mm (i, j) >= 4 then ' ' else c
    | (j, c) <- zip [0..] r
    ]
  | (i, r) <- zip [0..] bd
  ]
  where
    mm = M.fromList $ do
      (i, r) <- zip [0..] bd
      (j, c) <- zip [0..] r
      return $ ((i, j), c)

drops :: [String] -> [String]
drops = transpose . map (reverse . f . reverse) . transpose where
  f = take 13 . (++ repeat ' ') . filter (/= ' ')

next :: [String] -> [String]
next = drops . erase

main :: IO ()
main = do
  forM_ (take 20 $ iterate next input) $ \b -> do
    mapM_ putStrLn b
    putStrLn "======"

-----

input :: [String]
input =
  [ "  GYRR"
  , "RYYGYG"
  , "GYGYRR"
  , "RYGYRG"
  , "YGYRYG"
  , "GYRYRG"
  , "YGYRYR"
  , "YGYRYR"
  , "YRRGRG"
  , "RYGYGG"
  , "GRYGYR"
  , "GRYGYR"
  , "GRYGYR"
  ]