Find our JIT'ed function and run it
[kaleidoscope-hs.git] / 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