{-# 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 qualified LLVMUtils as U
import qualified LLVM.FFI.Core as FFI
import CodeGenMonad
import Instructions

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

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

-- Generate declarations for external C functions we're using to implement primitives
builtInFunctions = 
  let genericFun (name, arity) = do
  		func <- newNamedFunction (genericPrimitiveType arity) name
  		return (name, func)
  in Map.fromList <$> mapM genericFun
  		[ ("natElim", 3), ("finElim", 4), ("vecElim", 4), ("eqElim", 4),
  		  ("Cons", 2), ("Nil", 0), ("Succ",1), ("FSucc", 1), ("Refl", 1) ]

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
    
compile :: String -> Program -> IO (Module, Function)
compile moduleName prog =
    runCodeGen moduleName $ compile' prog

compile' :: Program -> CodeGen Function
compile' prog =
  let predeclare fname = do
        ftype <- liftIO $ makeType TGeneric
        decl <- newNamedFunction (constFunType ftype) fname
        return (fname, decl)
  in do decls <- Map.fromList <$> mapM (predeclare . fst) prog
        builtins <- builtInFunctions
        let env = Env { funcs = decls `Map.union` builtins }
        runReaderT (mapM_ compileFunc prog) env
        case Map.lookup "_main" decls of
        	Nothing -> fail "internal error: missing main function"
        	Just entryPoint -> return entryPoint

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 <- unboxInt =<< compileExp e1
    v2 <- unboxInt =<< compileExp e2
    n <- lift $ withCurrentBuilder $ \ bld -> U.withEmptyCString $ op bld v1 v2
    boxInt n

integerType = FFI.integerType 32

boxInt :: Value -> Compiler Value
boxInt n = lift $ do
    ptr <- malloc integerType
    store ptr n
    bitcast ptr opaquePtr

unboxInt box = lift $ do
	ptrToInt <- liftIO $ flip FFI.pointerType 0 <$> makeType TNat
	ptr <- bitcast box ptrToInt
	load ptr
    
compileExp :: Exp -> Compiler Value
compileExp (EInt n) =
	boxInt $ FFI.constInt integerType (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
    
compileExp (EGlobal fname) = do
    func <- lookDef fname
    lift $ call func []

compileExp (ELocal name) = lift $ getEnv name

compileExp exp@(EClosure env (argName, argType) (body, bodyT)) = 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).
    let envTypes = replicate (length env) opaquePtr
    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
    resultT <- liftIO $ makeType bodyT
    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.
    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
	-- compiling left-hand side gives us a closure
    closure <- compileExp func
    genericClosureT <- liftIO $ flip FFI.pointerType 0 <$> closureType opaquePtr opaquePtr
    genericClosure <- lift $ bitcast closure genericClosureT
 
    -- compiling right-hand side gives us some opaque value
    arg <- compileExp exp
    genericArg <- lift $ bitcast arg opaquePtr
    
    -- extract the fields from the closure...
    funcPtr <- lift $ load =<< getStructField genericClosure 0
    env  <- lift $ load =<< getStructField genericClosure 1 

    -- ...and invoke it
    lift $ call funcPtr [env, genericArg]
    
-- Primitives are just a straightforward function call in LLVM 
compileExp (EPrimitive primName args) = do
	primitive <- lookDef primName
	args' <- mapM compileExp args
	lift $ call primitive args'

-- Since we aren't doing type erasure, we need to pass around *something* for the
-- type arguments are runtime.
compileExp (EType) = compileExp (EInt 0xBADC0DE)

-- | 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
            load =<< getStructField envStruct n

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