Main.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2
   3 import AST as K -- K for Kaleidoscope
   4 import Utils
   5 import Control.Monad.Trans.Reader
   6 import Control.Monad.IO.Class
   7 import Data.String
   8 import qualified Data.Map as Map
   9 import qualified Data.Text.Lazy.IO as Text
  10 import LLVM.AST.Constant
  11 import LLVM.AST.Float
  12 import LLVM.AST.FloatingPointPredicate hiding (False, True)
  13 import LLVM.AST.Operand
  14 import LLVM.AST.Type as Type
  15 import LLVM.IRBuilder
  16 import LLVM.Pretty
  17 import System.IO
  18 import System.IO.Error
  19 import Text.Read (readMaybe)
  20
  21 main :: IO ()
  22 main = buildModuleT "main" repl >>= Text.hPutStrLn stderr . ("\n" <>) . ppll
  23
  24 repl :: ModuleBuilderT IO ()
  25 repl = do
  26   liftIO $ hPutStr stderr "ready> "
  27   mline <- liftIO $ catchIOError (Just <$> getLine) eofHandler
  28   case mline of
  29     Nothing -> return ()
  30     Just l -> do
  31       case readMaybe l of
  32         Nothing ->  liftIO $ hPutStrLn stderr "Couldn't parse"
  33         Just ast -> do
  34           hoist $ buildAST ast
  35           mostRecentDef >>= liftIO . Text.hPutStrLn stderr . ppll
  36       repl
  37   where
  38     eofHandler e
  39       | isEOFError e = return Nothing
  40       | otherwise = ioError e
  41
  42 type Binds = Map.Map String Operand
  43
  44 buildAST :: AST -> ModuleBuilder Operand
  45 buildAST (Function (Prototype nameStr paramStrs) body) = do
  46   let n = fromString nameStr
  47   function n params Type.double $ \ops -> do
  48     let binds = Map.fromList (zip paramStrs ops)
  49     flip runReaderT binds $ buildExpr body >>= ret
  50   where params = zip (repeat Type.double) (map fromString paramStrs)
  51
  52 buildAST (Extern (Prototype nameStr params)) =
  53   extern (fromString nameStr) (replicate (length params) Type.double) Type.double
  54
  55 buildAST (TopLevelExpr x) = function "__anon_expr" [] Type.double $
  56   const $ flip runReaderT mempty $ buildExpr x >>= ret
  57
  58 buildExpr :: Expr -> ReaderT Binds (IRBuilderT ModuleBuilder) Operand
  59 buildExpr (Num x) = pure $ ConstantOperand (Float (Double x))
  60 buildExpr (Var n) = do
  61   binds <- ask
  62   case binds Map.!? n of
  63     Just x -> pure x
  64     Nothing -> error $ "'" <> n <> "' doesn't exist in scope"
  65
  66 buildExpr (BinOp op a b) = do
  67   opA <- buildExpr a
  68   opB <- buildExpr b
  69   tmp <- instr opA opB
  70   if isCmp
  71     then uitofp tmp Type.double
  72     else return tmp
  73   where isCmp
  74           | Cmp _ <- op = True
  75           | otherwise = False
  76         instr = case op of
  77                   K.Add -> fadd
  78                   K.Sub -> fsub
  79                   K.Mul -> fmul
  80                   K.Cmp LT -> fcmp OLT
  81                   K.Cmp GT -> fcmp OGT
  82                   K.Cmp EQ -> fcmp OEQ