a simple pratt parser in Haskell

pratt.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}

import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import           MonadLib

type Tables t a = (PrefixMap t a, InfixMap t a)

newtype PrattM t r a = PrattM
  { runPrattM :: ReaderT (Tables t r) (StateT [t] (ExceptionT String Id)) a
  } deriving (Functor, Applicative, Monad)

data PrefixParselet t a = PParselet
  { pParse :: PrattM t a a }

data InfixParselet t a = IParselet
  { iParse     :: a -> PrattM t a a
  , precedence :: Int
  }

type PrefixMap t a = Map t (PrefixParselet t a)
type InfixMap  t a = Map t (InfixParselet t a)

throw :: String -> PrattM t r a
throw = PrattM . raise

ifFailed :: PrattM t r a -> String -> PrattM t r a
ifFailed (PrattM mote) err = PrattM (mapException (const err) mote)

next :: PrattM t r t
next = do
  lst <- PrattM get
  case lst of
    (x:xs) -> PrattM (set xs) >> return x
    []     -> throw "no more tokens"

peek :: PrattM t r t
peek = do
  lst <- PrattM get
  case lst of
    (x:_) -> return x
    []    -> throw "nothing to peek"

match :: (t -> Bool) -> PrattM t r ()
match p = do
  tok <- next
  if p tok
    then return ()
    else throw "failed match"

getPrefixParser :: Ord t => t -> PrattM t r (PrefixParselet t r)
getPrefixParser t = do
  pMap <- PrattM (fmap fst ask)
  case M.lookup t pMap of
    Just p  -> return p
    Nothing -> throw "Unable to find appropriate prefix parser"

getInfixParser :: Ord t => t -> PrattM t r (Maybe (InfixParselet t r))
getInfixParser t = do
  iMap <- PrattM (fmap snd ask)
  return (M.lookup t iMap)

runParse :: Ord t => Tables t a -> [t] -> Either String a
runParse rules toks
  = runId
  . runExceptionT
  . fmap fst
  . runStateT toks
  . runReaderT rules
  . runPrattM
  $ parsePrefix 0

parsePrefix :: Ord t => Int -> PrattM t a a
parsePrefix n = do
  tok <- next
  pre <- getPrefixParser tok
  expr <- pParse pre
  parseInfix n expr

parseInfix :: Ord t => Int -> a -> PrattM t a a
parseInfix n left = do
  tok <- peek
  infMb <- getInfixParser tok
  case infMb of
    Just inf | n < precedence inf -> do
                 _ <- next
                 left' <- iParse inf left
                 parseInfix n left'
    _ -> return left

data AST
  = Num
  | Pos AST
  | Neg AST
  | Add AST AST
  | Mul AST AST
    deriving (Eq, Show)

samp :: [String] -> Either String AST
samp = runParse (pres, infs) . (++ ["EOF"])
  where singleton r = PParselet (return r)
        prefix f  = PParselet $ do
          p <- parsePrefix 100
          return (f p)
        infix_ i f = IParselet go i
          where go left = do
                  right <- parsePrefix i
                  return (f left right)
        parens = PParselet $ do
          p <- parsePrefix 0
          match (==")") `ifFailed` "Unclosed right-paren"
          return p
        pres = M.fromList [ ("0", singleton Num)
                          , ("+", prefix Pos)
                          , ("-", prefix Neg)
                          , ("(", parens)
                          ]
        infs = M.fromList [ ("+", infix_ 10 Add)
                          , ("*", infix_ 20 Mul)
                          ]

main :: IO ()
main = do
  str <- getContents
  print (samp (words str))