    {-# LANGUAGE UndecidableInstances #-}

    -- | This module provides a monadic facility similar (and built on top
    -- of) "Futhark.FreshNames".  The removes the need for a (small) amount of
    -- boilerplate, at the cost of using some GHC extensions.  The idea is
    -- that if your compiler pass runs in a monad that is an instance of
    -- 'MonadFreshNames', you can automatically use the name generation
    -- functions exported by this module.
    module Futhark.MonadFreshNames
      ( MonadFreshNames (..),
        modifyNameSource,
        newName,
        newNameFromString,
        newVName,
        newIdent,
        newIdent',
        newParam,
        module Futhark.FreshNames,
      )
    where

    import Control.Monad.Except
    import Control.Monad.RWS.Lazy qualified
    import Control.Monad.RWS.Strict qualified
    import Control.Monad.Reader
    import Control.Monad.State.Lazy qualified
    import Control.Monad.State.Strict qualified
    import Control.Monad.Trans.Maybe qualified
    import Control.Monad.Writer.Lazy qualified
    import Control.Monad.Writer.Strict qualified
    import Futhark.FreshNames hiding (newName)
    import Futhark.FreshNames qualified as FreshNames
    import Futhark.IR.Syntax

    -- | A monad that stores a name source.  The following is a good
    -- instance for a monad in which the only state is a @NameSource vn@:
    --
    -- @
    --  instance MonadFreshNames vn MyMonad where
    --    getNameSource = get
    --    putNameSource = put
    -- @
    class (Monad m) => MonadFreshNames m where
      getNameSource :: m VNameSource
      putNameSource :: VNameSource -> m ()

    instance (Monad im) => MonadFreshNames (Control.Monad.State.Lazy.StateT VNameSource im) where
      getNameSource :: StateT VNameSource im VNameSource
    getNameSource = StateT VNameSource im VNameSource
    forall s (m :: * -> *). MonadState s m => m s
    Control.Monad.State.Lazy.get
      putNameSource :: VNameSource -> StateT VNameSource im ()
    putNameSource = VNameSource -> StateT VNameSource im ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    Control.Monad.State.Lazy.put

    instance (Monad im) => MonadFreshNames (Control.Monad.State.Strict.StateT VNameSource im) where
      getNameSource :: StateT VNameSource im VNameSource
    getNameSource = StateT VNameSource im VNameSource
    forall s (m :: * -> *). MonadState s m => m s
    Control.Monad.State.Strict.get
      putNameSource :: VNameSource -> StateT VNameSource im ()
    putNameSource = VNameSource -> StateT VNameSource im ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    Control.Monad.State.Strict.put

    instance
      (Monad im, Monoid w) =>
      MonadFreshNames (Control.Monad.RWS.Lazy.RWST r w VNameSource im)
      where
      getNameSource :: RWST r w VNameSource im VNameSource
    getNameSource = RWST r w VNameSource im VNameSource
    forall s (m :: * -> *). MonadState s m => m s
    Control.Monad.RWS.Lazy.get
      putNameSource :: VNameSource -> RWST r w VNameSource im ()
    putNameSource = VNameSource -> RWST r w VNameSource im ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    Control.Monad.RWS.Lazy.put

    instance
      (Monad im, Monoid w) =>
      MonadFreshNames (Control.Monad.RWS.Strict.RWST r w VNameSource im)
      where
      getNameSource :: RWST r w VNameSource im VNameSource
    getNameSource = RWST r w VNameSource im VNameSource
    forall s (m :: * -> *). MonadState s m => m s
    Control.Monad.RWS.Strict.get
      putNameSource :: VNameSource -> RWST r w VNameSource im ()
    putNameSource = VNameSource -> RWST r w VNameSource im ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    Control.Monad.RWS.Strict.put

    -- | Run a computation needing a fresh name source and returning a new
    -- one, using 'getNameSource' and 'putNameSource' before and after the
    -- computation.
    modifyNameSource :: (MonadFreshNames m) => (VNameSource -> (a, VNameSource)) -> m a
    modifyNameSource :: forall (m :: * -> *) a.
    MonadFreshNames m =>
    (VNameSource -> (a, VNameSource)) -> m a
    modifyNameSource VNameSource -> (a, VNameSource)
    m = do
      VNameSource
    src <- m VNameSource
    forall (m :: * -> *). MonadFreshNames m => m VNameSource
    getNameSource
      let (a
    x, VNameSource
    src') = VNameSource -> (a, VNameSource)
    m VNameSource
    src
      VNameSource
    src' VNameSource -> m () -> m ()
    forall a b. a -> b -> b
    `seq` VNameSource -> m ()
    forall (m :: * -> *). MonadFreshNames m => VNameSource -> m ()
    putNameSource VNameSource
    src'
      a -> m a
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure a
    x

    -- | Produce a fresh name, using the given name as a template.
    newName :: (MonadFreshNames m) => VName -> m VName
    newName :: forall (m :: * -> *). MonadFreshNames m => VName -> m VName
    newName = (VNameSource -> (VName, VNameSource)) -> m VName
    forall (m :: * -> *) a.
    MonadFreshNames m =>
    (VNameSource -> (a, VNameSource)) -> m a
    modifyNameSource ((VNameSource -> (VName, VNameSource)) -> m VName)
    -> (VName -> VNameSource -> (VName, VNameSource))
    -> VName
    -> m VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VNameSource -> VName -> (VName, VNameSource))
    -> VName -> VNameSource -> (VName, VNameSource)
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip VNameSource -> VName -> (VName, VNameSource)
    FreshNames.newName

    -- | As @newName@, but takes a 'String' for the name template.
    newNameFromString :: (MonadFreshNames m) => String -> m VName
    newNameFromString :: forall (m :: * -> *). MonadFreshNames m => String -> m VName
    newNameFromString String
    s = VName -> m VName
    forall (m :: * -> *). MonadFreshNames m => VName -> m VName
    newName (VName -> m VName) -> VName -> m VName
    forall a b. (a -> b) -> a -> b
    $ Name -> Int -> VName
    VName (String -> Name
    nameFromString String
    s) Int
    0

    -- | Produce a fresh 'VName', using the given base name as a template.
    newID :: (MonadFreshNames m) => Name -> m VName
    newID :: forall (m :: * -> *). MonadFreshNames m => Name -> m VName
    newID Name
    s = VName -> m VName
    forall (m :: * -> *). MonadFreshNames m => VName -> m VName
    newName (VName -> m VName) -> VName -> m VName
    forall a b. (a -> b) -> a -> b
    $ Name -> Int -> VName
    VName Name
    s Int
    0

    -- | Produce a fresh 'VName', using the given base name as a template.
    newVName :: (MonadFreshNames m) => String -> m VName
    newVName :: forall (m :: * -> *). MonadFreshNames m => String -> m VName
    newVName = Name -> m VName
    forall (m :: * -> *). MonadFreshNames m => Name -> m VName
    newID (Name -> m VName) -> (String -> Name) -> String -> m VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . String -> Name
    nameFromString

    -- | Produce a fresh 'Ident', using the given name as a template.
    newIdent ::
      (MonadFreshNames m) =>
      String ->
      Type ->
      m Ident
    newIdent :: forall (m :: * -> *).
    MonadFreshNames m =>
    String -> Type -> m Ident
    newIdent String
    s Type
    t = do
      VName
    s' <- Name -> m VName
    forall (m :: * -> *). MonadFreshNames m => Name -> m VName
    newID (Name -> m VName) -> Name -> m VName
    forall a b. (a -> b) -> a -> b
    $ String -> Name
    nameFromString String
    s
      Ident -> m Ident
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Ident -> m Ident) -> Ident -> m Ident
    forall a b. (a -> b) -> a -> b
    $ VName -> Type -> Ident
    Ident VName
    s' Type
    t

    -- | Produce a fresh 'Ident', using the given 'Ident' as a template,
    -- but possibly modifying the name.
    newIdent' ::
      (MonadFreshNames m) =>
      (String -> String) ->
      Ident ->
      m Ident
    newIdent' :: forall (m :: * -> *).
    MonadFreshNames m =>
    (String -> String) -> Ident -> m Ident
    newIdent' String -> String
    f Ident
    ident =
      String -> Type -> m Ident
    forall (m :: * -> *).
    MonadFreshNames m =>
    String -> Type -> m Ident
    newIdent
        (String -> String
    f (String -> String) -> String -> String
    forall a b. (a -> b) -> a -> b
    $ Name -> String
    nameToString (Name -> String) -> Name -> String
    forall a b. (a -> b) -> a -> b
    $ VName -> Name
    baseName (VName -> Name) -> VName -> Name
    forall a b. (a -> b) -> a -> b
    $ Ident -> VName
    identName Ident
    ident)
        (Ident -> Type
    identType Ident
    ident)

    -- | Produce a fresh 'Param', using the given name as a template.
    newParam ::
      (MonadFreshNames m) =>
      String ->
      dec ->
      m (Param dec)
    newParam :: forall (m :: * -> *) dec.
    MonadFreshNames m =>
    String -> dec -> m (Param dec)
    newParam String
    s dec
    t = do
      VName
    s' <- Name -> m VName
    forall (m :: * -> *). MonadFreshNames m => Name -> m VName
    newID (Name -> m VName) -> Name -> m VName
    forall a b. (a -> b) -> a -> b
    $ String -> Name
    nameFromString String
    s
      Param dec -> m (Param dec)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Param dec -> m (Param dec)) -> Param dec -> m (Param dec)
    forall a b. (a -> b) -> a -> b
    $ Attrs -> VName -> dec -> Param dec
    forall dec. Attrs -> VName -> dec -> Param dec
    Param Attrs
    forall a. Monoid a => a
    mempty VName
    s' dec
    t

    -- Utility instance defintions for MTL classes.  This requires
    -- UndecidableInstances, but saves on typing elsewhere.

    instance (MonadFreshNames m) => MonadFreshNames (ReaderT s m) where
      getNameSource :: ReaderT s m VNameSource
    getNameSource = m VNameSource -> ReaderT s m VNameSource
    forall (m :: * -> *) a. Monad m => m a -> ReaderT s m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift m VNameSource
    forall (m :: * -> *). MonadFreshNames m => m VNameSource
    getNameSource
      putNameSource :: VNameSource -> ReaderT s m ()
    putNameSource = m () -> ReaderT s m ()
    forall (m :: * -> *) a. Monad m => m a -> ReaderT s m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (m () -> ReaderT s m ())
    -> (VNameSource -> m ()) -> VNameSource -> ReaderT s m ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VNameSource -> m ()
    forall (m :: * -> *). MonadFreshNames m => VNameSource -> m ()
    putNameSource

    instance
      (MonadFreshNames m, Monoid s) =>
      MonadFreshNames (Control.Monad.Writer.Lazy.WriterT s m)
      where
      getNameSource :: WriterT s m VNameSource
    getNameSource = m VNameSource -> WriterT s m VNameSource
    forall (m :: * -> *) a. Monad m => m a -> WriterT s m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift m VNameSource
    forall (m :: * -> *). MonadFreshNames m => m VNameSource
    getNameSource
      putNameSource :: VNameSource -> WriterT s m ()
    putNameSource = m () -> WriterT s m ()
    forall (m :: * -> *) a. Monad m => m a -> WriterT s m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (m () -> WriterT s m ())
    -> (VNameSource -> m ()) -> VNameSource -> WriterT s m ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VNameSource -> m ()
    forall (m :: * -> *). MonadFreshNames m => VNameSource -> m ()
    putNameSource

    instance
      (MonadFreshNames m, Monoid s) =>
      MonadFreshNames (Control.Monad.Writer.Strict.WriterT s m)
      where
      getNameSource :: WriterT s m VNameSource
    getNameSource = m VNameSource -> WriterT s m VNameSource
    forall (m :: * -> *) a. Monad m => m a -> WriterT s m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift m VNameSource
    forall (m :: * -> *). MonadFreshNames m => m VNameSource
    getNameSource
      putNameSource :: VNameSource -> WriterT s m ()
    putNameSource = m () -> WriterT s m ()
    forall (m :: * -> *) a. Monad m => m a -> WriterT s m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (m () -> WriterT s m ())
    -> (VNameSource -> m ()) -> VNameSource -> WriterT s m ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VNameSource -> m ()
    forall (m :: * -> *). MonadFreshNames m => VNameSource -> m ()
    putNameSource

    instance
      (MonadFreshNames m) =>
      MonadFreshNames (Control.Monad.Trans.Maybe.MaybeT m)
      where
      getNameSource :: MaybeT m VNameSource
    getNameSource = m VNameSource -> MaybeT m VNameSource
    forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift m VNameSource
    forall (m :: * -> *). MonadFreshNames m => m VNameSource
    getNameSource
      putNameSource :: VNameSource -> MaybeT m ()
    putNameSource = m () -> MaybeT m ()
    forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (m () -> MaybeT m ())
    -> (VNameSource -> m ()) -> VNameSource -> MaybeT m ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VNameSource -> m ()
    forall (m :: * -> *). MonadFreshNames m => VNameSource -> m ()
    putNameSource

    instance
      (MonadFreshNames m) =>
      MonadFreshNames (ExceptT e m)
      where
      getNameSource :: ExceptT e m VNameSource
    getNameSource = m VNameSource -> ExceptT e m VNameSource
    forall (m :: * -> *) a. Monad m => m a -> ExceptT e m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift m VNameSource
    forall (m :: * -> *). MonadFreshNames m => m VNameSource
    getNameSource
      putNameSource :: VNameSource -> ExceptT e m ()
    putNameSource = m () -> ExceptT e m ()
    forall (m :: * -> *) a. Monad m => m a -> ExceptT e m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (m () -> ExceptT e m ())
    -> (VNameSource -> m ()) -> VNameSource -> ExceptT e m ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VNameSource -> m ()
    forall (m :: * -> *). MonadFreshNames m => VNameSource -> m ()
    putNameSource