Boggle solver

boggle.hs

{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeOperators #-}


import Control.Comonad
import Control.Monad
import Data.Functor.Reverse
import Data.List (unfoldr)
import Data.Maybe (maybeToList)
import qualified Data.Map as M


-----------------------------------------------------------
-- Composing comonads
-----------------------------------------------------------

newtype (f :. g) a = Compose { getCompose :: f (g a) }
    deriving (Eq, Ord, Show, Read, Functor, Foldable, Traversable)

instance (Applicative f, Applicative g) => Applicative (f :. g) where
    pure = Compose . pure . pure
    (Compose fgf) <*> (Compose fgx) = Compose $ (<*>) <$> fgf <*> fgx

instance (Comonad f, Traversable f, Comonad g, Applicative g) => Comonad (f :. g) where
    extract = extract . extract . getCompose
    duplicate = Compose . fmap (fmap Compose . sequenceA) . duplicate . fmap duplicate . getCompose


-----------------------------------------------------------
-- Non-empty list zippers
-----------------------------------------------------------

data LZipper a = LZipper (Reverse [] a) a [a]
    deriving (Eq, Ord, Show, Read, Functor, Foldable, Traversable)

mkZipper :: a -> [a] -> LZipper a
mkZipper = LZipper (Reverse [])

repeatZ :: a -> LZipper a
repeatZ x = LZipper (Reverse $ repeat x) x (repeat x)

ints :: LZipper Integer
ints = LZipper (Reverse [-1, -2 ..]) 0 [1..]

fwd, bwd :: LZipper a -> Maybe (LZipper a)
fwd (LZipper _ _ []) = Nothing
fwd (LZipper (Reverse xs) e (y:ys)) = Just $ LZipper (Reverse (e:xs)) y ys
bwd (LZipper (Reverse []) _ _) = Nothing
bwd (LZipper (Reverse (x:xs)) e ys) = Just $ LZipper (Reverse xs) x (e:ys)


-- pointwise application. Truncates ragged lists
instance Applicative LZipper where
    pure x = LZipper (Reverse $ repeat x) x (repeat x)
    (LZipper (Reverse bf) f ff) <*> (LZipper (Reverse bx) x fx) =
        LZipper (Reverse $ zipWith ($) bf bx) (f x) (zipWith ($) ff fx)


instance Comonad LZipper where
    extract (LZipper _ x _) = x
    duplicate z = LZipper (Reverse $ unfoldr (step bwd) z) z (unfoldr (step fwd) z)
        where step move = fmap (\y -> (y, y)) . move


-----------------------------------------------------------
-- Two-dimensional grid zippers
-----------------------------------------------------------

type Grid = LZipper :. LZipper

mkGrid :: (a, [a]) -> [(a, [a])] -> Grid a
mkGrid (x, xs) xss = Compose $ mkZipper (mkZipper x xs) $ map (uncurry mkZipper) xss

up, down, left, right :: Grid a -> Maybe (Grid a)
up = fmap Compose . bwd . getCompose
down = fmap Compose . fwd . getCompose
left = fmap Compose . traverse bwd . getCompose
right = fmap Compose . traverse fwd . getCompose


coords :: Grid (Integer, Integer)
coords = Compose $ LZipper (Reverse $ [row x | x <- [-1, -2..]])
                           (row 0)
                           [row x | x <- [1..]]
    where row n = repeatZ (n,) <*> ints

withCoords :: Grid a -> Grid ((Integer, Integer), a)
withCoords g = (,) <$> coords <*> g

-----------------------------------------------------------
-- Tries
-----------------------------------------------------------

data Trie a = Trie { isEnd :: Bool, children :: M.Map a (Trie a) }

empty :: Trie a
empty = Trie { isEnd = False, children = M.empty }

singleton :: [a] -> Trie a
singleton = foldr (\x t -> Trie { isEnd = False, children = M.singleton x t }) (empty { isEnd = True })

union :: Ord a => Trie a -> Trie a -> Trie a
union t1 t2 = Trie {
    isEnd = isEnd t1 || isEnd t2,
    children = M.unionWith union (children t1) (children t2)
}

mkTrie :: Ord a => [[a]] -> Trie a
mkTrie = foldr union empty . map singleton

getChild :: Ord a => a -> Trie a -> Maybe (Trie a)
getChild x t = M.lookup x (children t)


-----------------------------------------------------------
-- Boggle solver
-----------------------------------------------------------

wordsFromFocus :: (Eq coord, Ord a) => Trie a -> Grid (coord, a) -> [[a]]
wordsFromFocus t g = wordsFromFocus' [] t g
    where wordsFromFocus' seen t g
            | fst (extract g) `elem` seen = []
            | otherwise = do
                let (coord, x) = extract g
                nextT <- maybeToList $ getChild x t
                let results = do
                    move <- [up, down, left, right, up >=> left, up >=> right, down >=> left, down >=> right]
                    nextG <- maybeToList (move g)
                    wordsFromFocus' (coord:seen) nextT nextG
                result <- if isEnd nextT
                          then [] : results
                          else results
                return (x : result)
                        

allWords :: Ord a => Trie a -> Grid a -> [[a]]
allWords t g = concat $ fmap (wordsFromFocus t) $ duplicate (withCoords g)


-----------------------------------------------------------
-- test
-----------------------------------------------------------

boggle = mkGrid ('a', "bc") [('d', "ef"), ('g', "hi")]
dict = mkTrie (wordsInGrid ++ wordsNotInGrid)
    where wordsInGrid = ["abc", "def", "cfi", "abe"]
          wordsNotInGrid = ["aba", "agh", "adi", "gha"]

main = let x = allWords dict boggle
       in print x >> print (x == ["abe","abc","cfi","def"])