{-# LANGUAGE FlexibleContexts, RankNTypes #-}
module Compiler (compile) where
  
import qualified Data.Map as Map 
import Data.Map (Map, (!))
import Data.Maybe (fromJust)
import Control.Applicative ((<$>))
import AbsLambdaMini
import Data.Int
import LLVM.Core

type IntFunc = Function (IO Int32)
type Env = Map String IntFunc

compile :: String -> Program -> IO (Module, IntFunc)
compile name prog = do
  m <- newNamedModule name
  entry <- defineModule m $ do
    env <- compileProg prog
    wrapOutput (env ! "_main")
  return (m, entry)

wrapOutput :: IntFunc -> CodeGenModule IntFunc
wrapOutput entryPoint = createNamedFunction ExternalLinkage "main" $ do
  result <- call entryPoint
  printInt <- (externFunction "printInt" :: CodeGenFunction r (Function (Int32 -> IO ())))
  call printInt result
  ret (0 :: Int32)

compileProg :: Program -> CodeGenModule Env
compileProg (Prog fs) =
  let fnames = map (\(Func (FIdent name) _) -> name) fs
      predeclare :: String -> CodeGenModule (String, IntFunc)
      predeclare fname = do
        decl <- newNamedFunction InternalLinkage $ mangle fname
        return (mangle fname, decl)
  in do decls <- mapM predeclare fnames 
        let env = Map.fromList decls
        mapM_ (compileFunc env) fs
        return env

compileFunc env (Func (FIdent fname) exp) = 
  defineFunction (env ! mangle fname) $ compileExp env exp >>= ret
     
mangle = ("_" ++)
     
compileExp :: (IsInteger c, Num c, IsConst c) =>
              Map String (Function (IO c)) -> Exp -> CodeGenFunction r (Value c)
compileExp env (EInt n) = return $ valueOf $ fromIntegral n
compileExp env (EAdd e1 e2) = bind2 add (compileExp env e1) (compileExp env e2)
compileExp env (ESub e1 e2) = bind2 sub (compileExp env e1) (compileExp env e2)
compileExp env (EMul e1 e2) = bind2 mul (compileExp env e1) (compileExp env e2)
compileExp env (EDiv e1 e2) = bind2 sdiv (compileExp env e1) (compileExp env e2)

compileExp env (EVar (FIdent fname)) =  call $ env ! mangle fname 

-- This should be in Control.Monad
bind2 :: (Monad m) => (a -> b -> m c) -> m a -> m b -> m c
bind2 f m1 m2 = do
    x1 <- m1
    x2 <- m2
    f x1 x2