content: -- Copyright 2021 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 MTL1 where
  
  import Control.Monad.Reader
  import Control.Monad.Writer.Class
  import Control.Monad.State.Strict
  import Control.Monad.Trans.Maybe
  import Control.Applicative
  import Data.Foldable (toList)
  
  import Name
  import Err
  import Types.Top (Env)
  import Core (EnvReader (..), EnvExtender (..))
  import Util (SnocList (..), snoc, emptySnocList)
  
  class MonadTrans11 (t :: MonadKind1 -> MonadKind1) where
    lift11 :: Monad1 m => m n a -> t m n a
  
  -------------------- WriterT1 --------------------
  
  newtype WriterT1 (w :: E) (m :: MonadKind1) (n :: S) (a :: *) =
    WrapWriterT1 { runWriterT1' :: (StateT (w n) (m n) a) }
    deriving ( Functor, Applicative, Monad, MonadFail
             , Fallible, MonadIO)
  
  pattern WriterT1 :: ((w n) -> m n (a, w n)) -> WriterT1 w m n a
  pattern WriterT1 f = WrapWriterT1 (StateT f)
  {-# COMPLETE WriterT1 #-}
  
  runWriterT1 :: Monoid1 w => WriterT1 w m n a -> m n (a, w n)
  runWriterT1 = runWriterT1From mempty
  {-# INLINE runWriterT1 #-}
  
  runWriterT1From :: Monoid1 w => w n -> WriterT1 w m n a -> m n (a, w n)
  runWriterT1From w m = runStateT (runWriterT1' m) w
  {-# INLINE runWriterT1From #-}
  
  instance (Monad1 m, Monoid1 w) => MonadWriter (w n) (WriterT1 w m n) where
    writer (a, w) = WrapWriterT1 $ a <$ modify (<> w)
    {-# INLINE writer #-}
    tell w = WrapWriterT1 $ modify (<> w)
    {-# INLINE tell #-}
    listen (WrapWriterT1 m) = WrapWriterT1 $ do
      cur <- get
      put mempty
      ans <- m
      ext <- get
      put $ cur <> ext
      return (ans, ext)
    {-# INLINE listen #-}
    pass (WrapWriterT1 m) = WrapWriterT1 $ do
      cur <- get
      put mempty
      (ans, f) <- m
      ext <- get
      put $ cur <> f ext
      return ans
    {-# INLINE pass #-}
  
  instance Monoid1 w => MonadTrans11 (WriterT1 w) where
    lift11 = WrapWriterT1 . lift
    {-# INLINE lift11 #-}
  
  instance (SinkableE w, Monoid1 w, EnvReader m) => EnvReader (WriterT1 w m) where
    unsafeGetEnv = lift11 unsafeGetEnv
    {-# INLINE unsafeGetEnv #-}
  
  instance (SinkableE w, Monoid1 w, ScopeReader m) => ScopeReader (WriterT1 w m) where
    unsafeGetScope = lift11 unsafeGetScope
    {-# INLINE unsafeGetScope #-}
    getDistinct = lift11 getDistinct
    {-# INLINE getDistinct #-}
  
  instance ( SinkableE w, Monoid1 w
           , HoistableState w, EnvExtender m)
           => EnvExtender (WriterT1 w m) where
    refreshAbs ab cont = WriterT1 \s -> do
      (ans, Abs b new) <- refreshAbs ab \b e -> do
        (ans, new) <- runWriterT1From (sink s) $ cont b e
        return (ans, Abs b new)
      let new' = hoistState s b new
      return (ans, s <> new')
    {-# INLINE refreshAbs #-}
  
  -------------------- ReaderT1 --------------------
  
  newtype ReaderT1 (r :: E) (m :: MonadKind1) (n :: S) (a :: *) =
    ReaderT1 { runReaderT1' :: (ReaderT (r n) (m n) a) }
    deriving (Functor, Applicative, Monad, MonadFail, MonadReader (r n))
  
  runReaderT1 :: r n -> ReaderT1 r m n a -> m n a
  runReaderT1 r m = runReaderT (runReaderT1' m) r
  {-# INLINE runReaderT1 #-}
  
  instance MonadTrans11 (ReaderT1 r) where
    lift11 = ReaderT1 . lift
    {-# INLINE lift11 #-}
  
  deriving instance MonadWriter s (m n) => MonadWriter s (ReaderT1 r m n)
  
  deriving instance MonadState s (m n) => MonadState s (ReaderT1 r m n)
  
  instance (Monad1 m, Alternative1 m) => Alternative ((ReaderT1 r m) n) where
    empty = lift11 empty
    {-# INLINE empty #-}
    ReaderT1 (ReaderT m1) <|> ReaderT1 (ReaderT m2) =
      ReaderT1 $ ReaderT \r -> m1 r <|> m2 r
    {-# INLINE (<|>) #-}
  
  
  instance (SinkableE r, EnvReader m) => EnvReader (ReaderT1 r m) where
    unsafeGetEnv = lift11 unsafeGetEnv
    {-# INLINE unsafeGetEnv #-}
  
  instance (SinkableE r, ScopeReader m) => ScopeReader (ReaderT1 r m) where
    unsafeGetScope = lift11 unsafeGetScope
    {-# INLINE unsafeGetScope #-}
    getDistinct = lift11 getDistinct
    {-# INLINE getDistinct #-}
  
  instance (SinkableE r, ScopeExtender m) => ScopeExtender (ReaderT1 r m) where
    refreshAbsScope ab cont = ReaderT1 $ ReaderT \r -> do
      refreshAbsScope ab \b e -> runReaderT1 (sink r) $ cont b e
  
  instance (SinkableE r, EnvExtender m) => EnvExtender (ReaderT1 r m) where
    refreshAbs ab cont = ReaderT1 $ ReaderT \r -> do
      refreshAbs ab \b e -> runReaderT1 (sink r) $ cont b e
  
  instance (Monad1 m, Fallible (m n)) => Fallible (ReaderT1 r m n) where
    throwErr = lift11 . throwErr
  
  instance (Monad1 m, Catchable (m n)) => Catchable (ReaderT1 s m n) where
    catchErr (ReaderT1 m) f = ReaderT1 $ catchErr m (runReaderT1' . f)
  
  -------------------- StateT1 --------------------
  
  newtype StateT1 (s :: E) (m :: MonadKind1) (n :: S) (a :: *) =
    WrapStateT1 { runStateT1' :: (StateT (s n) (m n) a) }
    deriving ( Functor, Monad, MonadState (s n)
             , MonadFail, MonadIO)
  
  -- This is entirely standard, but we implement it explicitly to encourage GHC to inline.
  instance (Monad (m n), Applicative (m n)) => Applicative (StateT1 s m n) where
    pure = WrapStateT1 . pure
    {-# INLINE pure #-}
    (WrapStateT1 f) <*> (WrapStateT1 x) = WrapStateT1 $ f <*> x
    {-# INLINE (<*>) #-}
    liftA2 f (WrapStateT1 x) (WrapStateT1 y) = WrapStateT1 $ liftA2 f x y
    {-# INLINE liftA2 #-}
  
  pattern StateT1 :: ((s n) -> m n (a, s n)) -> StateT1 s m n a
  pattern StateT1 f = WrapStateT1 (StateT f)
  {-# COMPLETE StateT1 #-}
  
  type MonadState1 (e::E) (m::MonadKind1) = forall n. MonadState (e n) (m n)
  
  runStateT1 :: StateT1 s m n a -> s n -> m n (a, s n)
  runStateT1 = runStateT . runStateT1'
  {-# INLINE runStateT1 #-}
  
  evalStateT1 :: Monad1 m => StateT1 s m n a -> s n -> m n a
  evalStateT1 m s = fst <$> runStateT1 m s
  {-# INLINE evalStateT1 #-}
  
  instance MonadTrans11 (StateT1 s) where
    lift11 = WrapStateT1 . lift
    {-# INLINE lift11 #-}
  
  instance (SinkableE s, EnvReader m) => EnvReader (StateT1 s m) where
    unsafeGetEnv = lift11 unsafeGetEnv
    {-# INLINE unsafeGetEnv #-}
  
  instance (SinkableE s, ScopeReader m) => ScopeReader (StateT1 s m) where
    unsafeGetScope = lift11 unsafeGetScope
    {-# INLINE unsafeGetScope #-}
    getDistinct = lift11 getDistinct
    {-# INLINE getDistinct #-}
  
  instance (Monad1 m, Fallible (m n)) => Fallible (StateT1 s m n) where
    throwErr = lift11 . throwErr
  
  instance (Monad1 m, Catchable (m n)) => Catchable (StateT1 s m n) where
    catchErr (WrapStateT1 m) f = WrapStateT1 $ catchErr m (runStateT1' . f)
  
  instance (Monad1 m, Alternative1 m) => Alternative ((StateT1 s m) n) where
    empty = lift11 empty
    {-# INLINE empty #-}
    StateT1 m1 <|> StateT1 m2 = StateT1 \s -> m1 s <|> m2 s
    {-# INLINE (<|>) #-}
  
  class HoistableState (s::E) where
    hoistState :: BindsNames b => s n -> b n l -> s l -> s n
  
  instance (SinkableE s, EnvExtender m, HoistableState s) => EnvExtender (StateT1 s m) where
    refreshAbs ab cont = StateT1 \s -> do
      (ans, Abs b s') <- refreshAbs ab \b e -> do
        (ans, s') <- flip runStateT1 (sink s) $ cont b e
        return (ans, Abs b s')
      let s'' = hoistState s b s'
      return (ans, s'')
  
  instance HoistableState (LiftE a) where
    hoistState _ _ (LiftE x) = LiftE x
    {-# INLINE hoistState #-}
  
  instance HoistableState UnitE where
    hoistState _ _ UnitE = UnitE
    {-# INLINE hoistState #-}
  
  instance Show a => HoistableState (NameMap c a) where
    hoistState _ b m = hoistNameMap b m
    {-# INLINE hoistState #-}
  
  -------------------- ScopedT1 --------------------
  
  newtype ScopedT1 (s :: E) (m :: MonadKind1) (n :: S) (a :: *) =
    WrapScopedT1 { runScopedT1' :: StateT1 s m n a }
    deriving ( Functor, Monad, MonadState (s n), MonadFail
             , MonadTrans11, EnvReader, ScopeReader )
  
  -- This is entirely standard, but we implement it explicitly to encourage GHC to inline.
  instance (Monad (m n), Applicative (m n)) => Applicative (ScopedT1 s m n) where
    pure = WrapScopedT1 . pure
    {-# INLINE pure #-}
    (WrapScopedT1 f) <*> (WrapScopedT1 x) = WrapScopedT1 $ f <*> x
    {-# INLINE (<*>) #-}
    liftA2 f (WrapScopedT1 x) (WrapScopedT1 y) = WrapScopedT1 $ liftA2 f x y
    {-# INLINE liftA2 #-}
  
  pattern ScopedT1 :: ((s n) -> m n (a, s n)) -> ScopedT1 s m n a
  pattern ScopedT1 f = WrapScopedT1 (StateT1 f)
  {-# COMPLETE ScopedT1 #-}
  
  runScopedT1 :: Monad1 m => ScopedT1 s m n a -> s n -> m n a
  runScopedT1 m s = fst <$> runStateT1 (runScopedT1' m) s
  {-# INLINE runScopedT1 #-}
  
  deriving instance (Monad1 m, Fallible1 m) => Fallible (ScopedT1 s m n)
  deriving instance (Monad1 m, Catchable1 m) => Catchable (ScopedT1 s m n)
  
  instance (SinkableE s, EnvExtender m) => EnvExtender (ScopedT1 s m) where
    refreshAbs ab cont = ScopedT1 \s -> do
      ans <- refreshAbs ab \b e -> flip runScopedT1 (sink s) $ cont b e
      return (ans, s)
  
  -------------------- MaybeT1 --------------------
  
  newtype MaybeT1 (m :: MonadKind1) (n :: S) (a :: *) =
    MaybeT1 { runMaybeT1' :: (MaybeT (m n) a) }
    deriving (Functor, Applicative, Monad, Alternative)
  
  runMaybeT1 :: MaybeT1 m n a -> m n (Maybe a)
  runMaybeT1 = runMaybeT . runMaybeT1'
  {-# INLINE runMaybeT1 #-}
  
  instance MonadTrans11 MaybeT1 where
    lift11 = MaybeT1 . lift
    {-# INLINE lift11 #-}
  
  instance Monad (m n) => MonadFail (MaybeT1 m n) where
    fail s = MaybeT1 (fail s)
    {-# INLINE fail #-}
  
  instance Monad (m n) => Fallible (MaybeT1 m n) where
    throwErr _ = empty
  
  instance EnvReader m => EnvReader (MaybeT1 m) where
    unsafeGetEnv = lift11 unsafeGetEnv
    {-# INLINE unsafeGetEnv #-}
  
  instance ScopeReader m => ScopeReader (MaybeT1 m) where
    unsafeGetScope = lift11 unsafeGetScope
    {-# INLINE unsafeGetScope #-}
    getDistinct = lift11 getDistinct
    {-# INLINE getDistinct #-}
  
  instance EnvExtender m => EnvExtender (MaybeT1 m) where
    refreshAbs ab cont = MaybeT1 $ MaybeT $
      refreshAbs ab \b e -> runMaybeT $ runMaybeT1' $ cont b e
  
  -------------------- StreamWriter --------------------
  
  class Monad m => StreamWriter w m | m -> w where
    writeStream :: w -> m ()
  
  newtype StreamWriterT1 (w:: *) (m::MonadKind1) (n::S) (a:: *) =
    StreamWriterT1 { runStreamWriterT1' :: StateT1 (LiftE (SnocList w)) m n a }
    deriving (Functor, Applicative, Monad, MonadFail, MonadIO, ScopeReader, EnvReader)
  
  instance Monad1 m => StreamWriter w (StreamWriterT1 w m n) where
    writeStream w = StreamWriterT1 $ modify (\(LiftE ws) -> LiftE (ws `snoc` w))
    {-# INLINE writeStream #-}
  
  runStreamWriterT1 :: Monad1 m => StreamWriterT1 w m n a -> m n (a, [w])
  runStreamWriterT1 m = do
    (ans, LiftE ws) <- runStateT1 (runStreamWriterT1' m) (LiftE emptySnocList)
    return (ans, toList ws)
  {-# INLINE runStreamWriterT1 #-}
  
  -------------------- StreamReader --------------------
  
  class Monad m => StreamReader r m | m -> r where
    readStream :: m (Maybe r)
  
  newtype StreamReaderT1 (r:: *) (m::MonadKind1) (n::S) (a:: *) =
    StreamReaderT1 { runStreamReaderT1' :: StateT1 (LiftE [r]) m n a }
    deriving (Functor, Applicative, Monad, MonadFail, MonadIO, ScopeReader, EnvReader, MonadTrans11)
  
  instance Monad1 m => StreamReader r (StreamReaderT1 r m n) where
    readStream = StreamReaderT1 $ state \(LiftE rs) ->
      case rs of
        []       -> (Nothing, LiftE [])
        (r:rest) -> (Just r , LiftE rest)
    {-# INLINE readStream #-}
  
  runStreamReaderT1 :: Monad1 m => [r] -> StreamReaderT1 r m n a -> m n (a, [r])
  runStreamReaderT1 rs m = do
    (ans, LiftE rsRemaining) <- runStateT1 (runStreamReaderT1' m) (LiftE rs)
    return (ans, rsRemaining)
  {-# INLINE runStreamReaderT1 #-}
  
  -------------------- DiffState --------------------
  
  class MonoidE (d::E) where
    emptyE :: d n
    catE :: d n -> d n -> d n
  
  class MonoidE d => DiffStateE (s::E) (d::E) where
    updateDiffStateE :: Distinct n => Env n -> s n -> d n -> s n
  
  newtype DiffStateT1 (s::E) (d::E) (m::MonadKind1) (n::S) (a:: *) =
    DiffStateT1' { runDiffStateT1'' :: StateT (s n, d n) (m n) a }
    deriving ( Functor, Applicative, Monad, MonadFail, MonadIO
             , Fallible, Catchable)
  
  pattern DiffStateT1 :: ((s n, d n) -> m n (a, (s n, d n))) -> DiffStateT1 s d m n a
  pattern DiffStateT1 cont = DiffStateT1' (StateT cont)
  
  diffStateT1
    :: (EnvReader m, DiffStateE s d, MonoidE d)
    => (s n -> m n (a, d n)) -> DiffStateT1 s d m n a
  diffStateT1 cont = DiffStateT1 \(s, d) -> do
    (ans, d') <- cont s
    env <- unsafeGetEnv
    Distinct <- getDistinct
    return (ans, (updateDiffStateE env s d', catE d d'))
  {-# INLINE diffStateT1 #-}
  
  runDiffStateT1
    :: (EnvReader m, DiffStateE s d, MonoidE d)
    => s n -> DiffStateT1 s d m n a ->  m n (a, d n)
  runDiffStateT1 s (DiffStateT1' (StateT cont)) = do
    (ans, (_, d)) <- cont (s, emptyE)
    return (ans, d)
  {-# INLINE runDiffStateT1 #-}
  
  class (Monad1 m, MonoidE d)
        => MonadDiffState1 (m::MonadKind1) (s::E) (d::E) | m -> s, m -> d where
    withDiffState :: s n -> m n a -> m n (a, d n)
    updateDiffStateM :: d n -> m n ()
    getDiffState :: m n (s n)
  
  instance (EnvReader m, DiffStateE s d, MonoidE d) => MonadDiffState1 (DiffStateT1 s d m) s d where
    getDiffState = DiffStateT1' $ fst <$> get
    {-# INLINE getDiffState #-}
  
    withDiffState s cont = DiffStateT1' do
      (sOld, dOld) <- get
      put (s, emptyE)
      ans <- runDiffStateT1'' cont
      (_, dLocal) <- get
      put (sOld, dOld)
      return (ans, dLocal)
    {-# INLINE withDiffState #-}
  
    updateDiffStateM d = DiffStateT1' do
      (s, d') <- get
      env <- lift unsafeGetEnv
      Distinct <- lift getDistinct
      put (updateDiffStateE env s d, catE d d')
    {-# INLINE updateDiffStateM #-}
  
  instance MonoidE (ListE e) where
    emptyE = mempty
    catE = (<>)
  
  instance MonoidE (RListE e) where
    emptyE = mempty
    catE = (<>)
  
  instance (Monad1 m, Alternative1 m, MonoidE d) => Alternative ((DiffStateT1 s d m) n) where
    empty = DiffStateT1' $ StateT \_ -> empty
    {-# INLINE empty #-}
    DiffStateT1' (StateT m1) <|> DiffStateT1' (StateT m2) = DiffStateT1' $ StateT \s ->
      m1 s <|> m2 s
    {-# INLINE (<|>) #-}
  
  instance (ScopeReader m, MonoidE d) => ScopeReader (DiffStateT1 s d m) where
    unsafeGetScope = lift11 unsafeGetScope
    getDistinct = lift11 getDistinct
  
  instance (EnvReader m, MonoidE d) => EnvReader (DiffStateT1 s d m) where
    unsafeGetEnv = lift11 unsafeGetEnv
  
  instance MonadTrans11 (DiffStateT1 s d) where
    lift11 m = DiffStateT1' $ lift m
    {-# INLINE lift11 #-}
  