{-# LANGUAGE DeriveDataTypeable #-}
module Preprocessor (Program, ArithOp(..), CompareOp(..),
                     Exp(..), Type(..), preprocess, createClosures) where
  
import AbsLambdaMini (Type(..), FIdent(..))
import qualified AbsLambdaMini as Raw
import Data.Generics (gmapQ, gmapT, mkQ, mkT)
import Data.List (union, elemIndex, intersperse)
import Data.Data (Data)
import Data.Typeable (Typeable)
import Data.Map (Map, (!))
import qualified Data.Map as Map
import Data.Maybe (fromJust)

type Program = Map String Exp

data ArithOp = Add | Sub | Mult | Div
  deriving (Eq,Ord,Show,Data,Typeable)
-- Order here is important; same as LLVM enum (minus 32)
-- 'U' prefix is for unsigned version
data CompareOp = CompEq | CompNEq | CompUGT | CompUGEq | CompULT | CompULEq |
				 CompGT | CompGEq | CompLT | CompLEq
  				 deriving (Eq,Ord,Enum,Show,Data,Typeable)

data Exp =
   EGlobal String -- top level definition
 | ELocal String -- free and bound variables
 | EInt Integer
 | EBool Bool
 | EApp Exp Exp
 | ENeg Exp
 | EArith ArithOp   Exp Exp
 | EComp  CompareOp Exp Exp
 | EIf Exp Exp Exp
  -- Not used in compilation
 | ELambda (String, Type) Exp
 -- Created for compilation
 | EClosure [(String,Type)] (String,Type) Exp -- free vars, arg, body
  deriving (Eq,Ord,Data,Typeable,Show)

-- | Convert from BNFC-derived raw syntax in 'AbsLambdaMini' into something nicer.
preprocess :: Raw.Program -> Program
preprocess (Raw.Prog fs) = 
  renameMain $
  Map.fromList $
  map (\(Raw.Func (FIdent name) body) -> (name, ppExp [] body)) fs

-- The name "main" in LLVM designates the entry point of the program, which isn't
-- the same as the main function of the LambdaClassic program.
renameMain = Map.mapKeys $ \name ->
    case name of
        "main" -> "_main"
        other -> other

-- | Convert 'AbsLambdaMini.Exp' to 'Exp'. The only trick is determining if variable
--   references are global or local.
ppExp :: [(String, Type)] -> Raw.Exp -> Exp
ppExp ctx (Raw.EInt i) = EInt i
ppExp ctx (Raw.ETrue) = EBool True
ppExp ctx (Raw.EFalse) = EBool False
ppExp ctx (Raw.EApp e1 e2) = EApp (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.ENeg e) = ENeg $ ppExp ctx e
ppExp ctx (Raw.EMul e1 e2) = EArith Mult (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.EDiv e1 e2) = EArith Div (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.EAdd e1 e2) = EArith Add (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.ESub e1 e2) = EArith Sub (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.ELt  e1 e2) = EComp CompLT (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.EGt  e1 e2) = EComp CompGT (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.ELEq e1 e2) = EComp CompLEq (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.EGEq e1 e2) = EComp CompGEq (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.EEq  e1 e2) = EComp CompEq (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.ENEq e1 e2) = EComp CompNEq (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.EIf cond e1 e2) = EIf (ppExp ctx cond) (ppExp ctx e1) (ppExp ctx e2)
ppExp ctx (Raw.EVar (FIdent var)) =
  case lookup var ctx of
    Nothing -> EGlobal var
    Just _  -> ELocal var
ppExp ctx (Raw.ELambda (FIdent var) t e) = 
  ELambda arg $ ppExp (arg:ctx) e
    where arg = (var,t)

-- | We replace ELambda expression with
--   EClosure expressions that reference the free variables referenced in the lambda.
--   EClosure is interpreted by the 'Compiler' module to mean "create and initialize
--   a closure object here."
createClosures :: Program -> Program
createClosures p = Map.map (createClosures' []) p

createClosures' :: [(String, Type)] -> Exp -> Exp
createClosures' gamma (ELambda arg body) = 
  foldl createClosure (EClosure [] arg convertedSubExpr) $ freeVars (arg:gamma) convertedSubExpr
  where convertedSubExpr = createClosures' (arg:gamma) body
        createClosure (EClosure env arg' body') freeVar = EClosure (freeVar:env) arg' body'
createClosures' gamma exp = gmapT (mkT $ createClosures' gamma) exp

-- | find all free variables an expression given a binding context 'gamma' 
freeVars :: [(String,Type)] -> Exp -> [(String, Type)]
freeVars gamma (ELambda arg exp) = freeVars (arg:gamma) exp
-- properly created closures have no free variables
freeVars gamma (EClosure _ _ _) = []
-- It's a local, so we've got it in gamma somewhere. We only need to check against the
-- head of gamma to see if it's bound, since we're using single-variable args.
freeVars (arg@(argName, argT):gamma) (ELocal var) =
  if argName == var
    then []
    else [(var, fromJust $ lookup var gamma)]
freeVars gamma exp = unions $ gmapQ ([] `mkQ` (freeVars gamma)) exp
  where unions = foldl union []