1 {-# LANGUAGE OverloadedStrings #-}
3 import AST as K -- K for Kaleidoscope
6 import Control.Monad.Trans.Class
7 import Control.Monad.Trans.Reader
8 import Control.Monad.IO.Class
10 import qualified Data.Map as Map
11 import qualified Data.Text.Lazy.IO as Text
12 import LLVM.AST.AddrSpace
13 import LLVM.AST.Constant
15 import LLVM.AST.FloatingPointPredicate hiding (False, True)
16 import LLVM.AST.Operand
17 import LLVM.AST.Type as Type
21 import LLVM.PassManager
25 import System.IO.Error
26 import Text.Read (readMaybe)
30 withContext $ \ctx -> withHostTargetMachine $ \tm -> do
31 ast <- runReaderT (buildModuleT "main" repl) ctx
34 repl :: ModuleBuilderT (ReaderT Context IO) ()
36 liftIO $ hPutStr stderr "ready> "
37 mline <- liftIO $ catchIOError (Just <$> getLine) eofHandler
42 Nothing -> liftIO $ hPutStrLn stderr "Couldn't parse"
44 anon <- isAnonExpr <$> hoist (buildAST ast)
47 ast <- moduleSoFar "main"
49 liftIO $ withModuleFromAST ctx ast $ \mdl -> do
50 Text.hPutStrLn stderr $ ppll def
51 let spec = defaultCuratedPassSetSpec { optLevel = Just 3 }
52 -- this returns true if the module was modified
53 withPassManager spec $ flip runPassManager mdl
54 Text.hPutStrLn stderr . ("\n" <>) . ppllvm =<< moduleAST mdl
55 when anon (jit mdl >>= hPrint stderr)
57 when anon (removeDef def)
61 | isEOFError e = return Nothing
62 | otherwise = ioError e
63 isAnonExpr (ConstantOperand (GlobalReference _ "__anon_expr")) = True
66 jit :: Module -> IO Double
67 jit _mdl = putStrLn "Working on it!" >> return 0
69 type Binds = Map.Map String Operand
71 buildAST :: AST -> ModuleBuilder Operand
72 buildAST (Function (Prototype nameStr paramStrs) body) = do
73 let n = fromString nameStr
74 function n params Type.double $ \ops -> do
75 let binds = Map.fromList (zip paramStrs ops)
76 flip runReaderT binds $ buildExpr body >>= ret
77 where params = zip (repeat Type.double) (map fromString paramStrs)
79 buildAST (Extern (Prototype nameStr params)) =
80 extern (fromString nameStr) (replicate (length params) Type.double) Type.double
82 buildAST (TopLevelExpr x) = function "__anon_expr" [] Type.double $
83 const $ flip runReaderT mempty $ buildExpr x >>= ret
85 buildExpr :: Expr -> ReaderT Binds (IRBuilderT ModuleBuilder) Operand
86 buildExpr (Num x) = pure $ ConstantOperand (Float (Double x))
87 buildExpr (Var n) = do
89 case binds Map.!? n of
91 Nothing -> error $ "'" <> n <> "' doesn't exist in scope"
93 buildExpr (BinOp op a b) = do
98 then uitofp tmp Type.double
111 buildExpr (Call callee params) = do
112 paramOps <- mapM buildExpr params
113 let nam = fromString callee
114 -- get a pointer to the function
115 typ = FunctionType Type.double (replicate (length params) Type.double) False
116 ptrTyp = Type.PointerType typ (AddrSpace 0)
117 ref = GlobalReference ptrTyp nam
118 call (ConstantOperand ref) (zip paramOps (repeat []))