Main.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 {-# LANGUAGE RecursiveDo #-}
   3
   4 module Main where
   5
   6 import qualified AST
   7 import Control.Monad
   8 import Control.Monad.Trans.Class
   9 import qualified Data.Map as Map
  10 import qualified Data.Text.Lazy.IO as Text
  11 import Data.String
  12 import Foreign.Ptr
  13 import System.Exit
  14 import System.IO
  15 import LLVM.Context
  16 import LLVM.ExecutionEngine
  17 import LLVM.Module
  18 import LLVM.PassManager
  19 import LLVM.IRBuilder
  20 import LLVM.AST.AddrSpace
  21 import LLVM.AST.Constant
  22 import LLVM.AST.Float
  23 import LLVM.AST.FloatingPointPredicate hiding (False, True)
  24 import LLVM.AST.Operand
  25 import LLVM.AST.Type as Type
  26 import LLVM.AST.Typed
  27 import LLVM.Pretty
  28
  29 type ModuleBuilderE = ModuleBuilderT (Either String)
  30
  31 foreign import ccall "dynamic" exprFun :: FunPtr (IO Float) -> IO Float
  32
  33 main :: IO ()
  34 main = do
  35   AST.Program asts <- read <$> getContents
  36   let eitherMdl = buildModuleT "main" $ mapM buildAST asts
  37   case eitherMdl of
  38     Left err -> die err
  39     Right mdl -> withContext $ \ctx -> do
  40       hPutStrLn stderr "Before optimisation:"
  41       Text.hPutStrLn stderr (ppllvm mdl)
  42       withMCJIT ctx Nothing Nothing Nothing Nothing $ \mcjit ->
  43         withModuleFromAST ctx mdl $ \mdl' ->
  44           withPassManager defaultCuratedPassSetSpec $ \pm -> do
  45             runPassManager pm mdl' >>= guard
  46             hPutStrLn stderr "After optimisation:"
  47             Text.hPutStrLn stderr . ppllvm =<< moduleAST mdl'
  48             withModuleInEngine mcjit mdl' $ \emdl -> do
  49               Just f <- getFunction emdl "expr"
  50               let f' = castFunPtr f :: FunPtr (IO Float)
  51               exprFun f' >>= print
  52
  53 evalProg :: AST.Program -> IO (Maybe Float)
  54 evalProg (AST.Program asts) = do
  55   let eitherMdl = buildModuleT "main" $ mapM buildAST asts
  56   case eitherMdl of
  57     Left _ -> return Nothing
  58     Right mdl -> withContext $ \ctx ->
  59       withMCJIT ctx Nothing Nothing Nothing Nothing $ \mcjit ->
  60         withModuleFromAST ctx mdl $ \mdl' ->
  61           withModuleInEngine mcjit mdl' $ \emdl -> do
  62             Just f <- getFunction emdl "expr"
  63             let f' = castFunPtr f :: FunPtr (IO Float)
  64             Just <$> exprFun f'
  65
  66 -- | Builds up programs at the top-level of an LLVM Module
  67 -- >>> evalProg (read "31 - 5")
  68 -- Just 26.0
  69 buildAST :: AST.AST -> ModuleBuilderE Operand
  70 buildAST (AST.Function nameStr paramStrs body) = do
  71   let n = fromString nameStr
  72   function n params float $ \binds -> do
  73     let bindMap = Map.fromList (zip paramStrs binds)
  74     buildExpr bindMap body >>= ret
  75   where params = zip (repeat float) (map fromString paramStrs)
  76 buildAST (AST.Eval e) =
  77   function "expr" [] float $ \_ -> buildExpr mempty e >>= ret
  78
  79 -- | Builds up expressions, which are operands in LLVM IR
  80 -- >>> evalProg (read "def foo(x) x * 2; foo(6)")
  81 -- Just 12.0
  82 -- >>> evalProg (read "if 3 > 2 then 42 else 12")
  83 -- Just 42.0
  84 buildExpr :: Map.Map String Operand -> AST.Expr -> IRBuilderT ModuleBuilderE Operand
  85 buildExpr _ (AST.Num a) = pure $ ConstantOperand (Float (Single a))
  86 buildExpr binds (AST.Var n) = case binds Map.!? n of
  87   Just x -> pure x
  88   Nothing -> lift $ lift $ Left $ "'" <> n <> "' doesn't exist in scope"
  89
  90 buildExpr binds (AST.Call nameStr params) = do
  91   paramOps <- mapM (buildExpr binds) params
  92   let name = fromString nameStr
  93       -- get a pointer to the function
  94       typ = FunctionType float (replicate (length params) float) False
  95       ptrTyp = Type.PointerType typ (AddrSpace 0)
  96       ref = GlobalReference ptrTyp name
  97   call (ConstantOperand ref) (zip paramOps (repeat []))
  98
  99 buildExpr binds (AST.BinOp op a b) = do
 100   va <- buildExpr binds a
 101   vb <- buildExpr binds b
 102   let instr = case op of
 103                 AST.Add -> fadd
 104                 AST.Sub -> fsub
 105                 AST.Mul -> fmul
 106                 AST.Cmp GT -> fcmp OGT
 107                 AST.Cmp LT -> fcmp OLT
 108                 AST.Cmp EQ -> fcmp OEQ
 109   instr va vb
 110
 111 buildExpr binds (AST.If cond thenE elseE) = mdo
 112   _ifB <- block `named` "if"
 113
 114   condV <- buildExpr binds cond
 115   when (typeOf condV /= i1) $ lift $ lift $ Left "Not a boolean"
 116   condBr condV thenB elseB
 117
 118   thenB <- block `named` "then"
 119   thenOp <- buildExpr binds thenE
 120   br mergeB
 121
 122   elseB <- block `named` "else"
 123   elseOp <- buildExpr binds elseE
 124   br mergeB
 125
 126   mergeB <- block `named` "ifcont"
 127   phi [(thenOp, thenB), (elseOp, elseB)]