Compiler for an SSA register machine

Compiler2.hs

{-# LANGUAGE NoMonomorphismRestriction #-}

import Control.Applicative
import Control.Monad.State
import Control.Monad.Writer
import Data.IntMap (IntMap)
import qualified Data.IntMap as IM
import Text.ParserCombinators.UU (pChainl)
import Text.ParserCombinators.UU.BasicInstances(pSym)
import Text.ParserCombinators.UU.Utils


-- | Huttons razor
data Exp = Add Exp Exp
         | Lit Int
         deriving (Show)

-- | denotational semantics
eval :: Exp -> Int
eval (Add x y) = eval x + eval y
eval (Lit i)   = i

-- | Parser
pExp = foldr pChainl pAtom [Add <$ pSym '+']
  where
    pAtom =  Lit <$> pNatural
         <|> pParens pExp

parse :: String -> Exp
parse = runParser "Expression" pExp


----------------------------------
-- | SSA virtual-machine compiler

type Reg = Int

data Opd = IConst Int
         | IReg   Reg
        deriving (Show)

data Op = IAdd Reg Opd Opd
        deriving (Show)

type Gen = WriterT [Op] (State Reg)

new :: Gen Reg
new = modify succ >> get

gen :: Exp -> Gen Opd
gen (Add x y) = do
  o1 <- gen x
  o2 <- gen y
  r  <- new
  tell [IAdd r o1 o2]
  return $ IReg r
gen (Lit i) = return $ IConst i

comp :: Exp -> [Op]
comp = flip evalState 0 . execWriterT . gen


----------------------------------
-- | a vm implementation

type VM = State (IntMap Int)

exec :: Op -> VM ()
exec (IAdd r o1 o2) = do
  v1 <- opd o1
  v2 <- opd o2
  store r $ v1 + v2

opd :: Opd -> VM Int
opd (IConst i) = return i
opd (IReg r)   = gets $ IM.findWithDefault 0 r

store :: Reg -> Int -> VM ()
store r v = modify (IM.insert r v)

run :: [Op] -> IntMap Int
run = flip execState IM.empty . mapM_ exec