ASTEvalBench.hs

{-# language Strict, LambdaCase, GADTs #-}
{-# options_ghc -Wincomplete-patterns -O2 #-}

import Gauge -- package gauge
import Control.Monad

data Op
  = Add
  | NotEq

data Exp
 = Lit Int
 | Var Int
 | Set Int Exp
 | Bin Op Exp Exp
 | Do Exp Exp
 | While Exp Exp

ast :: Exp
ast =
  -- x = 100
  (Do (Set 0 (Lit 100))
      -- i = 1000
      (Do (Set 1 (Lit 1000))
          -- for (; i != 0;)
          (Do (While (Bin NotEq (Lit 0)
                                (Var 1))
                     -- x = (((x + 4) + x) + 3)
                     (Do (Set 0 (Bin Add (Bin Add (Bin Add (Var 0)
                                                           (Lit 4))
                                                  (Var 0))
                                         (Lit 3)))
                         -- x = ((x + 2) + 4)
                         (Do (Set 0 (Bin Add (Bin Add (Var 0)
                                                      (Lit 2))
                                             (Lit 4)))
                             -- i = i + (-1)
                             (Set 1 (Bin Add (Lit (-1))
                                             (Var 1))))))
              -- return x
              (Var 0))))

-- list-based env is faster than Map for small envs
data Env = Nil | Cons Int Int Env

empty = Nil

insert :: Int -> Int -> Env -> Env
insert x n Nil = Cons x n Nil
insert x n (Cons x' n' e)
  | x == x'   = Cons x' n e
  | otherwise = Cons x' n' (insert x n e)

var :: Int -> Env -> Int
var x Nil = undefined
var x (Cons x' n e)
  | x == x'   = n
  | otherwise = var x e

-- -- using Data.Map.Strict, it's a bit slower
-- type Env = M.Map Int Int

-- insert :: Int -> Int -> Env -> Env
-- insert = M.insert

-- var :: Int -> Env -> Int
-- var x e = e M.! x

-- empty = M.empty

data Res = Res Int Env

--------------------------------------------------------------------------------

eval :: Exp -> Env -> Res
eval t e = case t of
  Lit n      -> Res n e
  Var x      -> Res (var x e) e
  Set x t    -> case eval t e of
                  Res n e -> Res n (insert x n e)
  Bin op t u -> case eval t e of
                  Res n1 e -> case eval u e of
                    Res n2 e -> case op of
                      Add   -> Res (n1 + n2) e
                      NotEq -> Res (fromEnum (n1 /= n2)) e
  Do t u    -> case eval t e of
                 Res n e -> eval u e
  While t u -> let loop t u e = case eval t e of
                     Res 1 e -> case eval u e of
                       Res n e -> loop t u e
                     Res n e -> Res n e
               in loop t u e

evalDirect :: Exp -> Int
evalDirect t = case eval t empty of
  Res n _ -> n

--------------------------------------------------------------------------------

bin :: Op -> Int -> Int -> Int
bin = \case
  Add -> (+)
  NotEq -> \v1 v2 -> if v1 /= v2 then 1 else 0
{-# inline bin #-}

evalCont :: Exp -> Int
evalCont e = loop e empty (\_ r -> r)
  where
  loop :: Exp -> Env -> (Env -> Int -> Int) -> Int
  loop exp env cont =
    case exp of
      Lit v -> cont env v
      Var n -> cont env $! var n env
      Set n exp -> loop exp env (\env v -> (cont $! insert n v env) v)
      Bin op e1 e2 -> loop e1 env (\env v1 ->
        loop e2 env (\env v2 ->
          cont env $! (bin op v1 v2)))
      Do first rest -> loop first env (\env _ -> loop rest env cont)
      While condition body ->
        let while cond body env cont = loop cond env $ \env v -> case v of
              1 -> loop body env $ \env _ -> while cond body env cont
              _ -> cont env v
        in while condition body env cont

--------------------------------------------------------------------------------

-- This gets compiled to (# Env -> Res #). The unary unboxed tuple
-- has no overhead, but it is necessary to prevent GHC from lifting out
-- the Env argument, and converting the definition to almost the same
-- as the direct evaluator.
data Closure = Cl {app :: Env -> Res}

goCl :: Exp -> Closure
goCl t = case t of
  Lit n      -> Cl $ Res n
  Var x      -> Cl $ \e -> Res (var x e) e
  Set x t    -> case goCl t of
                  t -> Cl $ \e -> case app t e of
                    Res n e -> Res n (insert x n e)
  Bin op t u -> case goCl t of
                 t -> case goCl u of
                   u -> case op of
                     Add   -> Cl $ \e -> case app t e of
                       Res n1 e -> case app u e of
                         Res n2 e -> Res (n1 + n2) e
                     NotEq -> Cl $ \e -> case app t e of
                       Res n1 e -> case app u e of
                         Res n2 e -> Res (fromEnum (n1 /= n2)) e
  Do t u    -> case goCl t of
                 t -> case goCl u of
                   u -> Cl $ \e -> case app t e of
                     Res n e -> app u e
  While t u -> case goCl t of
                 t -> case goCl u of
                   u -> Cl $ \e ->
                     let loop t u e = case app t e of
                           Res 1 e -> case app u e of
                             Res n e -> loop t u e
                           Res n e -> Res n e
                     in loop t u e

evalClosure :: Exp -> Int
evalClosure t = case app (goCl t) empty of Res n _ -> n

--------------------------------------------------------------------------------

data EvalExec a where
  EvalBind :: EvalExec a -> (a -> EvalExec b) -> EvalExec b
  EvalReturn :: a -> EvalExec a
  EvalLookup :: Int -> EvalExec Int
  EvalSet :: Int -> Int -> EvalExec ()

instance Functor EvalExec where
  fmap = liftM; {-# inline fmap #-}
instance Applicative EvalExec where
  pure = return; {-# inline pure #-}
  (<*>) = ap; {-# inline (<*>) #-}
instance Monad EvalExec where
  return = EvalReturn; {-# inline return #-}
  (>>=) = EvalBind; {-# inline (>>=) #-}

goM :: Exp -> EvalExec Int
goM = \case
  Lit n      -> pure n
  Var x      -> EvalLookup x
  Set x t    -> do n <- goM t
                   EvalSet x n
                   pure n
  Bin op t u -> do n1 <- goM t
                   n2 <- goM u
                   case op of
                     Add   -> pure $! n1 + n2
                     NotEq -> pure $! fromEnum (n1 /= n2)
  Do t u     -> goM t >> goM u
  While t u  -> let loop t u = do
                      n <- goM t
                      case n of
                        1 -> goM u >> loop t u
                        _ -> pure n
                in loop t u

data ResM a = ResM a Env

execM :: EvalExec a -> Env -> ResM a
execM m e = case m of
  EvalBind m k -> case execM m e of ResM a e -> (execM $! k a) e
  EvalReturn a -> ResM a e
  EvalLookup x -> ResM (var x e) e
  EvalSet x n  -> ResM () (insert x n e)

evalM :: Exp -> Int
evalM t = case execM (goM t) empty of
  ResM n _ -> n

--------------------------------------------------------------------------------

main :: IO ()
main = defaultMain [
     bench "evalDirect"  $ whnf evalDirect  ast
   , bench "evalClosure" $ whnf evalClosure ast
   , bench "evalCont"    $ whnf evalCont    ast
   , bench "evalM"       $ whnf evalM       ast
  ]