content: {-# LANGUAGE RecursiveDo, FlexibleContexts #-}
  module Main where
  
  import Control.Applicative
  import Control.Exception
  import Control.DeepSeq
  import Criterion.Main
  import Data.Char
  import Data.Maybe
  import Text.Earley
  import qualified Text.Parsec as Parsec
  import qualified Text.Parsec.Pos as Parsec
  
  data Expr
    = Add Expr Expr
    | Mul Expr Expr
    | Var String
    deriving (Eq, Ord, Show)
  
  instance NFData Expr where
    rnf (Add a b) = rnf a `seq` rnf b
    rnf (Mul a b) = rnf a `seq` rnf b
    rnf (Var s)   = rnf s
  
  type Token = String
  
  tokenParens :: Bool -> [Token] -> [Token]
  tokenParens True s  = ["("] ++ s ++ [")"]
  tokenParens False s = s
  
  tokenExpr :: Int -> Expr -> [Token]
  tokenExpr _ (Var s) = [s]
  tokenExpr d (Add a b) = tokenParens (d > 0) $ tokenExpr 0 a ++ ["+"] ++ tokenExpr 1 b
  tokenExpr d (Mul a b) = tokenParens (d > 1) $ tokenExpr 1 a ++ ["*"] ++ tokenExpr 2 b
  
  linearSum :: Int -> Expr
  linearSum 1 = Var "x"
  linearSum n = Add (linearSum $ n - 1) (Var "x")
  
  treeSum :: Int -> Expr
  treeSum 1 = Var "x"
  treeSum n = let a = n `div` 2 -- will be at least 1
                  b = n - a
              in Add (treeSum a) (treeSum b)
  
  -- Earley parser
  
  expr :: Grammar r (Prod r String Token Expr)
  expr = mdo
    x1 <- rule $ Add <$> x1 <* namedToken "+" <*> x2
              <|> x2
              <?> "sum"
    x2 <- rule $ Mul <$> x2 <* namedToken "*" <*> x3
              <|> x3
              <?> "product"
    x3 <- rule $ Var <$> (satisfy isIdent <?> "identifier")
              <|> namedToken "(" *> x1 <* namedToken ")"
    return x1
  
  isIdent :: String -> Bool
  isIdent (x:_) = isAlpha x
  isIdent _     = False
  
  sepBy1 :: Prod r e t a -> Prod r e t op -> Grammar r (Prod r e t [a])
  sepBy1 p op = mdo
    ops <- rule $ pure [] <|> (:) <$ op <*> p <*> ops
    rule $ (:) <$> p <*> ops
  
  expr' :: Grammar r (Prod r String Token Expr)
  expr' = mdo
    let var = Var <$> satisfy isIdent <|> token "(" *> mul <* token ")"
    mul <- fmap (foldl1 Mul) <$> add `sepBy1` token "*"
    add <- fmap (foldl1 Add) <$> var `sepBy1` token "+"
    return mul
  
  parseEarley :: [Token] -> Maybe Expr
  parseEarley input = listToMaybe (fst (fullParses (parser expr) input))
  
  parseEarley' :: [Token] -> Maybe Expr
  parseEarley' input = listToMaybe (fst (fullParses (parser expr') input))
  
  -- Parsec parsec
  
  type Parsec = Parsec.Parsec [Token] ()
  
  parsecExpr :: Parsec Expr
  parsecExpr = mul
    where mul   = foldl1 Mul <$> add `Parsec.sepBy1` t "*"
          add   = foldl1 Add <$> var `Parsec.sepBy1` t "+"
          ident = Parsec.token id pos $ \y -> if isIdent y then Just (Var y) else Nothing
          var   = ident <|> (t "(" *> mul <* t ")")
          t x   = Parsec.token id pos $ \y -> if x == y then Just x else Nothing
          pos   = const (Parsec.initialPos "")
  
  parseParsec :: [Token] -> Maybe Expr
  parseParsec =  either (const Nothing) Just . Parsec.parse parsecExpr ""
  
  -- Our benchmark harness.
  
  linearInput :: Int -> (String, [Token])
  linearInput size = (show size, tokenExpr 0 $ linearSum size)
  
  treeInput :: Int -> (String, [Token])
  treeInput size = (show size, tokenExpr 0 $ treeSum size)
  
  inputBench :: (String, [Token]) -> Benchmark
  inputBench (name, input) = bench name $ nf id input
  
  earleyBench :: (String, [Token]) -> Benchmark
  earleyBench (name, input) = bench name $ nf parseEarley input
  
  earley'Bench :: (String, [Token]) -> Benchmark
  earley'Bench (name, input) = bench name $ nf parseEarley' input
  
  parsecBench :: (String, [Token]) -> Benchmark
  parsecBench (name, input) = bench name $ nf parseParsec input
  
  benchSizes :: [Int]
  benchSizes = [100, 200] -- [51, 101, 151, 201]
  
  linearInputs :: [(String, [Token])]
  linearInputs = map linearInput benchSizes
  
  treeInputs :: [(String, [Token])]
  treeInputs = map treeInput benchSizes
  
  main :: IO ()
  main = do
    evaluate (rnf linearInputs)
    evaluate (rnf treeInputs)
    defaultMain
      [ -- bgroup "inputs" $ map inputBench linearInputs 
        bgroup "earley" $ map earleyBench linearInputs
      , bgroup "earley'" $ map earley'Bench linearInputs
      , bgroup "parsec" $ map parsecBench linearInputs
      -- , bgroup "inputsTree" $ map inputBench treeInputs
      , bgroup "earleyTree" $ map earleyBench treeInputs
      , bgroup "earley'Tree" $ map earley'Bench treeInputs
      , bgroup "parsecTree" $ map parsecBench treeInputs
      ]
  
  