{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}

-- | The 'ZonkEnv' zonking environment, and the 'ZonkT' and 'ZonkBndrT'
-- monad transformers, for the final zonking to type in "GHC.Tc.Zonk.Type".
--
-- See Note [Module structure for zonking] in GHC.Tc.Zonk.Type.
module GHC.Tc.Zonk.Env
  ( -- * The 'ZonkEnv'
    ZonkEnv(..), getZonkEnv
  , ZonkFlexi(..)
  , initZonkEnv

    -- * The 'ZonkT' and 'ZonkBndrT' zonking monad transformers
  , ZonkT(ZonkT,runZonkT), ZonkBndrT(..)

    -- ** Going between 'ZonkT' and 'ZonkBndrT'
  , runZonkBndrT
  , noBinders, don'tBind

    -- ** Modifying and extending the 'ZonkEnv' in 'ZonkBndrT'
  , setZonkType
  , extendZonkEnv
  , extendIdZonkEnv, extendIdZonkEnvRec
  , extendTyZonkEnv

  )
  where

import GHC.Prelude

import GHC.Core.TyCo.Rep ( Type )
import GHC.Types.Var ( TyCoVar, Var, TyVar )

import GHC.Types.Var ( Id, isTyCoVar )
import GHC.Types.Var.Env

import GHC.Utils.Monad.Codensity
import GHC.Utils.Outputable

import Control.Monad.Fix         ( MonadFix(..) )
import Control.Monad.IO.Class    ( MonadIO(..) )
import Control.Monad.Trans.Class ( MonadTrans(..) )
import Data.Coerce               ( coerce )
import Data.IORef                ( IORef, newIORef )
import Data.List                 ( partition )

import GHC.Exts                  ( oneShot )

--------------------------------------------------------------------------------

-- | See Note [The ZonkEnv]
data ZonkEnv
  = ZonkEnv { ze_flexi       :: !ZonkFlexi
            , ze_tv_env      :: TyCoVarEnv TyCoVar
            , ze_id_env      :: IdEnv      Id
            , ze_meta_tv_env :: IORef (TyVarEnv Type) }

-- | How should we handle unfilled unification variables in the zonker?
--
-- See Note [Un-unified unification variables]
data ZonkFlexi
  = DefaultFlexi    -- ^ Default unbound unification variables to Any
  | SkolemiseFlexi  -- ^ Skolemise unbound unification variables
                    --   See Note [Zonking the LHS of a RULE]
  | RuntimeUnkFlexi -- ^ Used in the GHCi debugger
  | NoFlexi         -- ^ Panic on unfilled meta-variables
                    -- See Note [Error on unconstrained meta-variables]
                    -- in GHC.Tc.Utils.TcMType

instance Outputable ZonkEnv where
  ppr (ZonkEnv { ze_tv_env = tv_env
               , ze_id_env = id_env })
    = text "ZE" <+> braces (vcat
         [ text "ze_tv_env =" <+> ppr tv_env
         , text "ze_id_env =" <+> ppr id_env ])

{- Note [The ZonkEnv]
~~~~~~~~~~~~~~~~~~~~~
* ze_flexi :: ZonkFlexi says what to do with a
  unification variable that is still un-unified.
  See Note [Un-unified unification variables]

* ze_tv_env :: TyCoVarEnv TyCoVar promotes sharing. At a binding site
  of a tyvar or covar, we zonk the kind right away and add a mapping
  to the env. This prevents re-zonking the kind at every
  occurrence. But this is *just* an optimisation.

* ze_id_env : IdEnv Id promotes sharing among Ids, by making all
  occurrences of the Id point to a single zonked copy, built at the
  binding site.

  Unlike ze_tv_env, it is knot-tied: see extendIdZonkEnvRec.
  In a mutually recursive group
     rec { f = ...g...; g = ...f... }
  we want the occurrence of g to point to the one zonked Id for g,
  and the same for f.

  Because it is knot-tied, we must be careful to consult it lazily.

* ze_meta_tv_env: see Note [Sharing when zonking to Type]


Notes:
  * We must be careful never to put coercion variables (which are Ids,
    after all) in the knot-tied ze_id_env, because coercions can
    appear in types, and we sometimes inspect a zonked type in
    the GHC.Tc.Zonk.Type module.  [Question: where, precisely?]

  * An obvious suggestion would be to have one VarEnv Var to
    replace both ze_id_env and ze_tv_env, but that doesn't work
    because of the knot-tying stuff mentioned above.

Note [Un-unified unification variables]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
What should we do if we find a Flexi unification variable?
There are three possibilities:

* DefaultFlexi: this is the common case, in situations like
     length @alpha ([] @alpha)
  It really doesn't matter what type we choose for alpha.  But
  we must choose a type!  We can't leave mutable unification
  variables floating around: after typecheck is complete, every
  type variable occurrence must have a binding site.

  So we default it to 'Any' of the right kind.

  All this works for both type and kind variables (indeed
  the two are the same thing).

* SkolemiseFlexi: is a special case for the LHS of RULES.
  See Note [Zonking the LHS of a RULE]

* RuntimeUnkFlexi: is a special case for the GHCi debugger.
  It's a way to have a variable that is not a mutable
  unification variable, but doesn't have a binding site
  either.

* NoFlexi: See Note [Error on unconstrained meta-variables]
  in GHC.Tc.Utils.TcMType. This mode will panic on unfilled
  meta-variables.
-}

-- | A reader monad over 'ZonkEnv', for zonking computations which
-- don't modify the 'ZonkEnv' (e.g. don't bind any variables).
--
-- Use 'ZonkBndrT' when you need to modify the 'ZonkEnv' (e.g. to bind
-- a variable).
newtype ZonkT m a = ZonkT' { runZonkT :: ZonkEnv -> m a }

{- Note [Instances for ZonkT]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Below, we derive the following instances by hand:

  newtype ZonkT m a = ZonkT { runZonkT :: ZonkEnv -> m a }
    deriving (Functor, Applicative, Monad, MonadIO, MonadFix)
      via ReaderT ZonkEnv m
    deriving MonadTrans
      via ReaderT ZonkEnv

Why? Two reasons:

  1. To use oneShot. See Note [The one-shot state monad trick] in GHC.Utils.Monad.
  2. To be strict in the ZonkEnv. This allows us to worker-wrapper functions,
     passing them individual fields of the ZonkEnv instead of the whole record.
     When this happens, we avoid allocating a ZonkEnv, which is a win.
-}

-- See Note [The one-shot state monad trick] in GHC.Utils.Monad
{-# COMPLETE ZonkT #-}
pattern ZonkT :: forall m a. (ZonkEnv -> m a) -> ZonkT m a
pattern ZonkT m <- ZonkT' m
  where
    ZonkT m = ZonkT' (oneShot m)

-- See Note [Instances for ZonkT]
instance Functor m => Functor (ZonkT m) where
  fmap f (ZonkT g) = ZonkT $ \ !env -> fmap f (g env)
  a <$ ZonkT g     = ZonkT $ \ !env -> a <$ g env
  {-# INLINE fmap #-}
  {-# INLINE (<$) #-}

-- See Note [Instances for ZonkT]
instance Applicative m => Applicative (ZonkT m) where
  pure a = ZonkT (\ !_ -> pure a)
  ZonkT f <*> ZonkT x = ZonkT (\ !env -> f env <*> x env )
  ZonkT m *> f = ZonkT (\ !env -> m env *> runZonkT f env)
  {-# INLINE pure #-}
  {-# INLINE (<*>) #-}
  {-# INLINE (*>) #-}

-- See Note [Instances for ZonkT]
instance Monad m => Monad (ZonkT m) where
  ZonkT m >>= f =
    ZonkT (\ !env -> do { r <- m env
                        ; runZonkT (f r) env })
  (>>)   = (*>)
  {-# INLINE (>>=) #-}
  {-# INLINE (>>) #-}

-- See Note [Instances for ZonkT]
instance MonadIO m => MonadIO (ZonkT m) where
  liftIO f = ZonkT (\ !_ -> liftIO f)
  {-# INLINE liftIO #-}

-- See Note [Instances for ZonkT]
instance MonadTrans ZonkT where
  lift ma = ZonkT $ \ !_ -> ma
  {-# INLINE lift #-}

-- See Note [Instances for ZonkT]
instance MonadFix m => MonadFix (ZonkT m) where
  mfix f = ZonkT $ \ !r -> mfix $ oneShot $ \ a -> runZonkT (f a) r
  {-# INLINE mfix #-}

-- | Zonk binders, bringing them into scope in the inner computation.
--
-- Can be thought of as a state monad transformer @StateT ZonkEnv m a@,
-- but written in continuation-passing style.
--
-- See Note [Continuation-passing style for zonking].
newtype ZonkBndrT m a = ZonkBndrT { runZonkBndrT' :: forall r. (a -> ZonkT m r) -> ZonkT m r }
  deriving (Functor, Applicative, Monad, MonadIO, MonadFix)
    via Codensity (ZonkT m)
    -- See GHC.Utils.Monad.Codensity for the instance definitions.
    -- See Note [Continuation-passing style for zonking] for why we use
    -- continuation-passing style instead of a direct state monad.

{- Note [Continuation-passing style for zonking]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
While zonking, we sometimes need to modify the ZonkEnv. For example, when
zonking a binder with zonkTyBndrX, we extend the type variable ZonkEnv.

We could use direct state passing:

  zonkTyBndrX :: ZonkEnv -> TcTyVar -> TcM (ZonkEnv, TyVar)
  zonkTyBndrX ze tv =
    do { tv' <- ... ze tv
       ; let ze' = extendTyZonkEnv ze tv'
       ; return (ze', tv') }

but we can avoid allocating pairs by using continuation-passing style instead,
for example:

  zonkTyBndrX :: (ZonkEnv -> TcTyVar -> TcM r) -> ZonkEnv -> TcM r
  zonkTyBndrX k ze =
    do { tv' <- ... ze tv
       ; let ze' = extendTyZonkEnv ze tv
       ; k ze' tv' }

We thus define:

  newtype ZonkBndrT m a =
    ZonkBndrT { runZonkBndrT :: forall r. (a -> ZonkT m r) -> ZonkT m r }

which is the type of continuation-passing computations over ZonkT m = ReaderT ZonkEnv m.
We thus have:

  zonkTyBndrX :: TcTyVar -> ZonkBndrT TcM TyVar

which expresses the fact that zonkTyBndrX takes in a TcTyVar, returns a TyVar,
modifying the ZonkEnv state in the process. We can build computations out of it
by using runZonkBndrT and nesting. For example, zonking a type synonym:

  zonkTySynRHS :: [TcTyConBinder] -> [TcTyVar] -> ZonkT TcM ([TyConBinder], [TyVar])
  zonkTySynRHS tc_bndrs rhs_tc_ty =
    runZonkBndrT (zonkTyVarBindersX tc_bndrs) $ \ bndrs ->
      do { rhs_ty <- zonkTcTypeToTypeX rhs_tc_ty
         ; return (bndrs, rhs_ty) }

This is known as the codensity transformation, where

  newtype Codensity m a = Codensity { forall r. (a -> m r) -> m r }

expresses continuation-passing computations in the monad m.

Codensity (ReaderT s m) naturally corresponds to StateT s (Codensity m), and
the instances for Codensity reflect that, e.g.

  traverse :: (a -> ZonkBndrT m b) -> t a -> ZonkBndrT m (t b)

naturally behaves like mapAccumLM, accumulating changes to the ZonkEnv as
we go.
-}

-- | Zonk some binders and run the continuation.
--
-- Example:
--
-- > zonk (ForAllTy (Bndr tv vis) body_ty)
-- >  = runZonkBndrT (zonkTyBndrX tv) $ \ tv' ->
-- >    do { body_ty' <- zonkTcTypeToTypeX body_ty
-- >       ; return (ForAllTy (Bndr tv' vis) body_ty') }
--
-- See Note [Continuation-passing style for zonking].
runZonkBndrT :: ZonkBndrT m a -> forall r. (a -> ZonkT m r) -> ZonkT m r
runZonkBndrT (ZonkBndrT k) f = k (oneShot f)
{-# INLINE runZonkBndrT #-}

-- | Embed a computation that doesn't modify the 'ZonkEnv' into 'ZonkBndrT'.
noBinders :: Monad m => ZonkT m a -> ZonkBndrT m a
noBinders z = coerce $ toCodensity z
{-# INLINE noBinders #-}

-- | Run a nested computation that modifies the 'ZonkEnv',
-- without affecting the outer environment.
don'tBind :: Monad m => ZonkBndrT m a -> ZonkT m a
don'tBind (ZonkBndrT k) = fromCodensity (Codensity k)
{-# INLINE don'tBind #-}

initZonkEnv :: MonadIO m => ZonkFlexi -> ZonkT m b -> m b
initZonkEnv flexi thing_inside
  = do { mtv_env_ref <- liftIO $ newIORef emptyVarEnv
       ; let ze = ZonkEnv { ze_flexi = flexi
                          , ze_tv_env = emptyVarEnv
                          , ze_id_env = emptyVarEnv
                          , ze_meta_tv_env = mtv_env_ref }

       ; runZonkT thing_inside ze }
{-# INLINEABLE initZonkEnv #-} -- so it can be specialised

nestZonkEnv :: (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
nestZonkEnv f = ZonkBndrT $ \ k ->
  case k () of
    ZonkT g -> ZonkT (g . f)
{-# INLINE nestZonkEnv #-}

getZonkEnv :: Monad m => ZonkT m ZonkEnv
getZonkEnv = ZonkT return
{-# INLINE getZonkEnv #-}

-- | Extend the knot-tied environment.
extendIdZonkEnvRec :: [Var] -> ZonkBndrT m ()
extendIdZonkEnvRec ids =
  nestZonkEnv $
    \ ze@(ZonkEnv { ze_id_env = id_env }) ->
    -- NB: Don't look at the var to decide which env't to put it in. That
    -- would end up knot-tying all the env'ts.
      ze { ze_id_env = extendVarEnvList id_env [(id,id) | id <- ids] }
  -- Given coercion variables will actually end up here. That's OK though:
  -- coercion variables are never looked up in the knot-tied env't, so zonking
  -- them simply doesn't get optimised. No one gets hurt. An improvement (?)
  -- would be to do SCC analysis in zonkEvBinds and then only knot-tie the
  -- recursive groups. But perhaps the time it takes to do the analysis is
  -- more than the savings.

extendZonkEnv :: [Var] -> ZonkBndrT m ()
extendZonkEnv vars =
  nestZonkEnv $
    \ ze@(ZonkEnv { ze_tv_env = tyco_env, ze_id_env = id_env }) ->
      ze { ze_tv_env = extendVarEnvList tyco_env [(tv,tv) | tv <- tycovars]
         , ze_id_env = extendVarEnvList id_env   [(id,id) | id <- ids] }
  where
    (tycovars, ids) = partition isTyCoVar vars

extendIdZonkEnv :: Var -> ZonkBndrT m ()
extendIdZonkEnv id =
  nestZonkEnv $
    \ ze@(ZonkEnv { ze_id_env = id_env }) ->
      ze { ze_id_env = extendVarEnv id_env id id }

extendTyZonkEnv :: TyVar -> ZonkBndrT m ()
extendTyZonkEnv tv =
  nestZonkEnv $
    \ ze@(ZonkEnv { ze_tv_env = ty_env }) ->
      ze { ze_tv_env = extendVarEnv ty_env tv tv }

setZonkType :: ZonkFlexi -> ZonkT m a -> ZonkT m a
setZonkType flexi (ZonkT f) = ZonkT $ \ ze ->
  f $ ze { ze_flexi = flexi }