{-# LANGUAGE TypeFamilies, OverloadedStrings #-}
-- | Freshen shadowed names in the source file,
--   and rewrite anonymous binders to their named forms.
module DDC.Source.Tetra.Transform.Freshen
        ( type S, evalState, newName
        , freshenModule)
where
import DDC.Source.Tetra.Transform.Freshen.State
import DDC.Source.Tetra.Module
import DDC.Source.Tetra.Exp
import qualified Control.Monad.State    as S
import qualified Data.Set               as Set
import qualified Data.Map.Strict        as Map
import qualified Data.List              as List


-------------------------------------------------------------------------------
-- | Freshen the given module.
freshenModule :: Module Source -> S (Module Source)
freshenModule mm
 = do   tops'   <- freshenTops $ moduleTops mm
        return  $  mm { moduleTops = tops' }


-------------------------------------------------------------------------------
-- | Freshen a top level thing.
freshenTops :: [Top Source] -> S [Top Source]
freshenTops tops
 = do
        let (topCls, topRest) = List.partition isTopClause tops
        let (sps,    cls)     = unzip $ [(sp, cl) | TopClause sp cl <- topCls]
        cls'        <- freshenClauseGroup cls
        let topCls' =  zipWith TopClause sps cls'
        return  $ topRest ++ topCls'


-- | Check if this top level thing is a clause.
isTopClause :: Top Source -> Bool
isTopClause tt
 = case tt of
        TopClause{}     -> True
        _               -> False


-------------------------------------------------------------------------------
-- | Freshen a clause group.
freshenClauseGroup :: [Clause] -> S [Clause]
freshenClauseGroup cls
 =  freshenClauseGroupBinds cls $ \cls'
 -> mapM freshenClauseBody cls'


-- | Freshen the binders in a group of clauses.
--
--   In the group, we may have several clauses that define the same function.
--   If we freshen the name of one of those clauses, we need to use the same
--   name for the other associated ones.
--
--   All the new names for the clauses are in scope in the continuation,
--   which can then be used to freshen the bodies.
--
freshenClauseGroupBinds 
        :: [Clause] -> ([Clause] -> S a) -> S a

freshenClauseGroupBinds cls0 cont
 = do   envTopX <- S.gets stateEnvX
        go envTopX [] cls0
 where  
        go _envTopX clsAcc [] 
         = cont (reverse clsAcc)

        -- Signatures.
        go envTopX clsAcc ((SSig a b t) : clsRest)
         = do   -- Freshen the signature type.
                t'      <- freshenType t

                -- Check if we've already renamed this variable in the
                -- current clause group.
                envX    <- S.gets stateEnvX
                case b of
                 BName name
                  -- We've already renamed one of the clauses in the same
                  -- group to a new name, so use that for matching ones.
                  | Just name' <- Map.lookup name (envRename envX)
                  -> do let b'   = BName name'
                        let cls' = SSig a b' t'
                        go envTopX (cls' : clsAcc) clsRest

                 -- Bind the variable into the current environment,
                 -- renaming it if it was already visible in the context
                 -- of the current clause group.
                 _ 
                  -> bindCtxBX envTopX b $ \b'
                  -> do let cls' = SSig a b' t'
                        go envTopX (cls' : clsAcc) clsRest


        -- Binding Clauses.
        go envTopX clsAcc ((SLet a (XBindVarMT b mt) ps gxs) : clsRest)
         = do   -- Freshen the type on the binder.
                mt'     <- traverse freshenType mt
                
                -- Check if we've already renamed this variable in the
                -- current clause group.
                envX    <- S.gets stateEnvX
                case b of
                 BName name
                  -- We've already renamed one of the clauses in the same
                  -- group to a new name, so use that for matching ones.
                  |  Just name'  <- Map.lookup name (envRename envX)
                  -> do let b'   = BName name'
                        let cls' = SLet a (XBindVarMT b' mt') ps gxs
                        go envTopX (cls' : clsAcc) clsRest

                 -- Bind the variable into the current environment,
                 -- renaming it if it was already visible in the context
                 -- of the current clause group.
                 _
                  -> bindCtxBX envTopX b $ \b'
                  -> do let cls' = SLet a (XBindVarMT b' mt') ps gxs
                        go envTopX (cls' : clsAcc) clsRest


-- | Freshen a clause.
freshenClauseBody :: Clause -> S Clause
freshenClauseBody cl
 = case cl of
        SSig{}
         -> return cl

        SLet a b ps gxs
         -> mapFreshBinds bindParam ps $ \ps'
         -> SLet a b ps' <$> mapM freshenGuardedExp gxs


-------------------------------------------------------------------------------
-- | Freshen and bind a function parameter.
bindParam :: Param -> (Param -> S a) -> S a
bindParam pp cont
 = case pp of
        MType b mt
         -> bindBT b $ \b'
         -> do  mt'     <- traverse freshenType mt
                cont $ MType b' mt'

        MWitness b mt
         -> bindBX b $ \b'
         -> do  mt'     <- traverse freshenType mt
                cont $ MWitness b' mt'

        MValue p mt
         -> bindPat p $ \p'
         -> do  mt'     <- traverse freshenType mt
                cont $ MValue p' mt'


-------------------------------------------------------------------------------
-- | Freshen a guarded expression.
freshenGuardedExp :: GuardedExp -> S GuardedExp 
freshenGuardedExp gx
 = case gx of
        GGuard g gx'
         -> bindGuard g $ \g'
         -> GGuard g' <$> freshenGuardedExp gx'

        GExp x
         -> do  x'      <- freshenExp x
                return  $  GExp x'


-------------------------------------------------------------------------------
-- | Freshen an expression.
freshenExp :: Exp -> S Exp
freshenExp xx
 = case xx of
        XAnnot a x      -> XAnnot a <$> freshenExp x
        XVar u          -> XVar     <$> boundUX u
        XPrim{}         -> return xx
        XCon{}          -> return xx

        XLAM (XBindVarMT b mt) x
         -> do  mt'     <- traverse freshenType mt
                bindBT b $ \b'
                 -> XLAM (XBindVarMT b' mt') <$> freshenExp x

        XLam (XBindVarMT b mt) x
         -> do  mt'     <- traverse freshenType mt
                bindBX b $ \b'
                 -> XLam (XBindVarMT b' mt') <$> freshenExp x

        XApp x1 x2      -> XApp  <$> freshenExp x1 <*> freshenExp x2

        XLet lts x      
         -> bindLets lts $ \lts'
         -> XLet lts' <$> freshenExp x

        XCase x alts    -> XCase    <$> freshenExp x 
                                    <*> mapM freshenAltCase alts

        XCast c x       -> XCast    <$> freshenCast c <*> freshenExp x

        XType t         -> XType    <$> freshenType t

        XWitness w      -> XWitness <$> freshenWitness w

        XDefix a xs     -> XDefix a <$> mapM freshenExp xs

        XInfixOp{}      -> return xx

        XInfixVar{}     -> return xx

        XMatch a as x   -> XMatch a <$> mapM freshenAltMatch as 
                                    <*> freshenExp x

        XWhere a x cl   -> XWhere a <$> freshenExp x <*> freshenClauseGroup cl

        XLamPat a p mt x
         -> do  mt'     <- traverse freshenType mt
                bindPat p $ \p' 
                 -> XLamPat a p' mt' <$> freshenExp x

        XLamCase a as   
         -> XLamCase a   <$> mapM freshenAltCase as


-------------------------------------------------------------------------------
-- | Freshen and bind guards.
bindGuard   :: Guard -> (Guard -> S a) -> S a
bindGuard gg cont
 = case gg of
        GPat p x       
         -> bindPat p $ \p'
         -> cont =<< (GPat p' <$> freshenExp x)

        GPred x  
         -> cont =<< (GPred   <$> freshenExp x)

        GDefault 
         -> cont GDefault


-------------------------------------------------------------------------------
-- | Freshen and bind let expressions.
bindLets :: Lets -> (Lets -> S a) -> S a
bindLets lts cont
 = case lts of
        LLet (XBindVarMT b mt) x
         -> do  mt'     <- traverse freshenType mt        
                bindBX b $ \b' 
                 -> cont =<< (LLet (XBindVarMT b' mt') <$> freshenExp x)

        LRec bxs
         -> do  let (bs, xs)    = unzip bxs
                mapFreshBinds bindBVX bs $ \bs'
                 -> do  xs'     <- mapM freshenExp xs
                        cont (LRec $ zip bs' xs')

        LPrivate brs mt bwts
         -> do  mt'     <- traverse freshenType mt
                mapFreshBinds bindBT brs $ \brs'
                 -> do  let (bws, ts)  = unzip bwts
                        ts'     <- mapM freshenType ts
                        mapFreshBinds bindBX bws $ \bws'
                         -> cont (LPrivate brs' mt' $ zip bws' ts')

        LGroup cls
         -> cont =<< (LGroup <$> freshenClauseGroup cls)


-------------------------------------------------------------------------------
-- | Freshen a case alternative.
freshenAltCase :: AltCase -> S AltCase
freshenAltCase (AAltCase p gxs)
 =  bindPat p $ \p'
 -> AAltCase p' <$> mapM freshenGuardedExp gxs


-- | Freshen a match alternative
freshenAltMatch :: AltMatch -> S AltMatch
freshenAltMatch (AAltMatch gx)
 = AAltMatch <$> freshenGuardedExp gx


-------------------------------------------------------------------------------
-- | Freshen a cast.
freshenCast :: Cast -> S Cast
freshenCast cc
 = case cc of
        CastWeakenEffect t -> CastWeakenEffect <$> freshenType t
        CastPurify w       -> CastPurify       <$> freshenWitness w
        CastBox            -> return CastBox
        CastRun            -> return CastRun


-------------------------------------------------------------------------------
-- | Freshen a witness.
freshenWitness :: Witness -> S Witness
freshenWitness ww
 = case ww of
        WAnnot a w      -> WAnnot a <$> freshenWitness w
        WVar  u         -> WVar  <$> boundUX u
        WCon {}         -> return ww
        WApp  w1 w2     -> WApp  <$> freshenWitness w1 <*> freshenWitness w2
        WType t         -> WType <$> freshenType t


-------------------------------------------------------------------------------
-- | Freshen a type.
freshenType :: Type -> S Type
freshenType tt
 = case tt of
        TAnnot a t      -> TAnnot a <$> freshenType t

        TCon{}          -> return tt

        TVar u          -> TVar <$> boundUT u

        TAbs b t1 t2
         -> do  t1'     <- freshenType t1
                bindBT b $ \b'
                 -> TAbs b' t1' <$> freshenType t2

        TApp t1 t2      -> TApp <$> freshenType t1 <*> freshenType t2


-------------------------------------------------------------------------------
-- | Bind a pattern.
bindPat  :: Pat -> (Pat -> S a) -> S a
bindPat pp cont
 = case pp of
        PDefault
         -> cont pp

        PAt  b p
         -> bindBX  b $ \b'
         -> bindPat p $ \p'
         -> cont (PAt b' p')

        PVar b
         -> bindBX b  $ \b'
         -> cont (PVar b')

        PData dc ps     
         -> mapFreshBinds bindPat ps $ \ps' 
         -> cont (PData dc ps')


-------------------------------------------------------------------------------
-- | Bind a new type variable.
bindBT :: Bind -> (Bind -> S a) -> S a
bindBT b@BNone cont
 = cont b

bindBT BAnon cont
 = do   -- Create a new name for anonymous binders.
        name    <- newName "t"

        withModifiedEnvT
         (\envT -> envT { envStack    = name : envStack envT
                        , envStackLen = 1 + envStackLen envT })
         $ cont (BName name)

bindBT (BName n) cont
 =  S.get >>= \state0 
 -> case Set.member n (envNames $ stateEnvT state0) of
     -- If the binder does not shadow an existing one
     -- then don't bother rewriting it.
     False 
      ->     withModifiedEnvT
              (\envT -> envT { envNames = Set.insert n (envNames envT) })
              $ cont (BName n)

     -- The binder shadows an existing one, so rewrite it.
     True 
      -> do name    <- newName "t"

            -- Run the continuation in the extended environment.
            withModifiedEnvT
             (\envT -> envT { envNames  = Set.insert   name (envNames  envT) 
                            , envRename = Map.insert n name (envRename envT)})
             $ cont (BName name)


-------------------------------------------------------------------------------
-- | Bind a term variable with its attached type.
bindBVX :: BindVarMT -> (BindVarMT -> S a) -> S a
bindBVX (XBindVarMT b mt) cont
 = do   mt'     <- traverse freshenType mt
        bindBX b $ \b'
         -> cont (XBindVarMT b' mt')


-- | Bind a new term variable,
--   renaming it if it is already in the current environment.
bindBX :: Bind -> (Bind -> S a) -> S a
bindBX b cont
 = do   envX    <- S.gets stateEnvX
        bindCtxBX envX b cont


-- | Bind a new term variable,
--   renaming it if it is already in the given environment.
bindCtxBX :: Env -> Bind -> (Bind -> S a) -> S a

bindCtxBX _envCtx b@BNone cont
 = cont b

bindCtxBX _envCtx BAnon cont
 = do   -- Create a new name for anonymous binders.
        name    <- newName "x"

        withModifiedEnvX 
         (\envX -> envX { envStack    = name : envStack envX
                        , envStackLen = 1 + envStackLen envX })
         $ cont (BName name)

bindCtxBX envCtx (BName name) cont
 = case Set.member name (envNames envCtx) of
        -- If the binder does not shadow an existing one
        -- then don't bother rewriting it.
        False 
         -> do  withModifiedEnvX
                 (\envX -> envX
                        {  envNames = Set.insert name (envNames envX)})
                 $ cont (BName name)

        -- The binder shadows an existing one, so rewrite it.
        True 
         -> do  name'   <- newName name

                -- Run the continuation in the environment extended with the
                -- new name.
                withModifiedEnvX
                 (\envX -> envX 
                        {  envNames  = Set.insert      name' (envNames  envX) 
                        ,  envRename = Map.insert name name' (envRename envX)})
                 $ cont (BName name')


-------------------------------------------------------------------------------
-- | Rewrite bound type variable if needed.
boundUT :: Bound -> S Bound
boundUT uu
 = case uu of
        UName n
         -> do  envT    <- S.gets (envRename   . stateEnvT)
                case Map.lookup n envT of
                 Just n'  -> return $ UName n'
                 _        -> return uu

        UIx i
         -> do  stack   <- S.gets (envStack    . stateEnvT)
                len     <- S.gets (envStackLen . stateEnvT)
                if (i < len)
                 then   return $ UName (stack !! i)
                 else   return uu

        UHole{}
         -> return uu


-------------------------------------------------------------------------------
-- | Rewrite bound term variable if needed.
boundUX :: Bound -> S Bound
boundUX uu
 = case uu of
        UName n
         -> do  envX    <- S.gets (envRename   . stateEnvX)
                case Map.lookup n envX of
                 Just n'  -> return $ UName n'
                 _        -> return uu

        UIx i
         -> do  stack   <- S.gets (envStack    . stateEnvX)
                len     <- S.gets (envStackLen . stateEnvX)
                if (i < len)
                 then   return $ UName (stack !! i)
                 else   return uu

        UHole{}
         -> return uu