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.Reader
    ( catchR, catchR_
    , bracketR, bracketR_
    , withRetry
    ) where
  
  import Control.Exception
  import Control.Monad.Trans.Class
  import Control.Monad.Trans.Reader
  
  -- | Same as 'catch' from 'Control.Exception', but lifted into a ReaderT
  catchR
    :: Exception e
    => ReaderT r IO a
    -> (e -> ReaderT r IO a)
    -> ReaderT r IO a
  catchR (ReaderT action) handler = ReaderT $ \env ->
    action env `catch` \e -> runReaderT (handler e) env
  
  -- | Same as 'catchR', but handler doesn't take an argument
  catchR_
    :: ReaderT r IO a
    -> ReaderT r IO a
    -> ReaderT r IO a
  catchR_ action handler = action `catchR` \SomeException{} -> handler
  
  -- | Same as 'bracket' from 'Control.Exception', but lifted into a ReaderT
  bracketR
    :: ReaderT r IO a        -- ^ before
    -> (a -> ReaderT r IO b) -- ^ after
    -> (a -> ReaderT r IO c) -- ^ computation
    -> ReaderT r IO c
  bracketR before after thing = ReaderT $ \env ->
    bracket (runReaderT before env)
            (flip runReaderT env . after)
            (flip runReaderT env . thing)
  
  -- | Same as 'bracket_' from 'Control.Exception', but lifted into a ReaderT
  bracketR_
    :: ReaderT r IO a -- ^ before
    -> ReaderT r IO b -- ^ after
    -> ReaderT r IO c -- ^ computation
    -> ReaderT r IO c
  bracketR_ before after thing =
    bracketR before (const after) (const thing)
  
  -- | Retries the given computation up to n times if it fails.
  -- WARNING: only use this with actions that have NO SIDE EFFECTS on failure
  withRetry :: Int -> ReaderT r IO a -> ReaderT r IO a
  withRetry n action = foldl catchR action $ replicate n $
    \e@SomeException{} -> do lift $ print e ; action
  