{-# LANGUAGE DeriveDataTypeable #-}
module Compiler (compile) where
  
import qualified Data.Map as Map 
import Data.Map (Map, (!), keys)
import Data.List (elemIndex)
import Data.Maybe (fromJust)
import Control.Applicative ((<$>))
import Data.Int
import Control.Monad (liftM, liftM2, liftM3)
import Control.Monad.Reader
import Data.Generics (gmapQ, mkQ)
import Data.Typeable (Typeable, cast)
import Foreign.Marshal.Array (withArray)

import Preprocessor
import Typechecker (TypeDict)
import qualified LLVMUtils as U
import qualified LLVM.FFI.Core as FFI
import CodeGenMonad
import Instructions

data Env = Env { types :: TypeDict, funcs :: Map String Function }
type Compiler a = ReaderT Env CodeGen a

type FFIBinOp = FFI.BuilderRef -> FFI.ValueRef -> FFI.ValueRef -> U.CString -> IO FFI.ValueRef

{- A couple accessor functions for the Compiler state -}

lookDef :: String -> Compiler Function
lookDef name = do
  functions <- asks funcs
  case Map.lookup name functions of
    Just f -> return f
    Nothing -> fail $ "Unknown function " ++ name
  
lookType prog types name =
    Map.findWithDefault (error $ "internal error looking up " ++ show name ++ " in " ++ show types )
                        (EGlobal name) types
    
    
compile :: String -> Program -> TypeDict -> IO (Module, Function)
compile moduleName prog types =
    runCodeGen moduleName $ compile' prog types

compile' :: Program -> TypeDict -> CodeGen Function
compile' prog types =
  let predeclare fname = do
        ftype <- liftIO $ makeType $ lookType prog types fname
        decl <- newNamedFunction (constFunType ftype) fname
        return (fname, decl)
  in do decls <- Map.fromList <$> mapM predeclare (keys prog)
        let env = Env { types=types, funcs=decls }
        runReaderT (mapM_ compileFunc $ Map.toList prog) env
        return $ decls ! "_main"

compileFunc :: (String, Exp) -> Compiler ()
compileFunc (fname, exp) = do
  f <- lookDef fname
  env <- ask
  -- down-convert from Compiler to CodeGen monad
  let body = runReaderT (compileExp exp) env
  lift $ defineFunction f body Map.empty
      
compileBinOp :: FFIBinOp -> Exp -> Exp -> Compiler Value
compileBinOp op e1 e2 = do
    v1 <- compileExp e1
    v2 <- compileExp e2
    lift $ withCurrentBuilder $ \ bld ->
        U.withEmptyCString $ op bld v1 v2

compileExp (EBool b) = return $
    FFI.constInt (FFI.integerType 1) (if b then 1 else 0) (fromIntegral 1)
compileExp (EInt n) = return $
    FFI.constInt (FFI.integerType 32) (fromIntegral n) (fromIntegral 1)
compileExp (EArith Add e1 e2) = compileBinOp FFI.buildAdd e1 e2
compileExp (EArith Mult e1 e2) = compileBinOp FFI.buildMul e1 e2
compileExp (EArith Sub e1 e2) = compileBinOp FFI.buildSub e1 e2
compileExp (EArith Div e1 e2) = compileBinOp FFI.buildSDiv e1 e2
-- We borrow a trick from the LLVM package: the order of comparison ops
-- in CompareOp matches that used by LLVM internally.
compileExp (EComp op e1 e2) = compileBinOp cmp e1 e2
    where cmp = flip FFI.buildICmp $ fromIntegral $ fromEnum op + 32

-- If statements are pretty straightforward. The only trick is that we have to
-- create a new block to represent "the rest of the function" and also use a
-- phi instruction to determine which value to use going forward, depending on
-- which branch of the if we took.
compileExp exp@(EIf cond e1 e2) = do
    b1 <- lift newBasicBlock
    b2 <- lift newBasicBlock
    continuation <- lift newBasicBlock
    
    condVal <- compileExp cond
    lift $ withCurrentBuilder $ \bld -> FFI.buildCondBr bld condVal b1 b2
    
    lift $ defineBasicBlock b1
    v1 <- compileExp e1
    lift $ withCurrentBuilder $ \bld -> FFI.buildBr bld continuation
    
    lift $ defineBasicBlock b2
    v2 <- compileExp e2
    lift $ withCurrentBuilder $ \bld -> FFI.buildBr bld continuation
    
    lift $ defineBasicBlock continuation
    resultType <- asks $ fromJust . Map.lookup exp . types
    lift $ phi resultType [(v1, b1), (v2, b2)]
    
compileExp (EGlobal fname) = do
    fdict <- asks funcs
    lift $ call (fdict ! fname) []

compileExp (ELocal name) = lift $ getEnv name

compileExp exp@(EClosure env (argName, argType) body) = do
    -- Create and initialize a struct to hold the environment part of the closure. Also
    -- create a map of code to run when we want to access the environment (or the arg).
    envTypes <- liftIO $ mapM makeType $ map snd env
    envStructT <- liftIO $ makeStruct envTypes
    liftIO $ showType envStructT
    envStruct <- lift $ malloc envStructT
    envBindings <- lift $ initializeEnv envStruct argName env
    
    -- Create the function's type as an LLVM type object
    (TFun _ resultT)  <- asks (fromJust . Map.lookup exp . types)
    resultT' <- liftIO $ makeType resultT
    argT <- liftIO $ makeType argType
    let funcT = U.functionType False resultT' [FFI.pointerType envStructT 0, argT]
    
    -- Define the function we're creating the closure for (recurses within the
    -- CodeGen monad).
    func <- lift $ newFunction funcT
    env <- ask
    let bodyCode = runReaderT (compileExp body) env
    lift $ defineFunction func bodyCode envBindings
    
    -- Create and fill in the closure object. The calling function doesn't need to know
    -- the detailed type of the function or the closure. It just takes the function part
    -- and passes it the argument and environment. So we use the bitcast operation to
    -- hide these type details.
    
    closureT <- liftIO $ closureType argT resultT'
    closure  <- lift $ malloc closureT
    funcPart <- lift $ getStructField closure 0
    envPart  <- lift $ getStructField closure 1
    castedFuncPart <- lift $ bitcast funcPart $ FFI.pointerType (FFI.pointerType funcT 0) 0
    
    lift $ store castedFuncPart func 
    lift $ store envPart =<< bitcast envStruct opaquePtr
    
    return closure
      
compileExp appExp@(EApp func exp) = do
    closure <- compileExp func
    arg <- compileExp exp
    funcPtr <- lift $ load =<< getStructField closure 0
    env  <- lift $ load =<< getStructField closure 1 

    lift $ call funcPtr [env, arg]
    
-- | Create a map from names in the environment (argument and free variables) to
--   the LLVM code needed to access them. As a side effect, initializes an
--   environment struct for use in a closure.
initializeEnv :: Value -> String -> [(String, Type)] -> CodeGen (Map String (CodeGen Value))
initializeEnv envStruct argName env = do
    mapM_ initEnvField $ map fst env
    let env' = Map.fromList $ map (\(n, (name, _)) -> (name, envLookup n)) (zip [0..] env)
    return $ Map.insert argName (do func <- getFunction; return $ FFI.getParam func 1) env'
    where envLookup :: Int -> CodeGen Value
          envLookup n = do
            func <- getFunction
            let envStruct = FFI.getParam func 0
            structT <- liftIO $ showType =<< FFI.typeOf envStruct
            load =<< getStructField envStruct n

          initEnvField name = do
            superEntry <- getEnv name
            field <- getStructField envStruct $ fromJust $ elemIndex name (map fst env)
            store field superEntry