kaleidoscope-plus.hs

{-# LANGUAGE DeriveDataTypeable #-}

module Main where

import System.Console.CmdArgs.Implicit
import System.IO
import System.Environment
import System.Exit

import Data.Maybe
import Control.Monad
import Foreign.Marshal.Utils

import qualified Data.Text as T

import LLVM.Core
import LLVM.ExecutionEngine

import Data.Functor.Identity
import Text.Parsec.Combinator
import Text.Parsec.String hiding (Parser)
import Text.Parsec.Char hiding (upper)
import Text.Parsec.Prim
import qualified Text.Parsec.Token as PT
import Text.Parsec.Language (emptyDef)

import qualified Data.HashMap.Strict as H

type Parser = Parsec String (H.HashMap Char Int)

data KSToken =  Def | Extern 
             | Ident String | Number Double | Keyword Char 
             | If | Then | Else | For | In 
             | Binary | Unary

data KSExpr = -- | Terminals
              KSNumber Double
            | KSVariable String
            | KSUnaryOp Char KSExpr
            | KSBinaryOp Char KSExpr KSExpr
            | KSCall String [KSExpr]
            | KSIf KSExpr KSExpr KSExpr
            | KSFor String KSExpr KSExpr (Maybe KSExpr) KSExpr
            deriving (Show, Eq)

data KSProto =   KSProto String [String]
                | KSBinOpProto String [String] Integer
                deriving (Show, Eq)


data KSFunc =  KSFunc KSProto KSExpr
               deriving (Show, Eq)

data KSFuncOrExprList = KSFuncNext KSFunc KSFuncOrExprList
                      | KSEmpty
                      deriving (Show, Eq)

data KSProgram = KSProgram KSFuncOrExprList
                 deriving (Show, Eq)

configDef:: PT.GenLanguageDef String st Identity
configDef = emptyDef
    { PT.commentStart = ""
    , PT.commentEnd = ""
    , PT.commentLine = "#"
    , PT.nestedComments = False
    , PT.identStart = oneOf $ ['0'..'9'] ++ ['a'..'z'] ++ ['A'..'Z']
    , PT.identLetter = oneOf $ ['0'..'9'] ++ ['a'..'z'] ++ ['A'..'Z']
    , PT.reservedOpNames = []
    , PT.reservedNames = ["def", "extern", "if", "then", "else", "for", "in", "binary", "unary"]
    , PT.caseSensitive = True
    }


lexer :: PT.GenTokenParser String (H.HashMap Char Int) Identity
lexer = PT.makeTokenParser configDef


p_identifier :: ParsecT String (H.HashMap Char Int) Identity String
p_identifier = PT.identifier lexer

p_stringLiteral :: ParsecT String (H.HashMap Char Int) Identity String
p_stringLiteral = PT.stringLiteral lexer

p_whiteSpace :: ParsecT String (H.HashMap Char Int) Identity ()
p_whiteSpace = PT.whiteSpace lexer

p_parens :: ParsecT String (H.HashMap Char Int) Identity a -> ParsecT String (H.HashMap Char Int) Identity a
p_parens = PT.parens lexer

p_braces :: ParsecT String (H.HashMap Char Int) Identity a -> ParsecT String (H.HashMap Char Int) Identity a
p_braces = PT.braces lexer

p_symbol :: String -> ParsecT String (H.HashMap Char Int) Identity String
p_symbol = PT.symbol lexer

p_reserved :: String -> ParsecT String (H.HashMap Char Int) Identity ()
p_reserved = PT.reserved lexer

p_natural_or_float :: ParsecT String (H.HashMap Char Int) Identity (Either Integer Double)
p_natural_or_float = PT.naturalOrFloat lexer

p_commaSep1 :: ParsecT String (H.HashMap Char Int) Identity a -> ParsecT String (H.HashMap Char Int) Identity [a]
p_commaSep1 = PT.commaSep1 lexer

p_integer :: ParsecT String (H.HashMap Char Int) Identity Integer
p_integer = PT.integer lexer




ifCondParser :: Parser KSExpr
ifCondParser = do
    p_reserved "if"
    ifexpr <- exprParser
    p_reserved "then"
    thenexpr <- exprParser
    p_reserved "else"
    elseexpr <- exprParser
    return $ KSIf ifexpr thenexpr elseexpr 

forExprParser :: Parser KSExpr
forExprParser = do
    p_reserved "for"
    var_name <- p_identifier
    char '='
    expr1 <- exprParser
    char ','
    expr2 <- exprParser
    maybe_expr3 <- optionMaybe (do char ','; expr3 <- exprParser; return expr3)
    p_reserved "in"
    expr4 <- exprParser
    return $ KSFor var_name expr1 expr2 maybe_expr3 expr4 
    <?> "parse for loop"

binaryOpRhsParser :: Int -> KSExpr -> Parser KSExpr
binaryOpRhsParser expr_prec lhs = do
    (do
    p_whiteSpace
    bin_op <- lookAhead (oneOf "@%&^<>?+-*/")
    precTable <- getState
    let curr_bin_op_prec = fromJust $ H.lookup bin_op precTable

    if curr_bin_op_prec < expr_prec
        then return lhs
        else do 
            p_whiteSpace
            bin_op <- oneOf "@%&^<>?+-*/"
            p_whiteSpace
            expr1 <- unaryParser
            p_whiteSpace
            ((do
                next_bin_op <- lookAhead (oneOf "@%&^<>?+-*/")
                let next_bin_op_prec = fromJust $ H.lookup next_bin_op precTable
                if curr_bin_op_prec < next_bin_op_prec
                    then do 
                        expr2 <- binaryOpRhsParser (curr_bin_op_prec+1) expr1
                        res <- binaryOpRhsParser expr_prec (KSBinaryOp bin_op lhs expr2)
                        return res
                    else do
                        res <- binaryOpRhsParser expr_prec (KSBinaryOp bin_op lhs expr1)
                        return res)
                <|> (do 
                    res <- binaryOpRhsParser expr_prec (KSBinaryOp bin_op lhs expr1)
                    return res))
            
    )
    <|> return lhs
    <?> "parse binary op rhs"

unaryParser :: Parser KSExpr
unaryParser = do
    (do
    unary_op <- oneOf "!-"
    unaryParser
    )
    <|>
    (do 
    primaryParser
    )
    <?> "parse unary expression"


exprParser :: Parser KSExpr
exprParser = do 
    lhs <- unaryParser
    binaryOpRhsParser 0 lhs
    <?> "parse expression"

primaryParser :: Parser KSExpr
primaryParser = do
    p_whiteSpace
    (do
    numeric <- p_natural_or_float
    case numeric of 
        Left nat -> return $ KSNumber (fromIntegral nat)
        Right f -> return $ KSNumber f
    )
    <|>
    (do
    expr <- p_parens exprParser
    return $ expr
    )
    <|>
    (do
    try (do 
        p_whiteSpace
        func_name <- p_identifier
        func_args <- p_parens (exprParser `sepBy` (char ','))
        return $ KSCall func_name func_args 
        )
    <|>
    (do 
    var_name <- p_identifier
    return $ KSVariable var_name
    )
    )
    <|>
    (do
    expr <- ifCondParser
    return $ expr
    )
    <|>
    (do
    expr <- forExprParser
    return $ expr
    )
    <?> "parse primary expression"

protoParser :: Parser KSProto
protoParser = do
    (do 
    func_name <- p_identifier 
    func_args <- p_parens (do p_identifier `sepBy` (p_whiteSpace))
    return $ KSProto func_name func_args) 
    <|>
    (do 
    p_reserved "unary"
    op <- oneOf "!@%&^*<>?"
    precedence <- p_integer
    op_args <- p_parens (p_commaSep1 p_identifier)
    return $ KSBinOpProto [op] op_args precedence)
    <|>
    (do 
    p_reserved "binary"
    op <- oneOf "!@%&^*<>?"
    expr <- exprParser
    precedence <- p_integer
    op_args <- p_parens (p_commaSep1 p_identifier)
    return $ KSBinOpProto [op] op_args precedence)
    <?> "parse proto" 


funcParser :: Parser KSFunc
funcParser = do 
    p_reserved "def"
    proto <- protoParser
    expr <- exprParser
    return $ KSFunc proto expr
    <?> " parse function"


funcOrExprParser :: Parser KSFuncOrExprList
funcOrExprParser = do 
    (do
    p_whiteSpace
    eof
    return $ KSEmpty
    )
    <|>
    (do
    char ';'
    func_or_expr_list <- funcOrExprParser
    return func_or_expr_list
    )
    <|>
    (do 
    func <- funcParser
    func_or_expr_list <- funcOrExprParser
    return $ KSFuncNext func func_or_expr_list)
    <|>
    (do 
    expr <- exprParser
    func_or_expr_list <- funcOrExprParser
    return $ KSFuncNext (KSFunc (KSProto "" []) expr) func_or_expr_list)
    <?> "parse func or expr list"


sourceParser :: Parser KSProgram
sourceParser = do
    p_whiteSpace
    func_or_expr_list <- funcOrExprParser
    return $ KSProgram func_or_expr_list
    <?> "parse whole source file"


interactiveParser :: Parser (KSFunc)
interactiveParser = do
    (do 
    func <- funcParser
    char ';'
    return $ func)
    <|>
    (do 
    expr <- exprParser
    char ';'
    return $ KSFunc (KSProto "" []) expr)
    <?> "parse interactive mode"


codeGenExpr :: KSExpr -> CodeGenFunction Double (Value Double)
codeGenExpr (KSNumber num) = return $ valueOf num
codeGenExpr (KSVariable var_name) = undefined
codeGenExpr (KSUnaryOp op_name expr) = undefined
codeGenExpr (KSBinaryOp op_name lhs rhs) = do
    lhs_val <- codeGenExpr lhs
    rhs_val <- codeGenExpr rhs
    case op_name of 
        '+' -> do 
            r1 <- fadd lhs_val rhs_val
            return r1
        '-' -> do
            r1 <- fsub lhs_val rhs_val
            return r1
        '*' -> do 
            r1 <- fmul lhs_val rhs_val
            return r1
        '<' -> do
            r1 <- fcmp FPOLT lhs_val rhs_val
            r2 <- inttofp r1
            return $ r2
        _   -> undefined


codeGenExpr (KSCall func_name args) = undefined
codeGenExpr (KSIf ifexpr thenexpr elseexpr) = undefined
codeGenExpr (KSFor var_name assign_expr cond_expr term_expr body_expr) = undefined


codeGenFunc :: KSFunc -> CodeGenModule (Function (Double -> IO Double))
codeGenFunc (KSFunc proto expr) = do
    f <- codeGenProto proto
    defineFunction f $ \ _ -> do
        retval <- codeGenExpr expr
        ret retval
    return f


codeGenProto :: KSProto -> CodeGenModule (Function (Double -> IO Double))
codeGenProto (KSProto func_name args) = do
    f <- newNamedFunction InternalLinkage func_name
    return f

codeGenProto (KSBinOpProto func_name args prec) = undefined


data Opts = Opts
    { debug :: Bool
    , sourceFiles :: [FilePath]
    , interactive :: Bool
    } deriving (Data, Typeable, Show, Eq)


progOpts :: Opts
progOpts = Opts
    { sourceFiles = def &= args &= typ "source files"
    , debug = def &= help "debug mode"
    , interactive = def &= help "interactive mode"
    }

getOpts :: IO Opts
getOpts = cmdArgs $ progOpts
    &= summary (_PROGRAM_INFO ++ ", " ++ _COPYRIGHT)
    &= program _PROGRAM_NAME
    &= help _PROGRAM_DESC
    &= helpArg [explicit, name "help", name "h"]
    &= versionArg [explicit, name "version", name "v", summary _PROGRAM_INFO]

_PROGRAM_NAME = "kaleidoscope-plus"
_PROGRAM_VERSION = "0.0.1"
_PROGRAM_INFO = _PROGRAM_NAME ++ " version " ++ _PROGRAM_VERSION
_PROGRAM_DESC = "A dialect of kaleidoscope"
_COPYRIGHT = "BSD-3"


main :: IO ()
main = do
    hSetBuffering stdout LineBuffering
    hSetBuffering stdin LineBuffering
    opts <- getOpts

    let precTable = H.fromList [('<', 10), ('>', 10), ('+', 20), ('-', 20), ('*', 40)]

    if (not . null . sourceFiles) opts
    then do
        forM_ (sourceFiles $ opts) (\fileName -> 
            do prog_src <- readFile fileName
               case (runParser sourceParser precTable fileName prog_src) of
                 Left err -> print err
                 Right ksprog -> print ksprog
               putStrLn prog_src
            )
    else do
        initializeNativeTarget


        mod <- newModule
        mod_provider <- createModuleProviderForExistingModule mod
        fpm <- createFunctionPassManager mod_provider


        putStrLn "ready> "
        prog_src <- getContents
        case (runParser interactiveParser precTable "interactive" prog_src) of 
            Left err -> print err
            Right f -> do 
                print f
                v_func <- createModule $ codeGenFunc f
                dumpValue v_func