content: {-# LANGUAGE UnboxedTuples, PatternSynonyms, DerivingVia #-}
  module GHC.Types.Unique.DSM
    (
    -- * Threading a deterministic supply
      DUniqSupply
    , UniqDSM(UDSM)
    , DUniqResult
    , pattern DUniqResult
  
    -- ** UniqDSM and DUniqSupply operations
    , getUniqueDSM
    , runUniqueDSM
    , takeUniqueFromDSupply
    , initDUniqSupply
  
    -- ** Tag operations
    , newTagDUniqSupply
    , getTagDUniqSupply
  
    -- * A transfomer threading a deterministic supply
    , UniqDSMT(UDSMT)
  
    -- ** UniqDSMT operations
    , runUDSMT
    , withDUS
    , hoistUDSMT
    , liftUDSMT
  
    -- ** Tags
    , setTagUDSMT
  
    -- * Monad class for deterministic supply threading
    , MonadGetUnique(..)
    , MonadUniqDSM(..)
  
    )
    where
  
  import GHC.Exts (oneShot)
  import GHC.Prelude
  import GHC.Word
  import Control.Monad.Fix
  import GHC.Types.Unique
  import qualified GHC.Utils.Monad.State.Strict as Strict
  import qualified GHC.Types.Unique.Supply as USM
  import Control.Monad.IO.Class
  
  {-
  Note [Deterministic Uniques in the CG]
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  GHC produces fully deterministic object code. To achieve this, there is a key
  pass (detRenameCmmGroup) which renames all non-deterministic uniques in
  the Cmm code right after StgToCmm. See Note [Object determinism] for the big
  picture and some details.
  
  The code generation pipeline that processes this renamed, deterministic, Cmm,
  however, may still need to generate new uniques. If we were to resort to the
  non-deterministic unique supply used in the rest of the compiler, our renaming
  efforts would be for naught.
  
  Therefore, after having renamed Cmm deterministically, we must ensure that all
  uniques created by the code generation pipeline use a deterministic source of uniques.
  Most often, this means don't use `UniqSM` in the Cmm passes, use `UniqDSM`:
  
  `UniqDSM` is a pure state monad with an incrementing counter from which we
  source new uniques. Unlike `UniqSM`, there's no way to `split` the supply, but
  it turns out this was rarely really needed for code generation and migrating
  from UniqSM to UniqDSM was easy.
  
  Secondly, the `DUniqSupply` used to run a `UniqDSM` must be threaded through
  all passes to guarantee uniques in different passes are unique amongst them
  altogether.
  Specifically, the same `DUniqSupply` must be threaded through the CG Streaming
  pipeline, starting with Driver.Main calling `StgToCmm.codeGen`, `cmmPipeline`,
  `cmmToRawCmm`, and `codeOutput` in sequence.
  
  To thread resources through the `Stream` abstraction, we use the `UniqDSMT`
  transformer on top of `IO` as the Monad underlying the Stream. `UniqDSMT` will
  thread the `DUniqSupply` through every pass applied to the `Stream`, for every
  element. We use @type CgStream = Stream (UniqDSMT IO)@ for the Stream used in
  code generation which that carries through the deterministic unique supply.
  
  Unlike non-deterministic unique supplies which can be split into supplies using
  different tags, or where a new supply with a new tag can be brought from the
  void, a `DUniqSupply` needs to be sampled iteratively. To use a different tag
  during a specific pass (to more easily identify uniques created in it), the tag
  should be manually set and then reset on the unique supply. There's also the
  auxiliary `setTagUDSMT` which sets the tag for all uniques supplied in the given
  action, and resets it implicitly.
  
  See also Note [Object determinism] in GHC.StgToCmm
  -}
  
  -- See Note [Deterministic Uniques in the CG]
  newtype DUniqSupply = DUS Word64 -- supply uniques iteratively
  type DUniqResult result = (# result, DUniqSupply #)
  
  pattern DUniqResult :: a -> DUniqSupply -> (# a, DUniqSupply #)
  pattern DUniqResult x y = (# x, y #)
  {-# COMPLETE DUniqResult #-}
  
  -- | A monad which just gives the ability to obtain 'Unique's deterministically.
  -- There's no splitting.
  newtype UniqDSM result = UDSM' { unUDSM :: DUniqSupply -> DUniqResult result }
    deriving (Functor, Applicative, Monad) via (Strict.State DUniqSupply)
  
  instance MonadFix UniqDSM where
    mfix m = UDSM (\us0 -> let (r,us1) = runUniqueDSM us0 (m r) in DUniqResult r us1)
  
  -- See Note [The one-shot state monad trick] in GHC.Utils.Monad
  pattern UDSM :: (DUniqSupply -> DUniqResult a) -> UniqDSM a
  pattern UDSM m <- UDSM' m
    where
      UDSM m = UDSM' (oneShot $ \s -> m s)
  {-# COMPLETE UDSM #-}
  
  getUniqueDSM :: UniqDSM Unique
  getUniqueDSM = UDSM (\(DUS us0) -> DUniqResult (mkUniqueGrimily us0) (DUS $ us0+1))
  
  takeUniqueFromDSupply :: DUniqSupply -> (Unique, DUniqSupply)
  takeUniqueFromDSupply d =
    case unUDSM getUniqueDSM d of
      DUniqResult x y -> (x, y)
  
  -- | Initialize a deterministic unique supply with the given Tag and initial unique.
  initDUniqSupply :: Char -> Word64 -> DUniqSupply
  initDUniqSupply c firstUniq =
    let !tag = mkTag c
    in DUS (tag .|. firstUniq)
  
  runUniqueDSM :: DUniqSupply -> UniqDSM a -> (a, DUniqSupply)
  runUniqueDSM ds (UDSM f) =
    case f ds of
      DUniqResult uq us -> (uq, us)
  
  -- | Set the tag of uniques generated from this deterministic unique supply
  newTagDUniqSupply :: Char -> DUniqSupply -> DUniqSupply
  newTagDUniqSupply c (DUS w) = DUS $ getKey $ newTagUnique (mkUniqueGrimily w) c
  
  -- | Get the tag uniques generated from this deterministic unique supply would have
  getTagDUniqSupply :: DUniqSupply -> Char
  getTagDUniqSupply (DUS w) = fst $ unpkUnique (mkUniqueGrimily w)
  
  -- | Get a unique from a monad that can access a unique supply.
  --
  -- Crucially, because 'MonadGetUnique' doesn't allow you to get the
  -- 'UniqSupply' (unlike 'MonadUnique'), an instance such as 'UniqDSM' can use a
  -- deterministic unique supply to return deterministic uniques without allowing
  -- for the 'UniqSupply' to be shared.
  class Monad m => MonadGetUnique m where
    getUniqueM :: m Unique
  
  instance MonadGetUnique UniqDSM where
    getUniqueM = getUniqueDSM
  
  -- non deterministic instance
  instance MonadGetUnique USM.UniqSM where
    getUniqueM = USM.getUniqueM
  
  --------------------------------------------------------------------------------
  -- UniqDSMT
  --------------------------------------------------------------------------------
  
  -- | Transformer version of 'UniqDSM' to use when threading a deterministic
  -- uniq supply over a Monad. Specifically, it is used in the `Stream` of Cmm
  -- decls.
  newtype UniqDSMT m result = UDSMT' (DUniqSupply -> m (result, DUniqSupply))
    deriving (Functor)
  
  -- Similar to GHC.Utils.Monad.State.Strict, using Note [The one-shot state monad trick]
  -- Using the one-shot trick is necessary for performance.
  -- Using transfomer's strict `StateT` regressed some performance tests in 1-2%.
  -- The one-shot trick here fixes those regressions.
  
  pattern UDSMT :: (DUniqSupply -> m (result, DUniqSupply)) -> UniqDSMT m result
  pattern UDSMT m <- UDSMT' m
    where
      UDSMT m = UDSMT' (oneShot $ \s -> m s)
  {-# COMPLETE UDSMT #-}
  
  instance Monad m => Applicative (UniqDSMT m) where
    pure x = UDSMT $ \s -> pure (x, s)
    UDSMT f <*> UDSMT x = UDSMT $ \s0 -> do
      (f', s1) <- f s0
      (x', s2) <- x s1
      pure (f' x', s2)
  
  instance Monad m => Monad (UniqDSMT m) where
    UDSMT x >>= f = UDSMT $ \s0 -> do
      (x', s1) <- x s0
      case f x' of UDSMT y -> y s1
  
  instance MonadIO m => MonadIO (UniqDSMT m) where
    liftIO x = UDSMT $ \s -> (,s) <$> liftIO x
  
  instance Monad m => MonadGetUnique (UniqDSMT m) where
    getUniqueM = UDSMT $ \us -> do
      let (u, us') = takeUniqueFromDSupply us
      return (u, us')
  
  -- | Set the tag of the running @UniqDSMT@ supply to the given tag and run an action with it.
  -- All uniques produced in the given action will use this tag, until the tag is changed
  -- again.
  setTagUDSMT :: Monad m => Char {-^ Tag -} -> UniqDSMT m a -> UniqDSMT m a
  setTagUDSMT tag (UDSMT act) = UDSMT $ \us -> do
    let origtag = getTagDUniqSupply us
        new_us  = newTagDUniqSupply tag us
    (a, us') <- act new_us
    let us'_origtag = newTagDUniqSupply origtag us'
        -- restore original tag
    return (a, us'_origtag)
  
  -- | Like 'runUniqueDSM' but for 'UniqDSMT'
  runUDSMT :: DUniqSupply -> UniqDSMT m a -> m (a, DUniqSupply)
  runUDSMT dus (UDSMT st) = st dus
  
  -- | Lift an IO action that depends on, and threads through, a unique supply
  -- into UniqDSMT IO.
  withDUS :: (DUniqSupply -> IO (a, DUniqSupply)) -> UniqDSMT IO a
  withDUS f = UDSMT $ \us -> do
    (a, us') <- liftIO (f us)
    return (a, us')
  
  -- | Change the monad underyling an applied @UniqDSMT@, i.e. transform a
  -- @UniqDSMT m@ into a @UniqDSMT n@ given @m ~> n@.
  hoistUDSMT :: (forall x. m x -> n x) -> UniqDSMT m a -> UniqDSMT n a
  hoistUDSMT nt (UDSMT m) = UDSMT $ \s -> nt (m s)
  
  -- | Lift a monadic action @m a@ into an @UniqDSMT m a@
  liftUDSMT :: Functor m => m a -> UniqDSMT m a
  liftUDSMT m = UDSMT $ \s -> (,s) <$> m
  
  --------------------------------------------------------------------------------
  -- MonadUniqDSM
  --------------------------------------------------------------------------------
  
  class (Monad m) => MonadUniqDSM m where
    -- | Lift a pure 'UniqDSM' action into a 'MonadUniqDSM' such as 'UniqDSMT'
    liftUniqDSM :: UniqDSM a -> m a
  
  instance MonadUniqDSM UniqDSM where
    liftUniqDSM = id
  
  instance Monad m => MonadUniqDSM (UniqDSMT m) where
    liftUniqDSM act = UDSMT $ \us -> pure $ runUniqueDSM us act
  