Word search solver

word-search.hs

import Data.Array
import Data.Maybe
import Data.List.Split
import System.Environment
import System.IO

direction "n" = \(r, c) -> (r - 1, c)
direction "s" = \(r, c) -> (r + 1, c)
direction "e" = \(r, c) -> (r, c + 1)
direction "w" = \(r, c) -> (r, c - 1)
direction (x:xs) = direction [x] . direction xs

findWords ws words directions = concatMap wordsAt (indices ws)
    where
      wordsAt idx             = concatMap (inDirectionsAt idx) words
      inDirectionsAt idx word = mapMaybe (try idx word) directions
      try idx word d          = if search word idx d then answer word idx d else Nothing
      answer word idx d       = Just (word, idx, end d idx (tail word))
      end d                   = foldl $ \a b -> d a
      search "" _ _           = True
      search (c:cs) idx d     = inRange (bounds ws) idx && c == ws ! idx && search cs (d idx) d

grid rows = array ((0,0), (max_r, max_c)) (zip indices characters)
    where
      indices    = [ (r, c) | r <- [0..max_r], c <- [0..max_c]]
      characters = concat rows
      max_r      = length rows - 1
      max_c      = length (head rows) - 1

directions = map direction . splitOn ","

readFiles = do
  [ grid_file, words_file, directions_file ] <- getArgs
  grid_text       <- readFile grid_file
  words_text      <- readFile words_file
  directions_text <- readFile directions_file
  return (grid (lines grid_text), lines words_text, directions (head (lines directions_text)))

showSolutions (grid, words, directions) = mapM print (findWords grid words directions)

main = readFiles >>= showSolutions