{-# 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 { ZonkEnv -> ZonkFlexi
ze_flexi       :: !ZonkFlexi
            , ZonkEnv -> TyCoVarEnv TyCoVar
ze_tv_env      :: TyCoVarEnv TyCoVar
            , ZonkEnv -> TyCoVarEnv TyCoVar
ze_id_env      :: IdEnv      Id
            , ZonkEnv -> IORef (TyVarEnv Type)
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 -> SDoc
ppr (ZonkEnv { ze_tv_env :: ZonkEnv -> TyCoVarEnv TyCoVar
ze_tv_env = TyCoVarEnv TyCoVar
tv_env
               , ze_id_env :: ZonkEnv -> TyCoVarEnv TyCoVar
ze_id_env = TyCoVarEnv TyCoVar
id_env })
    = String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"ZE" SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc
braces ([SDoc] -> SDoc
forall doc. IsDoc doc => [doc] -> doc
vcat
         [ String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"ze_tv_env =" SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> TyCoVarEnv TyCoVar -> SDoc
forall a. Outputable a => a -> SDoc
ppr TyCoVarEnv TyCoVar
tv_env
         , String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"ze_id_env =" SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> TyCoVarEnv TyCoVar -> SDoc
forall a. Outputable a => a -> SDoc
ppr TyCoVarEnv TyCoVar
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' { forall (m :: * -> *) a. ZonkT m a -> ZonkEnv -> m a
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 $mZonkT :: forall {r} {m :: * -> *} {a}.
ZonkT m a -> ((ZonkEnv -> m a) -> r) -> ((# #) -> r) -> r
$bZonkT :: forall (m :: * -> *) a. (ZonkEnv -> m a) -> ZonkT m a
ZonkT m <- ZonkT' m
  where
    ZonkT ZonkEnv -> m a
m = (ZonkEnv -> m a) -> ZonkT m a
forall (m :: * -> *) a. (ZonkEnv -> m a) -> ZonkT m a
ZonkT' ((ZonkEnv -> m a) -> ZonkEnv -> m a
forall a b. (a -> b) -> a -> b
oneShot ZonkEnv -> m a
m)

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

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

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

-- See Note [Instances for ZonkT]
instance MonadIO m => MonadIO (ZonkT m) where
  liftIO :: forall a. IO a -> ZonkT m a
liftIO IO a
f = (ZonkEnv -> m a) -> ZonkT m a
forall (m :: * -> *) a. (ZonkEnv -> m a) -> ZonkT m a
ZonkT (\ !ZonkEnv
_ -> IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO a
f)
  {-# INLINE liftIO #-}

-- See Note [Instances for ZonkT]
instance MonadTrans ZonkT where
  lift :: forall (m :: * -> *) a. Monad m => m a -> ZonkT m a
lift m a
ma = (ZonkEnv -> m a) -> ZonkT m a
forall (m :: * -> *) a. (ZonkEnv -> m a) -> ZonkT m a
ZonkT ((ZonkEnv -> m a) -> ZonkT m a) -> (ZonkEnv -> m a) -> ZonkT m a
forall a b. (a -> b) -> a -> b
$ \ !ZonkEnv
_ -> m a
ma
  {-# INLINE lift #-}

-- See Note [Instances for ZonkT]
instance MonadFix m => MonadFix (ZonkT m) where
  mfix :: forall a. (a -> ZonkT m a) -> ZonkT m a
mfix a -> ZonkT m a
f = (ZonkEnv -> m a) -> ZonkT m a
forall (m :: * -> *) a. (ZonkEnv -> m a) -> ZonkT m a
ZonkT ((ZonkEnv -> m a) -> ZonkT m a) -> (ZonkEnv -> m a) -> ZonkT m a
forall a b. (a -> b) -> a -> b
$ \ !ZonkEnv
r -> (a -> m a) -> m a
forall a. (a -> m a) -> m a
forall (m :: * -> *) a. MonadFix m => (a -> m a) -> m a
mfix ((a -> m a) -> m a) -> (a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ (a -> m a) -> a -> m a
forall a b. (a -> b) -> a -> b
oneShot ((a -> m a) -> a -> m a) -> (a -> m a) -> a -> m a
forall a b. (a -> b) -> a -> b
$ \ a
a -> ZonkT m a -> ZonkEnv -> m a
forall (m :: * -> *) a. ZonkT m a -> ZonkEnv -> m a
runZonkT (a -> ZonkT m a
f a
a) ZonkEnv
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 { forall (m :: * -> *) a.
ZonkBndrT m a -> forall r. (a -> ZonkT m r) -> ZonkT m r
runZonkBndrT' :: forall r. (a -> ZonkT m r) -> ZonkT m r }
  deriving ((forall a b. (a -> b) -> ZonkBndrT m a -> ZonkBndrT m b)
-> (forall a b. a -> ZonkBndrT m b -> ZonkBndrT m a)
-> Functor (ZonkBndrT m)
forall a b. a -> ZonkBndrT m b -> ZonkBndrT m a
forall a b. (a -> b) -> ZonkBndrT m a -> ZonkBndrT m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (m :: * -> *) a b. a -> ZonkBndrT m b -> ZonkBndrT m a
forall (m :: * -> *) a b.
(a -> b) -> ZonkBndrT m a -> ZonkBndrT m b
$cfmap :: forall (m :: * -> *) a b.
(a -> b) -> ZonkBndrT m a -> ZonkBndrT m b
fmap :: forall a b. (a -> b) -> ZonkBndrT m a -> ZonkBndrT m b
$c<$ :: forall (m :: * -> *) a b. a -> ZonkBndrT m b -> ZonkBndrT m a
<$ :: forall a b. a -> ZonkBndrT m b -> ZonkBndrT m a
Functor, Functor (ZonkBndrT m)
Functor (ZonkBndrT m) =>
(forall a. a -> ZonkBndrT m a)
-> (forall a b.
    ZonkBndrT m (a -> b) -> ZonkBndrT m a -> ZonkBndrT m b)
-> (forall a b c.
    (a -> b -> c) -> ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m c)
-> (forall a b. ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m b)
-> (forall a b. ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m a)
-> Applicative (ZonkBndrT m)
forall a. a -> ZonkBndrT m a
forall a b. ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m a
forall a b. ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m b
forall a b. ZonkBndrT m (a -> b) -> ZonkBndrT m a -> ZonkBndrT m b
forall a b c.
(a -> b -> c) -> ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m c
forall (m :: * -> *). Functor (ZonkBndrT m)
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall (m :: * -> *) a. a -> ZonkBndrT m a
forall (m :: * -> *) a b.
ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m a
forall (m :: * -> *) a b.
ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m b
forall (m :: * -> *) a b.
ZonkBndrT m (a -> b) -> ZonkBndrT m a -> ZonkBndrT m b
forall (m :: * -> *) a b c.
(a -> b -> c) -> ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m c
$cpure :: forall (m :: * -> *) a. a -> ZonkBndrT m a
pure :: forall a. a -> ZonkBndrT m a
$c<*> :: forall (m :: * -> *) a b.
ZonkBndrT m (a -> b) -> ZonkBndrT m a -> ZonkBndrT m b
<*> :: forall a b. ZonkBndrT m (a -> b) -> ZonkBndrT m a -> ZonkBndrT m b
$cliftA2 :: forall (m :: * -> *) a b c.
(a -> b -> c) -> ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m c
liftA2 :: forall a b c.
(a -> b -> c) -> ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m c
$c*> :: forall (m :: * -> *) a b.
ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m b
*> :: forall a b. ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m b
$c<* :: forall (m :: * -> *) a b.
ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m a
<* :: forall a b. ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m a
Applicative, Applicative (ZonkBndrT m)
Applicative (ZonkBndrT m) =>
(forall a b.
 ZonkBndrT m a -> (a -> ZonkBndrT m b) -> ZonkBndrT m b)
-> (forall a b. ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m b)
-> (forall a. a -> ZonkBndrT m a)
-> Monad (ZonkBndrT m)
forall a. a -> ZonkBndrT m a
forall a b. ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m b
forall a b. ZonkBndrT m a -> (a -> ZonkBndrT m b) -> ZonkBndrT m b
forall (m :: * -> *). Applicative (ZonkBndrT m)
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
forall (m :: * -> *) a. a -> ZonkBndrT m a
forall (m :: * -> *) a b.
ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m b
forall (m :: * -> *) a b.
ZonkBndrT m a -> (a -> ZonkBndrT m b) -> ZonkBndrT m b
$c>>= :: forall (m :: * -> *) a b.
ZonkBndrT m a -> (a -> ZonkBndrT m b) -> ZonkBndrT m b
>>= :: forall a b. ZonkBndrT m a -> (a -> ZonkBndrT m b) -> ZonkBndrT m b
$c>> :: forall (m :: * -> *) a b.
ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m b
>> :: forall a b. ZonkBndrT m a -> ZonkBndrT m b -> ZonkBndrT m b
$creturn :: forall (m :: * -> *) a. a -> ZonkBndrT m a
return :: forall a. a -> ZonkBndrT m a
Monad, Monad (ZonkBndrT m)
Monad (ZonkBndrT m) =>
(forall a. IO a -> ZonkBndrT m a) -> MonadIO (ZonkBndrT m)
forall a. IO a -> ZonkBndrT m a
forall (m :: * -> *).
Monad m =>
(forall a. IO a -> m a) -> MonadIO m
forall (m :: * -> *). MonadIO m => Monad (ZonkBndrT m)
forall (m :: * -> *) a. MonadIO m => IO a -> ZonkBndrT m a
$cliftIO :: forall (m :: * -> *) a. MonadIO m => IO a -> ZonkBndrT m a
liftIO :: forall a. IO a -> ZonkBndrT m a
MonadIO, Monad (ZonkBndrT m)
Monad (ZonkBndrT m) =>
(forall a. (a -> ZonkBndrT m a) -> ZonkBndrT m a)
-> MonadFix (ZonkBndrT m)
forall a. (a -> ZonkBndrT m a) -> ZonkBndrT m a
forall (m :: * -> *).
Monad m =>
(forall a. (a -> m a) -> m a) -> MonadFix m
forall (m :: * -> *). MonadIO m => Monad (ZonkBndrT m)
forall (m :: * -> *) a.
MonadIO m =>
(a -> ZonkBndrT m a) -> ZonkBndrT m a
$cmfix :: forall (m :: * -> *) a.
MonadIO m =>
(a -> ZonkBndrT m a) -> ZonkBndrT m a
mfix :: forall a. (a -> ZonkBndrT m a) -> ZonkBndrT m a
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 :: forall (m :: * -> *) a.
ZonkBndrT m a -> forall r. (a -> ZonkT m r) -> ZonkT m r
runZonkBndrT (ZonkBndrT forall r. (a -> ZonkT m r) -> ZonkT m r
k) a -> ZonkT m r
f = (a -> ZonkT m r) -> ZonkT m r
forall r. (a -> ZonkT m r) -> ZonkT m r
k ((a -> ZonkT m r) -> a -> ZonkT m r
forall a b. (a -> b) -> a -> b
oneShot a -> ZonkT m r
f)
{-# INLINE runZonkBndrT #-}

-- | Embed a computation that doesn't modify the 'ZonkEnv' into 'ZonkBndrT'.
noBinders :: Monad m => ZonkT m a -> ZonkBndrT m a
noBinders :: forall (m :: * -> *) a. Monad m => ZonkT m a -> ZonkBndrT m a
noBinders ZonkT m a
z = Codensity (ZonkT m) a -> ZonkBndrT m a
forall a b. Coercible a b => a -> b
coerce (Codensity (ZonkT m) a -> ZonkBndrT m a)
-> Codensity (ZonkT m) a -> ZonkBndrT m a
forall a b. (a -> b) -> a -> b
$ ZonkT m a -> Codensity (ZonkT m) a
forall (m :: * -> *) a. Monad m => m a -> Codensity m a
toCodensity ZonkT m a
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 :: forall (m :: * -> *) a. Monad m => ZonkBndrT m a -> ZonkT m a
don'tBind (ZonkBndrT forall r. (a -> ZonkT m r) -> ZonkT m r
k) = Codensity (ZonkT m) a -> ZonkT m a
forall (m :: * -> *) a. Monad m => Codensity m a -> m a
fromCodensity ((forall r. (a -> ZonkT m r) -> ZonkT m r) -> Codensity (ZonkT m) a
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> Codensity m a
Codensity (a -> ZonkT m r) -> ZonkT m r
forall r. (a -> ZonkT m r) -> ZonkT m r
k)
{-# INLINE don'tBind #-}

initZonkEnv :: MonadIO m => ZonkFlexi -> ZonkT m b -> m b
initZonkEnv :: forall (m :: * -> *) b. MonadIO m => ZonkFlexi -> ZonkT m b -> m b
initZonkEnv ZonkFlexi
flexi ZonkT m b
thing_inside
  = do { IORef (TyVarEnv Type)
mtv_env_ref <- IO (IORef (TyVarEnv Type)) -> m (IORef (TyVarEnv Type))
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef (TyVarEnv Type)) -> m (IORef (TyVarEnv Type)))
-> IO (IORef (TyVarEnv Type)) -> m (IORef (TyVarEnv Type))
forall a b. (a -> b) -> a -> b
$ TyVarEnv Type -> IO (IORef (TyVarEnv Type))
forall a. a -> IO (IORef a)
newIORef TyVarEnv Type
forall a. VarEnv a
emptyVarEnv
       ; let ze :: ZonkEnv
ze = ZonkEnv { ze_flexi :: ZonkFlexi
ze_flexi = ZonkFlexi
flexi
                          , ze_tv_env :: TyCoVarEnv TyCoVar
ze_tv_env = TyCoVarEnv TyCoVar
forall a. VarEnv a
emptyVarEnv
                          , ze_id_env :: TyCoVarEnv TyCoVar
ze_id_env = TyCoVarEnv TyCoVar
forall a. VarEnv a
emptyVarEnv
                          , ze_meta_tv_env :: IORef (TyVarEnv Type)
ze_meta_tv_env = IORef (TyVarEnv Type)
mtv_env_ref }

       ; ZonkT m b -> ZonkEnv -> m b
forall (m :: * -> *) a. ZonkT m a -> ZonkEnv -> m a
runZonkT ZonkT m b
thing_inside ZonkEnv
ze }
{-# INLINEABLE initZonkEnv #-} -- so it can be specialised

nestZonkEnv :: (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
nestZonkEnv :: forall (m :: * -> *). (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
nestZonkEnv ZonkEnv -> ZonkEnv
f = (forall r. (() -> ZonkT m r) -> ZonkT m r) -> ZonkBndrT m ()
forall (m :: * -> *) a.
(forall r. (a -> ZonkT m r) -> ZonkT m r) -> ZonkBndrT m a
ZonkBndrT ((forall r. (() -> ZonkT m r) -> ZonkT m r) -> ZonkBndrT m ())
-> (forall r. (() -> ZonkT m r) -> ZonkT m r) -> ZonkBndrT m ()
forall a b. (a -> b) -> a -> b
$ \ () -> ZonkT m r
k ->
  case () -> ZonkT m r
k () of
    ZonkT ZonkEnv -> m r
g -> (ZonkEnv -> m r) -> ZonkT m r
forall (m :: * -> *) a. (ZonkEnv -> m a) -> ZonkT m a
ZonkT (ZonkEnv -> m r
g (ZonkEnv -> m r) -> (ZonkEnv -> ZonkEnv) -> ZonkEnv -> m r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ZonkEnv -> ZonkEnv
f)
{-# INLINE nestZonkEnv #-}

getZonkEnv :: Monad m => ZonkT m ZonkEnv
getZonkEnv :: forall (m :: * -> *). Monad m => ZonkT m ZonkEnv
getZonkEnv = (ZonkEnv -> m ZonkEnv) -> ZonkT m ZonkEnv
forall (m :: * -> *) a. (ZonkEnv -> m a) -> ZonkT m a
ZonkT ZonkEnv -> m ZonkEnv
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
{-# INLINE getZonkEnv #-}

-- | Extend the knot-tied environment.
extendIdZonkEnvRec :: [Var] -> ZonkBndrT m ()
extendIdZonkEnvRec :: forall (m :: * -> *). [TyCoVar] -> ZonkBndrT m ()
extendIdZonkEnvRec [TyCoVar]
ids =
  (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
forall (m :: * -> *). (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
nestZonkEnv ((ZonkEnv -> ZonkEnv) -> ZonkBndrT m ())
-> (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
forall a b. (a -> b) -> a -> b
$
    \ ze :: ZonkEnv
ze@(ZonkEnv { ze_id_env :: ZonkEnv -> TyCoVarEnv TyCoVar
ze_id_env = TyCoVarEnv TyCoVar
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.
      ZonkEnv
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 :: forall (m :: * -> *). [TyCoVar] -> ZonkBndrT m ()
extendZonkEnv [TyCoVar]
vars =
  (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
forall (m :: * -> *). (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
nestZonkEnv ((ZonkEnv -> ZonkEnv) -> ZonkBndrT m ())
-> (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
forall a b. (a -> b) -> a -> b
$
    \ ze :: ZonkEnv
ze@(ZonkEnv { ze_tv_env :: ZonkEnv -> TyCoVarEnv TyCoVar
ze_tv_env = TyCoVarEnv TyCoVar
tyco_env, ze_id_env :: ZonkEnv -> TyCoVarEnv TyCoVar
ze_id_env = TyCoVarEnv TyCoVar
id_env }) ->
      ZonkEnv
ze { ze_tv_env = extendVarEnvList tyco_env [(tv,tv) | tv <- tycovars]
         , ze_id_env = extendVarEnvList id_env   [(id,id) | id <- ids] }
  where
    ([TyCoVar]
tycovars, [TyCoVar]
ids) = (TyCoVar -> Bool) -> [TyCoVar] -> ([TyCoVar], [TyCoVar])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition TyCoVar -> Bool
isTyCoVar [TyCoVar]
vars

extendIdZonkEnv :: Var -> ZonkBndrT m ()
extendIdZonkEnv :: forall (m :: * -> *). TyCoVar -> ZonkBndrT m ()
extendIdZonkEnv TyCoVar
id =
  (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
forall (m :: * -> *). (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
nestZonkEnv ((ZonkEnv -> ZonkEnv) -> ZonkBndrT m ())
-> (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
forall a b. (a -> b) -> a -> b
$
    \ ze :: ZonkEnv
ze@(ZonkEnv { ze_id_env :: ZonkEnv -> TyCoVarEnv TyCoVar
ze_id_env = TyCoVarEnv TyCoVar
id_env }) ->
      ZonkEnv
ze { ze_id_env = extendVarEnv id_env id id }

extendTyZonkEnv :: TyVar -> ZonkBndrT m ()
extendTyZonkEnv :: forall (m :: * -> *). TyCoVar -> ZonkBndrT m ()
extendTyZonkEnv TyCoVar
tv =
  (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
forall (m :: * -> *). (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
nestZonkEnv ((ZonkEnv -> ZonkEnv) -> ZonkBndrT m ())
-> (ZonkEnv -> ZonkEnv) -> ZonkBndrT m ()
forall a b. (a -> b) -> a -> b
$
    \ ze :: ZonkEnv
ze@(ZonkEnv { ze_tv_env :: ZonkEnv -> TyCoVarEnv TyCoVar
ze_tv_env = TyCoVarEnv TyCoVar
ty_env }) ->
      ZonkEnv
ze { ze_tv_env = extendVarEnv ty_env tv tv }

setZonkType :: ZonkFlexi -> ZonkT m a -> ZonkT m a
setZonkType :: forall (m :: * -> *) a. ZonkFlexi -> ZonkT m a -> ZonkT m a
setZonkType ZonkFlexi
flexi (ZonkT ZonkEnv -> m a
f) = (ZonkEnv -> m a) -> ZonkT m a
forall (m :: * -> *) a. (ZonkEnv -> m a) -> ZonkT m a
ZonkT ((ZonkEnv -> m a) -> ZonkT m a) -> (ZonkEnv -> m a) -> ZonkT m a
forall a b. (a -> b) -> a -> b
$ \ ZonkEnv
ze ->
  ZonkEnv -> m a
f (ZonkEnv -> m a) -> ZonkEnv -> m a
forall a b. (a -> b) -> a -> b
$ ZonkEnv
ze { ze_flexi = flexi }