Tagless final games

tagless.hs

-- No LANGUAGE pragmas! This is all barebones Haskell.
-- $ ghc --make -O2 -fllvm -fforce-recomp tagless.hs

import Control.Applicative
import Control.Monad.Reader hiding (fix)
import Control.Monad.State hiding (fix)

import Criterion.Main

-- An example of how multiple backends can be achieved for an EDSL
-- through the final tagless style, including a backend that amounts
-- to a direct Haskell implementation. 


-- | Oleg calls this the 'Symantics' class, as the class definition
-- comprises the syntax, and instances of the class comprise various
-- semantics.
class Exp repr where
    int :: Integer -> repr Integer
    bool :: Bool -> repr Bool
    lam :: (repr a -> repr b) -> repr (a -> b)
    app :: repr (a -> b) -> repr a -> repr b
    fix :: (repr a -> repr a) -> repr a
    add :: repr Integer -> repr Integer -> repr Integer
    sub :: repr Integer -> repr Integer -> repr Integer
    mul :: repr Integer -> repr Integer -> repr Integer
    leq :: repr Integer -> repr Integer -> repr Bool
    if_ :: repr Bool -> repr a -> repr a -> repr a

-- | Factorial function in our embedded language, polymorphic in its
-- backend.
fact :: (Exp repr) => repr (Integer -> Integer)
fact = fix $ \f ->
         lam $ \n ->
             if_ (leq n (int 0))
                 (int 1)
                 (mul n (app f (sub n (int 1))))

-- | For comparison, the straightforward fact using the same operators
-- as the embedded version.
fact' :: Integer -> Integer
fact' n = if n <= 0 then 1 else n * (fact' (n-1))

--------------------------------------------------------------------------------
-- A first-order representation of the abstract syntax, suitable for
-- running through a compiler.

-- | Since our embedded syntax uses HOAS, we don't need to have pretty
-- variable names.
type Var = Int

-- | Mostly-standard first-order representation, although Fix might be
-- a bit strange, I'm not sure.
data ExpE = IntE Integer
          | BoolE Bool
          | VarE Var
          | LamE Var ExpE
          | AppE ExpE ExpE
          | FixE Var ExpE
          | AddE ExpE ExpE
          | SubE ExpE ExpE
          | MulE ExpE ExpE
          | LeqE ExpE ExpE
          | IfE ExpE ExpE ExpE
            deriving (Eq, Ord, Show)

-- | Runtime representation of values, still first-order.
data ValV = IntV Integer
          | BoolV Bool
          | ClosV Var ExpE Env
            deriving (Eq, Ord, Show)

-- | Association list runtime environment.
type Env = [(Var, ValV)]

-- | Evaluator in the 'Reader' monad. Might actually be less clean
-- than the standard 311 definition. Note that since '((->) env)' is
-- already a monad, we don't need 'Reader', but it has nice
-- combinators like 'ask' and 'local'.
evalExpE :: ExpE -> Reader Env ValV
evalExpE (IntE n)  = return $ IntV n
evalExpE (BoolE b) = return $ BoolV b
evalExpE (VarE x)  = do
  env <- ask
  case lookup x env of
    Just val -> return val
    Nothing -> error ("unbound variable " ++ show x)
evalExpE (LamE x body) = ClosV x body <$> ask
evalExpE (AppE rator rand) = do
  p <- evalExpE rator
  a <- evalExpE rand
  case p of
    ClosV x body env -> do 
      local (const ((x,a):env)) (evalExpE body)
    _ -> error ("attempt to apply non-procedure " ++ show p)
evalExpE (FixE x body) = do
  -- this is tricky, and relies on Haskell's laziness to tie the
  -- knot. If you're playing around at the REPL and try to print an
  -- env with a fixpoint in it, you'll be printing for a very long
  -- time...
  env <- ask
  let rclos = runReader (evalExpE (FixE x body)) ((x,rclos):env)
  local ((x,rclos):) (evalExpE body)
evalExpE (AddE x y) = do
  v1 <- evalExpE x
  v2 <- evalExpE y
  case (v1, v2) of
    (IntV x', IntV y') -> return $ IntV (x' + y')
    _ -> error "non-int arguments to add"
evalExpE (SubE x y) = do
  v1 <- evalExpE x
  v2 <- evalExpE y
  case (v1, v2) of
    (IntV x', IntV y') -> return $ IntV (x' - y')
    _ -> error "non-int arguments to sub"
evalExpE (MulE x y) = do
  v1 <- evalExpE x
  v2 <- evalExpE y
  case (v1, v2) of
    (IntV x', IntV y') -> return $ IntV (x' * y')
    _ -> error "non-int arguments to mul"
evalExpE (LeqE x y) = do
  v1 <- evalExpE x
  v2 <- evalExpE y
  case (v1, v2) of
    (IntV x', IntV y') -> return $ BoolV (x' <= y')
    _ -> error "non-int arguments to leq"
evalExpE (IfE t c a) = do
  b <- evalExpE t
  case b of
    BoolV True  -> evalExpE c
    BoolV False -> evalExpE a
    _ -> error "non-boolean test in if"

-- | The tricky part is how we interpret the HOAS of our embedded
-- language into the first-order 'ExpE' type. This approach
-- instantiates 'Exp' with a state monad containing a counter where we
-- can get fresh variable names. The phantom type is necessary to
-- match the kind of 'repr', and ensures type safety of constructed
-- terms.
newtype ExpS a = ES { unES :: State Var ExpE }

instance Exp ExpS where
  int       = ES . return . IntE
  bool      = ES . return . BoolE
  lam f     = ES $ do 
    -- 'lam' and 'fix' are the interesting cases where we need to go
    -- from HOAS to first-order. We do this by creating a fresh
    -- variable, wrapping it in the 'ExpS' State monad computation,
    -- passing that to the function of '(repr a -> repr b)', and
    -- running the resulting computation. Then, we record what
    -- variable we used, and save the resulting computation.
    fresh <- get
    modify (1+)
    body <- unES (f (ES (return (VarE fresh))))
    return (LamE fresh body)
  app f x   = ES $ AppE <$> unES f <*> unES x
  fix f     = ES $ do 
    -- see 'lam' case
    fresh <- get
    modify (1+)
    body <- unES (f (ES (return (VarE fresh))))
    return (FixE fresh body)
  add x y   = ES $ AddE <$> unES x <*> unES y
  sub x y   = ES $ SubE <$> unES x <*> unES y
  mul x y   = ES $ MulE <$> unES x <*> unES y
  leq x y   = ES $ LeqE <$> unES x <*> unES y
  if_ t c a = ES $ IfE  <$> unES t <*> unES c <*> unES a

-- | Build the first-order representation by starting the fresh
-- variable counter at 0 and running the 'State' computation.
buildExpE :: ExpS a -> ExpE
buildExpE c = evalState (unES c) 0

-- | Tie it all together: convert HOAS -> first-order, and then
-- interpret in the empty environment.
valofE :: ExpS a -> ValV
valofE e = runReader (evalExpE (buildExpE e)) []

--------------------------------------------------------------------------------
-- Interpretation as ordinary Haskell code

-- | 'newtype' introduces an isomorphism rather than a new algebraic
-- type. For the purposes of typechecking, it is the same as using
-- 'data', but there is no runtime cost of boxing/unboxing.
newtype Hask a = H { unH :: a }

-- | Defining the 'Exp' instance is just a matter of boxing and
-- unboxing when needed. Remember that there is no cost to these at
-- runtime, so if you look past the calls to 'H' and 'unH', you'll see
-- that this really does just use Haskell's primitives.
instance Exp Hask where
  int       = H
  bool      = H
  lam f     = H $ \x -> unH (f (H x))
  app f x   = H $ (unH f) (unH x)
  fix f     = f (fix f)
  add x y   = H $ unH x + unH y
  sub x y   = H $ unH x - unH y
  mul x y   = H $ unH x * unH y
  leq x y   = H $ unH x <= unH y
  if_ t c a = H $ if (unH t) then (unH c) else (unH a)

-- | Evaluating a 'Hask' interpretation is trivial: just unbox it.
valofH :: Hask a -> a
valofH = unH

--------------------------------------------------------------------------------
-- Fat fibs for benchmarking

{-

-- SPECIALIZE pragmas don't seem to make much difference

{-# SPECIALIZE int :: Integer -> Hask Integer #-}
{-# SPECIALIZE bool :: Bool -> Hask Bool #-}
{-# SPECIALIZE lam :: (Hask a -> Hask b) -> Hask (a -> b) #-}
{-# SPECIALIZE app :: Hask (a -> b) -> Hask a -> Hask b #-}
{-# SPECIALIZE fix :: (Hask a -> Hask a) -> Hask a #-}
{-# SPECIALIZE add :: Hask Integer -> Hask Integer -> Hask Integer #-}
{-# SPECIALIZE sub :: Hask Integer -> Hask Integer -> Hask Integer #-}
{-# SPECIALIZE mul :: Hask Integer -> Hask Integer -> Hask Integer #-}
{-# SPECIALIZE leq :: Hask Integer -> Hask Integer -> Hask Bool #-}
{-# SPECIALIZE if_ :: Hask Bool -> Hask a -> Hask a -> Hask a #-}

{-# SPECIALIZE int :: Integer -> ExpS Integer #-}
{-# SPECIALIZE bool :: Bool -> ExpS Bool #-}
{-# SPECIALIZE lam :: (ExpS a -> ExpS b) -> ExpS (a -> b) #-}
{-# SPECIALIZE app :: ExpS (a -> b) -> ExpS a -> ExpS b #-}
{-# SPECIALIZE fix :: (ExpS a -> ExpS a) -> ExpS a #-}
{-# SPECIALIZE add :: ExpS Integer -> ExpS Integer -> ExpS Integer #-}
{-# SPECIALIZE sub :: ExpS Integer -> ExpS Integer -> ExpS Integer #-}
{-# SPECIALIZE mul :: ExpS Integer -> ExpS Integer -> ExpS Integer #-}
{-# SPECIALIZE leq :: ExpS Integer -> ExpS Integer -> ExpS Bool #-}
{-# SPECIALIZE if_ :: ExpS Bool -> ExpS a -> ExpS a -> ExpS a #-}

{-# SPECIALIZE fib :: Integer -> ExpS Integer #-}
{-# SPECIALIZE fib :: Integer -> Hask Integer #-}

-}

fib :: (Exp repr) => Integer -> repr Integer
fib n = app (fix $ \f ->
                 lam $ \n ->
                     if_ (leq n (int 0))
                         (int 0)
                         (if_ (leq n (int 1))
                              (int 1)
                              (add (app f (sub n (int 1)))
                                   (app f (sub n (int 2))))))
            (int n)

fib' :: Integer -> Integer
fib' n = let func f n = if n <= 0
                        then 0
                        else if n <= 1
                             then 1
                             else (f (n-1)) + (f (n-2))
             fix' f = f (fix' f)
         in (fix' func) n

main = let n = 30 in
       defaultMain [
           bgroup "fib" [ bench "ExpS" $ whnf (valofE . fib) n
                        , bench "Hask" $ whnf (valofH . fib) n
                        , bench "native" $ whnf fib' n
                        ]
       ]