content: -- Copyright 2023 Google LLC
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  
  {-# LANGUAGE UndecidableInstances #-}
  
  module MonadUtil (
    DefuncState (..), LabelReader (..), SingletonLabel (..), FreshNames (..),
    runFreshNameT, FreshNameT (..), Logger (..), LogLevel (..), getIOLogger, CanSetIOLogger (..),
    IOLoggerT (..), runIOLoggerT, LoggerT (..), runLoggerT,
    IOLogger (..), HasIOLogger (..), captureIOLogs) where
  
  import Control.Monad.Reader
  import Control.Monad.State.Strict
  import Control.Monad.Writer.Strict
  import Data.IORef
  
  import Err
  
  -- === Defunctionalized state ===
  -- Interface for state whose allowable updates are specified by a data type.
  -- Useful for `IncState`, for specifying read-only env components, or
  -- generally for specifying certain constraints on updates.
  
  class DefuncState d m | m -> d where
    update :: d -> m ()
  
  class LabelReader (l :: * -> *) m | m -> l where
    getl :: l a -> m a
  
  data SingletonLabel a b where
    It :: SingletonLabel a a
  
  -- === Fresh name monad ===
  
  -- Used for ad-hoc names with no nested binders that don't need to be treated
  -- carefully using the whole "foil" name system.
  
  class Monad m => FreshNames a m | m -> a where
    freshName :: m a
  
  newtype FreshNameT m a = FreshNameT { runFreshNameT' :: StateT Int m a }
          deriving (Functor, Applicative, Monad, MonadIO)
  
  instance MonadIO m => FreshNames Int (FreshNameT m) where
    freshName = FreshNameT do
      fresh <- get
      put (fresh + 1)
      return fresh
  
  instance FreshNames a m => FreshNames a (ReaderT r m) where
    freshName = lift freshName
  
  runFreshNameT :: MonadIO m => FreshNameT m a -> m a
  runFreshNameT cont = evalStateT (runFreshNameT' cont) 0
  
  -- === Logging monad ===
  
  data IOLogger w = IOLogger { ioLogLevel  :: LogLevel
                             , ioLogAction :: w -> IO () }
  data LogLevel = NormalLogLevel | DebugLogLevel
  
  class (Monoid w, Monad m) => Logger w m | m -> w where
    emitLog :: w -> m ()
    getLogLevel :: m LogLevel
  
  newtype IOLoggerT w m a = IOLoggerT { runIOLoggerT' :: ReaderT (IOLogger w) m a }
          deriving (Functor, Applicative, Monad, MonadIO, Fallible, MonadFail, Catchable)
  
  class Monad m => HasIOLogger w m | m -> w where
    getIOLogAction :: Monad m => m (w -> IO ())
  
  class Monad m => CanSetIOLogger w m | m -> w where
    withIOLogAction :: Monad m => (w -> IO ()) -> m a -> m a
  
  instance (Monoid w, MonadIO m) => HasIOLogger w (IOLoggerT w m) where
    getIOLogAction = IOLoggerT $ asks ioLogAction
  
  instance (Monoid w, MonadIO m) => CanSetIOLogger w (IOLoggerT w m) where
    withIOLogAction logger (IOLoggerT m) = IOLoggerT do
      local (\r -> r { ioLogAction = logger }) m
  
  instance (Monoid w, MonadIO m) => Logger w (IOLoggerT w m) where
    emitLog w = do
      logger <- getIOLogAction
      liftIO $ logger w
    getLogLevel = IOLoggerT $ asks ioLogLevel
  
  getIOLogger :: (HasIOLogger w m, Logger w m) => m (IOLogger w)
  getIOLogger = IOLogger <$> getLogLevel <*> getIOLogAction
  
  runIOLoggerT :: (Monoid w, MonadIO m) => LogLevel -> (w -> IO ()) -> IOLoggerT w m a -> m a
  runIOLoggerT logLevel write cont = runReaderT (runIOLoggerT' cont) (IOLogger logLevel write)
  
  newtype LoggerT w m a = LoggerT { runLoggerT' :: WriterT w m a }
          deriving (Functor, Applicative, Monad, MonadIO)
  
  instance (Monoid w, Monad m) => Logger w (LoggerT w m) where
    emitLog w = LoggerT $ tell w
    getLogLevel = return NormalLogLevel
  
  runLoggerT :: (Monoid w, Monad m) => LoggerT w m a -> m (a, w)
  runLoggerT cont = runWriterT (runLoggerT' cont)
  
  captureIOLogs
    :: forall w m a. (Monoid w, MonadIO m, HasIOLogger w m, CanSetIOLogger w m)
    => m a -> m (a, w)
  captureIOLogs cont = do
    ref <- liftIO $ newIORef (mempty :: w)
    ans <- withIOLogAction (\w -> modifyIORef ref (<> w)) cont
    w <- liftIO $ readIORef ref
    return (ans, w)
  
  