{-# 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 -> 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