llvm-hs Kaleidoscope Tutorial

module AST where

import Data.Char

import Text.Read 

import Text.ParserCombinators.ReadP hiding ((+++), choice)

data Expr = Num Float

          | Var String

          | BinOp BinOp Expr Expr

          | Call String [Expr]

  deriving Show

data BinOp = Add | Sub | Mul | Cmp Ordering

  deriving Show

instance Read Expr where

  readPrec = parens $ choice [ parseNum

                             , parseVar

                             , parseCall

                             , parseBinOp "<" 10 (Cmp LT)

                             , parseBinOp "+" 20 Add

                             , parseBinOp "-" 20 Sub

                             , parseBinOp "*" 40 Mul

                             ]

    where parseNum = Num <$> readPrec

          parseVar = Var <$> lift (munch1 isAlpha)

          parseBinOp s prc op = prec prc $ do

            a <- step readPrec

            lift $ do

              skipSpaces

              string s

              skipSpaces

            b <- readPrec

            return (BinOp op a b)

          parseCall = do

            func <- lift (munch1 isAlpha)

            params <- lift $ between (char '(') (char ')') $

                        sepBy (readS_to_P reads)

                              (skipSpaces >> char ',' >> skipSpaces)

            return (Call func params)

data Prototype = Prototype String [String]

  deriving Show

instance Read Prototype where

  readPrec = lift $ do

    name <- munch1 isAlpha

    params <- between (char '(') (char ')') $

                sepBy (munch1 isAlpha) skipSpaces

    return (Prototype name params)

data AST = Function Prototype Expr

         | Extern Prototype

         | TopLevelExpr Expr

  deriving Show

instance Read AST where

  readPrec = parseFunction +++ parseExtern +++ parseTopLevel

    where parseFunction = do

            lift $ string "def" >> skipSpaces

            Function <$> readPrec <*> readPrec

          parseExtern = do

            lift $ string "extern" >> skipSpaces

            Extern <$> readPrec

          parseTopLevel = TopLevelExpr <$> readPrec

import AST

import System.IO

import Text.Read

main = do

  hPutStr stderr "ready> "

  ast <- (readMaybe <$> getLine) :: IO (Maybe AST)

  case ast of

    Just x -> hPrint stderr x

    Nothing ->  hPutStrLn stderr "Couldn't parse"

  main

data Expr = Num Double

                             , parseBinOp ">" 10 (Cmp GT)

                             , parseBinOp "==" 10 (Cmp EQ)

{-# LANGUAGE OverloadedStrings #-}

import AST as K -- K for Kaleidoscope

import Utils

import Control.Monad.Trans.Reader

import Control.Monad.IO.Class

import Data.String

import qualified Data.Map as Map

import qualified Data.Text.Lazy.IO as Text

import LLVM.AST.AddrSpace

import LLVM.AST.Constant

import LLVM.AST.Float

import LLVM.AST.FloatingPointPredicate hiding (False, True)

import LLVM.AST.Operand

import LLVM.AST.Type as Type

import LLVM.IRBuilder

import LLVM.Pretty

import System.IO.Error

import Text.Read (readMaybe)

main :: IO ()

main = buildModuleT "main" repl >>= Text.hPutStrLn stderr . ("\n" <>) . ppll

repl :: ModuleBuilderT IO ()

repl = do

  liftIO $ hPutStr stderr "ready> "

  mline <- liftIO $ catchIOError (Just <$> getLine) eofHandler

  case mline of

    Nothing -> return ()

    Just l -> do

      case readMaybe l of

        Nothing ->  liftIO $ hPutStrLn stderr "Couldn't parse"

        Just ast -> do

          hoist $ buildAST ast

          mostRecentDef >>= liftIO . Text.hPutStrLn stderr . ppll

      repl

  where

    eofHandler e

      | isEOFError e = return Nothing

      | otherwise = ioError e

type Binds = Map.Map String Operand

buildAST :: AST -> ModuleBuilder Operand

buildAST (Function (Prototype nameStr paramStrs) body) = do

  let n = fromString nameStr

  function n params Type.double $ \ops -> do

    let binds = Map.fromList (zip paramStrs ops)

    flip runReaderT binds $ buildExpr body >>= ret

  where params = zip (repeat Type.double) (map fromString paramStrs)

buildAST (Extern (Prototype nameStr params)) =

  extern (fromString nameStr) (replicate (length params) Type.double) Type.double

buildAST (TopLevelExpr x) = function "__anon_expr" [] Type.double $

  const $ flip runReaderT mempty $ buildExpr x >>= ret

buildExpr :: Expr -> ReaderT Binds (IRBuilderT ModuleBuilder) Operand

buildExpr (Num x) = pure $ ConstantOperand (Float (Double x))

buildExpr (Var n) = do

  binds <- ask

  case binds Map.!? n of

    Just x -> pure x

    Nothing -> error $ "'" <> n <> "' doesn't exist in scope"

buildExpr (BinOp op a b) = do

  opA <- buildExpr a

  opB <- buildExpr b

  tmp <- instr opA opB

  if isCmp

    then uitofp tmp Type.double

    else return tmp

  where isCmp

          | Cmp _ <- op = True

          | otherwise = False

        instr = case op of

                  K.Add -> fadd

                  K.Sub -> fsub

                  K.Mul -> fmul

                  K.Cmp LT -> fcmp OLT

                  K.Cmp GT -> fcmp OGT

                  K.Cmp EQ -> fcmp OEQ

buildExpr (Call callee params) = do

  paramOps <- mapM buildExpr params

  let nam = fromString callee

      -- get a pointer to the function

      typ = FunctionType Type.double (replicate (length params) Type.double) False

      ptrTyp = Type.PointerType typ (AddrSpace 0)

      ref = GlobalReference ptrTyp nam

  call (ConstantOperand ref) (zip paramOps (repeat []))

{-|

Shoving away gross stuff into this one module.

-}

module Utils where

import Control.Monad.Trans.State

import Data.Functor.Identity

import LLVM.AST

import LLVM.IRBuilder.Module

import LLVM.IRBuilder.Internal.SnocList

mostRecentDef :: Monad m => ModuleBuilderT m Definition

mostRecentDef = last . getSnocList . builderDefs <$> liftModuleState get

hoist :: Monad m => ModuleBuilder a -> ModuleBuilderT m a

hoist m = ModuleBuilderT $ StateT $

  return . runIdentity . runStateT (unModuleBuilderT m)

import Control.Monad

import Control.Monad.Trans.Class

import Foreign.Ptr

import LLVM.Context

import LLVM.Module

import LLVM.OrcJIT

import LLVM.OrcJIT.CompileLayer

import LLVM.PassManager

import LLVM.Target

foreign import ccall "dynamic" mkFun :: FunPtr (IO Double) -> IO Double

data JITEnv = JITEnv

  { jitEnvContext :: Context

  , jitEnvCompileLayer :: IRCompileLayer ObjectLinkingLayer

  , jitEnvModuleKey :: ModuleKey

  }

main =

  withContext $ \ctx -> withHostTargetMachineDefault $ \tm ->

    withExecutionSession $ \exSession ->

      withSymbolResolver exSession (SymbolResolver symResolver) $ \symResolverPtr ->

        withObjectLinkingLayer exSession (const $ pure symResolverPtr) $ \linkingLayer ->

          withIRCompileLayer linkingLayer tm $ \compLayer ->

            withModuleKey exSession $ \mdlKey -> do

              let env = JITEnv ctx compLayer mdlKey

              _ast <- runReaderT (buildModuleT "main" repl) env

              return ()

-- This can eventually be used to resolve external functions, e.g. a stdlib call

symResolver :: MangledSymbol -> IO (Either JITSymbolError JITSymbol)

symResolver sym = undefined

repl :: ModuleBuilderT (ReaderT JITEnv IO) ()

          anon <- isAnonExpr <$> hoist (buildAST ast)

          def <- mostRecentDef

          llvmAst <- moduleSoFar "main"

          ctx <- lift $ asks jitEnvContext

          env <- lift ask

          liftIO $ withModuleFromAST ctx llvmAst $ \mdl -> do

            Text.hPutStrLn stderr $ ppll def

            let spec = defaultCuratedPassSetSpec { optLevel = Just 3 }

            -- this returns true if the module was modified

            withPassManager spec $ flip runPassManager mdl

            when anon (jit env mdl >>= hPrint stderr)

          when anon (removeDef def)

    isAnonExpr (ConstantOperand (GlobalReference _ "__anon_expr")) = True

    isAnonExpr _ = False

jit :: JITEnv -> Module -> IO Double

jit JITEnv{jitEnvCompileLayer=compLayer, jitEnvModuleKey=mdlKey} mdl =

  withModule compLayer mdlKey mdl $ do

    mangled <- mangleSymbol compLayer "__anon_expr"

    Right (JITSymbol fPtr _) <- findSymbolIn compLayer mdlKey mangled False

    mkFun (castPtrToFunPtr (wordPtrToPtr fPtr))

import Data.ByteString.Short (ShortByteString)

import Data.List

moduleSoFar :: MonadModuleBuilder m => ShortByteString -> m Module

moduleSoFar nm = do

  s <- liftModuleState get

  let ds = getSnocList (builderDefs s)

  return $ defaultModule { moduleName = nm, moduleDefinitions = ds }

removeDef :: MonadModuleBuilder m => Definition -> m ()

removeDef def = liftModuleState (modify update)

  where

    update (ModuleBuilderState defs typeDefs) =

      let newDefs = SnocList (delete def (getSnocList defs))

      in ModuleBuilderState newDefs typeDefs