{-# Language TypeSynonymInstances,FlexibleInstances,MultiParamTypeClasses,FunctionalDependencies,RankNTypes,FlexibleContexts,KindSignatures,ScopedTypeVariables #-}

{-

 Inlining functions on values and moving let-bindings (of variables) to the top of the function. (4th step)

 Assumption:
  
  - let-bindings are in the order consistent of their dependencies

-}

module Inlining where

import Spec


import Control.Monad.State
import Data.Maybe
import Data.List
import Numeric (showHex)
import Debug.Trace

import ASTData
import TypeChecker

type DBinding = (DVarName, DSmplDef DASTData)
type DBindings = [DBinding]
type DEnvI = (DBindings, [Int], DUnique, [[DSmplDef DASTData]])


-- entry point
runInlining :: DProgramSpec DASTData -> DUnique -> (DProgramSpec DASTData, DUnique)
runInlining p uid = let (p', (_,_,uid',_)) = runState (inlining p) ([],[],uid,[[]])
                    in (p', uid')

class Inliable a where
    inlining :: a -> State DEnvI a  -- inliging functions, and storing let-bindings contained in functions to the env.
    inlining a = do return a  -- by default, returns the input

getBinds :: State DEnvI DBindings
getBinds = do xx <- get
              let (bds,_,_,_) = xx
              return bds

addBinds :: DBindings -> State DEnvI ()
addBinds bds =
    do xx <- get
       let (bds', ks, i, vbds) = xx
       put (bds++bds', length bds:ks, i, vbds)

popBinds :: State DEnvI ()
popBinds =
    do xx <- get
       let (bds, k:ks, i, vbds) = xx
       put (drop k bds, ks, i, vbds)

getNewName :: String -> State DEnvI DVarName
getNewName s = 
    do xx <- get 
       let (bds, ks, i, vbds) = xx
           (n, i') = genNewName i (getBaseName s)
       put (bds, ks, i', vbds)
       return n

addVarBinds :: [DSmplDef DASTData] -> State DEnvI ()
addVarBinds vbds = 
    do xx <- get
       let (bds, ks, i, vbds':vbdss) = xx
       put (bds, ks, i, (reverse vbds ++ vbds'):vbdss)

getVarBinds :: State DEnvI [DSmplDef DASTData]
getVarBinds = 
    do xx <- get
       let (bds, ks, i, vbds:vbdss) = xx
       put (bds, ks, i, vbdss)
       return (reverse vbds)

saveVarBinds :: State DEnvI ()
saveVarBinds = 
    do (bds, ks, i, vbdss) <- get
       put (bds, ks, i, []:vbdss)

instance Inliable (DProgramSpec DASTData) where
    inlining (DProgramSpec rs p a) = 
        do p' <- inlining p
           return (DProgramSpec rs p' a)

instance Inliable (DRecordSpec DASTData) 
instance Inliable (DType DASTData) 
instance Inliable (DConst DASTData)

isFunDefG (DGDefSmpl def a) = isFunDef def
isFunDefG _ = False

isSmplDef (DGDefSmpl def a) = True
isSmplDef _ = False

isFunDefGV _ = False

isFunDef (DDefFun f vs ds e a) = True
isFunDef _ = False

getSmplDef (DGDefSmpl def a) = def


addFunDefBindG :: (DGroundDef DASTData) -> State DEnvI ()
addFunDefBindG def = 
    do case def of 
         (DGDefSmpl def' a) -> addFunDefBind def'
         _ -> return ()


addFunDefBind :: (DSmplDef DASTData) -> State DEnvI ()
addFunDefBind def = 
    do case def of 
         (DDefFun f vs ds e a) -> addBinds [(getName f, def)]
         _ -> return ()

popFunDefBindG :: (DGroundDef DASTData) -> State DEnvI ()
popFunDefBindG def = 
    do case def of 
         (DGDefSmpl def a) -> popFunDefBind def
         _ -> return ()

popFunDefBindGV :: (DDefGraphVar DASTData) -> State DEnvI ()
popFunDefBindGV def = 
    do return ()

popFunDefBind :: (DSmplDef DASTData) -> State DEnvI ()
popFunDefBind def = 
    do case def of 
         (DDefFun f vs ds e a) -> popBinds
         _ -> return ()


inliningLetG defs e =
 do defs' <- mapM (\def -> do def' <- inlining def; addFunDefBindG def'; return def') defs
    e' <- inlining e
    mapM (\x -> do popFunDefBindG x) defs
    let defs'' = filter (not . isFunDefG) defs'
        defs''' = filter (not . isSmplDef) defs''
        vdefs = map getSmplDef $ filter (isSmplDef) defs''
    addVarBinds vdefs
    return (defs''', e')       

-- no room for inlining
inliningLetGV defs e = return (defs, e)

inliningLet defs e =
 do defs' <- mapM (\def -> do def' <- inlining def; addFunDefBind def'; return def') defs
    let defs'' = filter (not . isFunDef) defs'
    addVarBinds defs''
    e' <- inlining e
    mapM (\x -> do popFunDefBind x) defs
    return (e')       

instance Inliable (DProg DASTData) where
    inlining (DProg f defs e a) = 
        do (defs', e') <- inliningLetG defs e
           vdefs <- getVarBinds               -- the inner var defs
           let vdefs' = map (\def -> DGDefSmpl def (getData def)) vdefs
           return (DProg f (vdefs' ++ defs') e' a)

instance Inliable (DGroundDef DASTData) where
    inlining (DGDefVI d a) = do d' <- inlining d
                                return (DGDefVI d' a)
    inlining (DGDefVC d a) = do d' <- inlining d
                                return (DGDefVC d' a)
    inlining (DGDefGV d a) = do d' <- inlining d
                                return (DGDefGV d' a)
    inlining (DGDefGF d a) = do d' <- inlining d
                                return (DGDefGF d' a)
    inlining (DGDefSmpl d a) = do d' <- inlining d
                                  return (DGDefSmpl d' a)

instance Inliable (DDefVertComp DASTData) where
    inlining (DDefVertComp f defs e a) = 
        do saveVarBinds
           e' <- inliningLet defs e
           vdefs <- getVarBinds  -- the inner var definitions
           return (DDefVertComp f vdefs e' a)

instance Inliable (DDefVertInit DASTData) where
    inlining (DDefVertInit f defs e a) = 
        do saveVarBinds
           e' <- inliningLet defs e 
           vdefs <- getVarBinds  -- the inner var definitions
           return (DDefVertInit f vdefs e' a)

instance Inliable (DDefGraphVar DASTData) where
    inlining (DDefGraphVar v e a) = 
        do (_, e') <- inliningLetG [] e 
           return (DDefGraphVar v e' a)

instance Inliable (DDefGraphFun DASTData) where
    inlining (DDefGraphFun f v defs e a) = 
        do (defs', e') <- inliningLetGV defs e 
           return (DDefGraphFun f v defs' e' a)

instance Inliable (DSmplDef DASTData) where
    inlining (DDefFun f vs defs e a) = 
        do saveVarBinds    -- save the outer var definitions
           e' <- inliningLet defs e
           vdefs' <- getVarBinds   -- get the inner var definitions
           return (DDefFun f vs vdefs' e' a)
    inlining (DDefVar v defs e a) = 
        do e' <- inliningLet defs e 
           return (DDefVar v [] e' a)  -- no inner definitions
    inlining (DDefTuple vs defs e a) = 
        do e' <- inliningLet defs e 
           return (DDefTuple vs [] e' a)  -- no inner definitions

instance Inliable (DTermination DASTData) where
    inlining (DTermF a) = return (DTermF a)
    inlining (DTermI e a) = 
        do e' <- inlining e
           return (DTermI e' a)
    inlining (DTermU e a) = 
        do e' <- inlining e
           return (DTermU e' a)

instance Inliable (DGraphExpr DASTData)

refreshNameV :: DVar DASTData -> State DEnvI (DVar DASTData)
refreshNameV (DVar v a) =
    do n <- getNewName v
       return (DVar n a)

refreshName :: DSmplDef DASTData -> State DEnvI (DSmplDef DASTData)
refreshName (DDefVar v [] e a) = 
    do v' <- refreshNameV v
       return (DDefVar v' [] e a)
refreshName (DDefTuple vs [] e a) = 
    do vs' <- mapM refreshNameV vs
       return (DDefTuple vs' [] e a)

refreshNameV2 :: (DASTData, DVar DASTData) -> State DEnvI (DVar DASTData)
refreshNameV2 (a, DVar v _) =
    do n <- getNewName v
       return (DVar n a)

replaceNamesD :: [(DVarName, DVarName)] -> [(DVarName, DExpr DASTData)] -> DSmplDef DASTData -> DSmplDef DASTData
replaceNamesD subsv subse (DDefVar v [] e a) = (DDefVar v [] (replaceNames subsv subse e) a)
replaceNamesD subsv subse (DDefTuple vs [] e a) = (DDefTuple vs [] (replaceNames subsv subse e) a)

replaceNames :: [(DVarName, DVarName)] -> [(DVarName, DExpr DASTData)] -> DExpr DASTData -> DExpr DASTData
replaceNames subsv subse e = rec e
  where
    rec (DIf p t e a) = 
        let p' = rec p
            t' = rec t
            e' = rec e
        in (DIf p' t' e' a)
    rec (DTuple es a) =
        let es' = map rec es 
        in (DTuple es' a)
    rec (DFunAp f es a) =
        let es' = map rec es
        in (DFunAp f es' a)
    rec (DConsAp c es a) =
        let es' = map rec es
        in (DConsAp c es' a)
    rec (x@(DFieldAcc t fs a)) = x
    rec (x@(DFieldAccE e fs a)) = x
    rec (DAggr a' e g es a) =
        let e'  = rec e
            es' = map rec es
        in (DAggr a' e' g es' a)
    rec (x@(DVExp (DVar v av) a)) = 
        case lookup v subsv of
          Just v' -> (DVExp (DVar v' av) a)
          Nothing -> case lookup v subse of
                       Just e -> e
                       Nothing -> x
    rec (x@(DCExp c a)) = x


updateType :: forall (t :: * -> *). 
            (DAdditionalData2 (t DASTData) DASTData DASTData (t DASTData), 
             DAdditionalData (t DASTData) DASTData
            ) 
            => Substitution -> t DASTData -> t DASTData
updateType s x = mapData (\a -> a { typeOf = apply s (typeOf a)}) x

instantiation :: DFun DASTData -> [DExpr DASTData] -> DASTData -> State DEnvI (DExpr DASTData)
instantiation f es a = 
    do bds <- getBinds
       case lookup (getName f) bds of
         Just (DDefFun f vs defs e a') ->         -- defs are of DefVar/DDefTuple
             do let ets = map (typeOf.getData) es      -- types of the arguments
                    eas = map (getData) es             -- types and deps of the arguments 
                defs' <- mapM refreshName defs         -- give them new names
                vs' <- mapM refreshNameV2 (zip eas vs) -- new variable names with the argument types
                let t = typeOf a                       -- type of the expression
                    tf = typeOf (getData f)            -- type of the function
                    subs = unify [(tf, typeFunction (ets++[t]))] -- the type-substitution
                    defs'' = map (updateType subs) defs'
                    subsv = zip (concatMap getNames defs) (concatMap getNames defs'') -- [(oldvar,newvar)]
                    defs2 = zipWith (\v e -> DDefVar v [] e (getData v)) vs' es  -- arguments are bound to the new variables
                    subse = zip (map getName vs) (map (\v -> DVExp v (getData v)) vs')  -- [(arg, exp)]
                    defs''' = map (replaceNamesD subsv subse) defs''
                    e' = replaceNames subsv subse e
                    e'' =  updateType subs e'
                addVarBinds (defs''' ++ defs2)        -- these are moved to immediately under the VertComp/VertInit
                return e''
         Nothing -> return (DFunAp f es a)  -- built-in functions
  
instance Inliable (DExpr DASTData) where
    inlining (DIf p t e a) = 
        do p' <- inlining p
           t' <- inlining t
           e' <- inlining e
           return (DIf p' t' e' a)
    inlining (DTuple es a) =
        do es' <- mapM inlining es 
           return (DTuple es' a)
    inlining (DFunAp f es a) = -- inlining!
        do es' <- mapM inlining es
           instantiation f es' a
    inlining (DConsAp c es a) =
        do c' <- inlining c
           es' <- mapM inlining es
           return (DConsAp c' es' a)
    inlining (DFieldAcc t fs a) =
        do return (DFieldAcc t fs a)
    inlining (DFieldAccE e fs a) =
        do return (DFieldAccE e fs a)
    inlining (DAggr a' e g es a) =
        do e' <- inlining e
           es' <- mapM inlining es
           return (DAggr a' e' g es' a)
    inlining (DVExp v a) =
        do return (DVExp v a)
    inlining (DCExp c a) = 
        do return (DCExp c a)

instance Inliable (DGen DASTData) 
instance Inliable (DEdge DASTData)
instance Inliable (DTableExpr DASTData)
instance Inliable (DField DASTData) 
instance Inliable (DFun DASTData) 
instance Inliable (DVar DASTData) 
instance Inliable (DConstructor DASTData)