content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  module Util.Defer (
    Defer, Embed(..), now, later, immediately, lift
  ) where
  
  import Control.Applicative
  import Control.Monad.Catch
  import Control.Monad.STM
  import Control.Monad.State.Strict
  
  -- | Used to link an STM transaction and deferred IO actions to be run when that
  -- transaction has been completed.
  newtype Defer io stm a = Defer (StateT (io ()) stm a)
    deriving(
      Functor, Applicative, Alternative, Monad,
      MonadTrans, MonadThrow, MonadCatch)
  
  -- | Run a given STM transaction inside a 'Defer'.
  now :: Monad stm => stm a -> Defer io stm a
  now = lift
  
  -- | Used inside 'immediately' actions to append deferred IO actions to run
  -- after the STM transaction has been completed.
  later :: (Applicative io, Monad stm) => io () -> Defer io stm ()
  later io = Defer $ modify' (*> io)
  
  class Embed m1 m2 where
    embed :: m1 a -> m2 a
  
  instance Embed STM IO where
    embed = atomically
  
  -- | Run an STM transaction and then its deferred IO action.
  immediately :: (Monad io, Monad stm, Embed stm io) => Defer io stm a -> io a
  immediately (Defer action) = do
    (x,deferred) <- embed $ runStateT action $ return ()
    deferred
    return x
  