{-# LANGUAGE GeneralizedNewtypeDeriving #-}

-- | Naming things is hard
module Data.Record.Anon.Internal.Plugin.Source.NamingT (
    -- * Monad definition
    NamingT -- opaque
  , runNamingT
  , runNamingHsc
    -- * Key features of NamingT
  , useName
  , fresh
    -- * Convenience derived features
  , freshVar
  ) where

import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor
import Data.IORef
import Data.Set (Set)

import qualified Data.Set as Set

import Data.Record.Anon.Internal.Plugin.Source.GhcShim

{-------------------------------------------------------------------------------
  Monad definition
-------------------------------------------------------------------------------}

data Env = Env {
      Env -> IORef NameCache
envNameCache :: IORef NameCache
    }

-- | Naming things is hard
--
-- The 'NamingT' monad transformer that provides two things:
--
-- 1. Keep track of the imports we need for the names that we use.
-- 2. Generation of fresh names.
newtype NamingT m a = WrapNamingT {
      NamingT m a -> StateT (Set ModuleName) (ReaderT Env m) a
unwrapNamingT :: StateT (Set ModuleName) (ReaderT Env m) a
    }
  deriving (a -> NamingT m b -> NamingT m a
(a -> b) -> NamingT m a -> NamingT m b
(forall a b. (a -> b) -> NamingT m a -> NamingT m b)
-> (forall a b. a -> NamingT m b -> NamingT m a)
-> Functor (NamingT m)
forall a b. a -> NamingT m b -> NamingT m a
forall a b. (a -> b) -> NamingT m a -> NamingT m b
forall (m :: * -> *) a b.
Functor m =>
a -> NamingT m b -> NamingT m a
forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> NamingT m a -> NamingT m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> NamingT m b -> NamingT m a
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> NamingT m b -> NamingT m a
fmap :: (a -> b) -> NamingT m a -> NamingT m b
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> NamingT m a -> NamingT m b
Functor, Functor (NamingT m)
a -> NamingT m a
Functor (NamingT m)
-> (forall a. a -> NamingT m a)
-> (forall a b. NamingT m (a -> b) -> NamingT m a -> NamingT m b)
-> (forall a b c.
    (a -> b -> c) -> NamingT m a -> NamingT m b -> NamingT m c)
-> (forall a b. NamingT m a -> NamingT m b -> NamingT m b)
-> (forall a b. NamingT m a -> NamingT m b -> NamingT m a)
-> Applicative (NamingT m)
NamingT m a -> NamingT m b -> NamingT m b
NamingT m a -> NamingT m b -> NamingT m a
NamingT m (a -> b) -> NamingT m a -> NamingT m b
(a -> b -> c) -> NamingT m a -> NamingT m b -> NamingT m c
forall a. a -> NamingT m a
forall a b. NamingT m a -> NamingT m b -> NamingT m a
forall a b. NamingT m a -> NamingT m b -> NamingT m b
forall a b. NamingT m (a -> b) -> NamingT m a -> NamingT m b
forall a b c.
(a -> b -> c) -> NamingT m a -> NamingT m b -> NamingT m c
forall (m :: * -> *). Monad m => Functor (NamingT m)
forall (m :: * -> *) a. Monad m => a -> NamingT m a
forall (m :: * -> *) a b.
Monad m =>
NamingT m a -> NamingT m b -> NamingT m a
forall (m :: * -> *) a b.
Monad m =>
NamingT m a -> NamingT m b -> NamingT m b
forall (m :: * -> *) a b.
Monad m =>
NamingT m (a -> b) -> NamingT m a -> NamingT m b
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> NamingT m a -> NamingT m b -> NamingT m c
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
<* :: NamingT m a -> NamingT m b -> NamingT m a
$c<* :: forall (m :: * -> *) a b.
Monad m =>
NamingT m a -> NamingT m b -> NamingT m a
*> :: NamingT m a -> NamingT m b -> NamingT m b
$c*> :: forall (m :: * -> *) a b.
Monad m =>
NamingT m a -> NamingT m b -> NamingT m b
liftA2 :: (a -> b -> c) -> NamingT m a -> NamingT m b -> NamingT m c
$cliftA2 :: forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> NamingT m a -> NamingT m b -> NamingT m c
<*> :: NamingT m (a -> b) -> NamingT m a -> NamingT m b
$c<*> :: forall (m :: * -> *) a b.
Monad m =>
NamingT m (a -> b) -> NamingT m a -> NamingT m b
pure :: a -> NamingT m a
$cpure :: forall (m :: * -> *) a. Monad m => a -> NamingT m a
$cp1Applicative :: forall (m :: * -> *). Monad m => Functor (NamingT m)
Applicative, Applicative (NamingT m)
a -> NamingT m a
Applicative (NamingT m)
-> (forall a b. NamingT m a -> (a -> NamingT m b) -> NamingT m b)
-> (forall a b. NamingT m a -> NamingT m b -> NamingT m b)
-> (forall a. a -> NamingT m a)
-> Monad (NamingT m)
NamingT m a -> (a -> NamingT m b) -> NamingT m b
NamingT m a -> NamingT m b -> NamingT m b
forall a. a -> NamingT m a
forall a b. NamingT m a -> NamingT m b -> NamingT m b
forall a b. NamingT m a -> (a -> NamingT m b) -> NamingT m b
forall (m :: * -> *). Monad m => Applicative (NamingT m)
forall (m :: * -> *) a. Monad m => a -> NamingT m a
forall (m :: * -> *) a b.
Monad m =>
NamingT m a -> NamingT m b -> NamingT m b
forall (m :: * -> *) a b.
Monad m =>
NamingT m a -> (a -> NamingT m b) -> NamingT m b
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
return :: a -> NamingT m a
$creturn :: forall (m :: * -> *) a. Monad m => a -> NamingT m a
>> :: NamingT m a -> NamingT m b -> NamingT m b
$c>> :: forall (m :: * -> *) a b.
Monad m =>
NamingT m a -> NamingT m b -> NamingT m b
>>= :: NamingT m a -> (a -> NamingT m b) -> NamingT m b
$c>>= :: forall (m :: * -> *) a b.
Monad m =>
NamingT m a -> (a -> NamingT m b) -> NamingT m b
$cp1Monad :: forall (m :: * -> *). Monad m => Applicative (NamingT m)
Monad)

instance MonadTrans NamingT where
  lift :: m a -> NamingT m a
lift = StateT (Set ModuleName) (ReaderT Env m) a -> NamingT m a
forall (m :: * -> *) a.
StateT (Set ModuleName) (ReaderT Env m) a -> NamingT m a
WrapNamingT (StateT (Set ModuleName) (ReaderT Env m) a -> NamingT m a)
-> (m a -> StateT (Set ModuleName) (ReaderT Env m) a)
-> m a
-> NamingT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ReaderT Env m a -> StateT (Set ModuleName) (ReaderT Env m) a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ReaderT Env m a -> StateT (Set ModuleName) (ReaderT Env m) a)
-> (m a -> ReaderT Env m a)
-> m a
-> StateT (Set ModuleName) (ReaderT Env m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> ReaderT Env m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift

runNamingT :: Functor m => IORef NameCache -> NamingT m a -> m (a, [ModuleName])
runNamingT :: IORef NameCache -> NamingT m a -> m (a, [ModuleName])
runNamingT IORef NameCache
ncVar =
      ((a, Set ModuleName) -> (a, [ModuleName]))
-> m (a, Set ModuleName) -> m (a, [ModuleName])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Set ModuleName -> [ModuleName])
-> (a, Set ModuleName) -> (a, [ModuleName])
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Set ModuleName -> [ModuleName]
forall a. Set a -> [a]
Set.toList)
    (m (a, Set ModuleName) -> m (a, [ModuleName]))
-> (NamingT m a -> m (a, Set ModuleName))
-> NamingT m a
-> m (a, [ModuleName])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ReaderT Env m (a, Set ModuleName) -> Env -> m (a, Set ModuleName))
-> Env
-> ReaderT Env m (a, Set ModuleName)
-> m (a, Set ModuleName)
forall a b c. (a -> b -> c) -> b -> a -> c
flip ReaderT Env m (a, Set ModuleName) -> Env -> m (a, Set ModuleName)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT Env
env
    (ReaderT Env m (a, Set ModuleName) -> m (a, Set ModuleName))
-> (NamingT m a -> ReaderT Env m (a, Set ModuleName))
-> NamingT m a
-> m (a, Set ModuleName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StateT (Set ModuleName) (ReaderT Env m) a
 -> Set ModuleName -> ReaderT Env m (a, Set ModuleName))
-> Set ModuleName
-> StateT (Set ModuleName) (ReaderT Env m) a
-> ReaderT Env m (a, Set ModuleName)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT (Set ModuleName) (ReaderT Env m) a
-> Set ModuleName -> ReaderT Env m (a, Set ModuleName)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT  Set ModuleName
forall a. Set a
Set.empty
    (StateT (Set ModuleName) (ReaderT Env m) a
 -> ReaderT Env m (a, Set ModuleName))
-> (NamingT m a -> StateT (Set ModuleName) (ReaderT Env m) a)
-> NamingT m a
-> ReaderT Env m (a, Set ModuleName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NamingT m a -> StateT (Set ModuleName) (ReaderT Env m) a
forall (m :: * -> *) a.
NamingT m a -> StateT (Set ModuleName) (ReaderT Env m) a
unwrapNamingT
  where
    env :: Env
    env :: Env
env = Env :: IORef NameCache -> Env
Env { envNameCache :: IORef NameCache
envNameCache = IORef NameCache
ncVar }

runNamingHsc :: NamingT Hsc a -> Hsc (a, [ModuleName])
runNamingHsc :: NamingT Hsc a -> Hsc (a, [ModuleName])
runNamingHsc NamingT Hsc a
ma = do
    HscEnv
env <- Hsc HscEnv
getHscEnv
    IORef NameCache -> NamingT Hsc a -> Hsc (a, [ModuleName])
forall (m :: * -> *) a.
Functor m =>
IORef NameCache -> NamingT m a -> m (a, [ModuleName])
runNamingT (HscEnv -> IORef NameCache
hsc_NC HscEnv
env) NamingT Hsc a
ma

{-------------------------------------------------------------------------------
  Key features of NamingT
-------------------------------------------------------------------------------}

useName :: Monad m => RdrName -> NamingT m ()
useName :: RdrName -> NamingT m ()
useName (Qual ModuleName
modl OccName
_) = StateT (Set ModuleName) (ReaderT Env m) () -> NamingT m ()
forall (m :: * -> *) a.
StateT (Set ModuleName) (ReaderT Env m) a -> NamingT m a
WrapNamingT (StateT (Set ModuleName) (ReaderT Env m) () -> NamingT m ())
-> StateT (Set ModuleName) (ReaderT Env m) () -> NamingT m ()
forall a b. (a -> b) -> a -> b
$ (Set ModuleName -> Set ModuleName)
-> StateT (Set ModuleName) (ReaderT Env m) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (ModuleName -> Set ModuleName -> Set ModuleName
forall a. Ord a => a -> Set a -> Set a
Set.insert ModuleName
modl)
useName RdrName
_otherwise    = [Char] -> NamingT m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"useName: expected qualified name"

fresh :: MonadIO m => SrcSpan -> RdrName -> NamingT m RdrName
fresh :: SrcSpan -> RdrName -> NamingT m RdrName
fresh SrcSpan
l RdrName
name = StateT (Set ModuleName) (ReaderT Env m) RdrName
-> NamingT m RdrName
forall (m :: * -> *) a.
StateT (Set ModuleName) (ReaderT Env m) a -> NamingT m a
WrapNamingT (StateT (Set ModuleName) (ReaderT Env m) RdrName
 -> NamingT m RdrName)
-> StateT (Set ModuleName) (ReaderT Env m) RdrName
-> NamingT m RdrName
forall a b. (a -> b) -> a -> b
$ do
    IORef NameCache
ncVar <- (Env -> IORef NameCache)
-> StateT (Set ModuleName) (ReaderT Env m) (IORef NameCache)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> IORef NameCache
envNameCache
    IO RdrName -> StateT (Set ModuleName) (ReaderT Env m) RdrName
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO RdrName -> StateT (Set ModuleName) (ReaderT Env m) RdrName)
-> IO RdrName -> StateT (Set ModuleName) (ReaderT Env m) RdrName
forall a b. (a -> b) -> a -> b
$ IORef NameCache
-> (NameCache -> (NameCache, RdrName)) -> IO RdrName
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef NameCache
ncVar NameCache -> (NameCache, RdrName)
aux
  where
    aux :: NameCache -> (NameCache, RdrName)
    aux :: NameCache -> (NameCache, RdrName)
aux NameCache
nc = (
          NameCache
nc { nsUniqs :: UniqSupply
nsUniqs = UniqSupply
us }
        , 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 -> [Char] -> OccName
mkOccName (OccName -> NameSpace
occNameSpace OccName
n) ([Char] -> OccName) -> ([Char] -> [Char]) -> [Char] -> OccName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Char]
"_" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++) ([Char] -> OccName) -> [Char] -> OccName
forall a b. (a -> b) -> a -> b
$ OccName -> [Char]
occNameString OccName
n

{-------------------------------------------------------------------------------
  Derived convenience functions
-------------------------------------------------------------------------------}

freshVar :: MonadIO m => SrcSpan -> String -> NamingT m RdrName
freshVar :: SrcSpan -> [Char] -> NamingT m RdrName
freshVar SrcSpan
l = SrcSpan -> RdrName -> NamingT m RdrName
forall (m :: * -> *).
MonadIO m =>
SrcSpan -> RdrName -> NamingT m RdrName
fresh SrcSpan
l (RdrName -> NamingT m RdrName)
-> ([Char] -> RdrName) -> [Char] -> NamingT m RdrName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OccName -> RdrName
mkRdrUnqual (OccName -> RdrName) -> ([Char] -> OccName) -> [Char] -> RdrName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> OccName
mkVarOcc