{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Data.Record.Internal.GHC.Fresh (
    MonadFresh(..)
  , runFreshHsc
  ) where

import Data.IORef
import Control.Monad.Reader

import Data.Record.Internal.GHC.Shim

class Monad m => MonadFresh m where
  -- | Construct a fresh name for use in term level expressions
  --
  -- NOTES:
  --
  -- o These names should be used for module exports.
  -- o These names should be used for exactly /one/ binder.
  -- o The resulting name has the same 'NameSpace' as the argument.
  freshName :: LRdrName -> m LRdrName

newtype Fresh a = WrapFresh { Fresh a -> ReaderT (IORef NameCache) IO a
unwrapFresh :: ReaderT (IORef NameCache) IO a }
  deriving newtype (a -> Fresh b -> Fresh a
(a -> b) -> Fresh a -> Fresh b
(forall a b. (a -> b) -> Fresh a -> Fresh b)
-> (forall a b. a -> Fresh b -> Fresh a) -> Functor Fresh
forall a b. a -> Fresh b -> Fresh a
forall a b. (a -> b) -> Fresh a -> Fresh b
forall (f :: Type -> Type).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> Fresh b -> Fresh a
$c<$ :: forall a b. a -> Fresh b -> Fresh a
fmap :: (a -> b) -> Fresh a -> Fresh b
$cfmap :: forall a b. (a -> b) -> Fresh a -> Fresh b
Functor, Functor Fresh
a -> Fresh a
Functor Fresh
-> (forall a. a -> Fresh a)
-> (forall a b. Fresh (a -> b) -> Fresh a -> Fresh b)
-> (forall a b c. (a -> b -> c) -> Fresh a -> Fresh b -> Fresh c)
-> (forall a b. Fresh a -> Fresh b -> Fresh b)
-> (forall a b. Fresh a -> Fresh b -> Fresh a)
-> Applicative Fresh
Fresh a -> Fresh b -> Fresh b
Fresh a -> Fresh b -> Fresh a
Fresh (a -> b) -> Fresh a -> Fresh b
(a -> b -> c) -> Fresh a -> Fresh b -> Fresh c
forall a. a -> Fresh a
forall a b. Fresh a -> Fresh b -> Fresh a
forall a b. Fresh a -> Fresh b -> Fresh b
forall a b. Fresh (a -> b) -> Fresh a -> Fresh b
forall a b c. (a -> b -> c) -> Fresh a -> Fresh b -> Fresh c
forall (f :: Type -> Type).
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
<* :: Fresh a -> Fresh b -> Fresh a
$c<* :: forall a b. Fresh a -> Fresh b -> Fresh a
*> :: Fresh a -> Fresh b -> Fresh b
$c*> :: forall a b. Fresh a -> Fresh b -> Fresh b
liftA2 :: (a -> b -> c) -> Fresh a -> Fresh b -> Fresh c
$cliftA2 :: forall a b c. (a -> b -> c) -> Fresh a -> Fresh b -> Fresh c
<*> :: Fresh (a -> b) -> Fresh a -> Fresh b
$c<*> :: forall a b. Fresh (a -> b) -> Fresh a -> Fresh b
pure :: a -> Fresh a
$cpure :: forall a. a -> Fresh a
$cp1Applicative :: Functor Fresh
Applicative, Applicative Fresh
a -> Fresh a
Applicative Fresh
-> (forall a b. Fresh a -> (a -> Fresh b) -> Fresh b)
-> (forall a b. Fresh a -> Fresh b -> Fresh b)
-> (forall a. a -> Fresh a)
-> Monad Fresh
Fresh a -> (a -> Fresh b) -> Fresh b
Fresh a -> Fresh b -> Fresh b
forall a. a -> Fresh a
forall a b. Fresh a -> Fresh b -> Fresh b
forall a b. Fresh a -> (a -> Fresh b) -> Fresh b
forall (m :: Type -> Type).
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
return :: a -> Fresh a
$creturn :: forall a. a -> Fresh a
>> :: Fresh a -> Fresh b -> Fresh b
$c>> :: forall a b. Fresh a -> Fresh b -> Fresh b
>>= :: Fresh a -> (a -> Fresh b) -> Fresh b
$c>>= :: forall a b. Fresh a -> (a -> Fresh b) -> Fresh b
$cp1Monad :: Applicative Fresh
Monad)

instance MonadFresh Fresh where
  freshName :: LRdrName -> Fresh LRdrName
freshName (L SrcSpan
l RdrName
name) = ReaderT (IORef NameCache) IO LRdrName -> Fresh LRdrName
forall a. ReaderT (IORef NameCache) IO a -> Fresh a
WrapFresh (ReaderT (IORef NameCache) IO LRdrName -> Fresh LRdrName)
-> ReaderT (IORef NameCache) IO LRdrName -> Fresh LRdrName
forall a b. (a -> b) -> a -> b
$ (IORef NameCache -> IO LRdrName)
-> ReaderT (IORef NameCache) IO LRdrName
forall r (m :: Type -> Type) a. (r -> m a) -> ReaderT r m a
ReaderT ((IORef NameCache -> IO LRdrName)
 -> ReaderT (IORef NameCache) IO LRdrName)
-> (IORef NameCache -> IO LRdrName)
-> ReaderT (IORef NameCache) IO LRdrName
forall a b. (a -> b) -> a -> b
$ \IORef NameCache
nc_var ->
      IORef NameCache
-> (NameCache -> (NameCache, LRdrName)) -> IO LRdrName
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef NameCache
nc_var NameCache -> (NameCache, LRdrName)
aux
    where
      aux :: NameCache -> (NameCache, LRdrName)
      aux :: NameCache -> (NameCache, LRdrName)
aux NameCache
nc = (
            NameCache
nc { nsUniqs :: UniqSupply
nsUniqs = UniqSupply
us }
          , SrcSpan -> RdrName -> LRdrName
forall l e. l -> e -> GenLocated l e
L SrcSpan
l (RdrName -> LRdrName) -> RdrName -> LRdrName
forall a b. (a -> b) -> a -> b
$ Name -> RdrName
Exact (Name -> RdrName) -> Name -> RdrName
forall a b. (a -> b) -> a -> b
$
              Unique -> OccName -> SrcSpan -> Name
mkInternalName Unique
newUniq (OccName -> OccName
newOccName (RdrName -> OccName
rdrNameOcc RdrName
name)) SrcSpan
l
          )
        where
          (Unique
newUniq, UniqSupply
us) = UniqSupply -> (Unique, UniqSupply)
takeUniqFromSupply (NameCache -> UniqSupply
nsUniqs NameCache
nc)

      -- Even when we generate fresh names, ghc can still complain about name
      -- shadowing, because this check only considers the 'OccName', not the
      -- unique. We therefore prefix the name with an underscore to avoid the
      -- warning.
      newOccName :: OccName -> OccName
      newOccName :: OccName -> OccName
newOccName OccName
n = NameSpace -> String -> OccName
mkOccName (OccName -> NameSpace
occNameSpace OccName
n) (String -> OccName) -> (String -> String) -> String -> OccName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"_" String -> String -> String
forall a. [a] -> [a] -> [a]
++) (String -> OccName) -> String -> OccName
forall a b. (a -> b) -> a -> b
$ OccName -> String
occNameString OccName
n

runFresh :: Fresh a -> IORef NameCache -> IO a
runFresh :: Fresh a -> IORef NameCache -> IO a
runFresh = ReaderT (IORef NameCache) IO a -> IORef NameCache -> IO a
forall r (m :: Type -> Type) a. ReaderT r m a -> r -> m a
runReaderT (ReaderT (IORef NameCache) IO a -> IORef NameCache -> IO a)
-> (Fresh a -> ReaderT (IORef NameCache) IO a)
-> Fresh a
-> IORef NameCache
-> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fresh a -> ReaderT (IORef NameCache) IO a
forall a. Fresh a -> ReaderT (IORef NameCache) IO a
unwrapFresh

runFreshHsc :: Fresh a -> Hsc a
runFreshHsc :: Fresh a -> Hsc a
runFreshHsc Fresh a
fa = do
    HscEnv
env <- Hsc HscEnv
getHscEnv
    IO a -> Hsc a
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO a -> Hsc a) -> IO a -> Hsc a
forall a b. (a -> b) -> a -> b
$ Fresh a -> IORef NameCache -> IO a
forall a. Fresh a -> IORef NameCache -> IO a
runFresh Fresh a
fa (HscEnv -> IORef NameCache
hsc_NC HscEnv
env)