    ------------------------------------------------------------------------
    -- The Agda standard library
    --
    -- The indexed state monad
    ------------------------------------------------------------------------

    {-# OPTIONS --cubical-compatible --safe #-}

    module Effect.Monad.State.Indexed where

    open import Effect.Applicative.Indexed
      using (IFun; RawIApplicative; RawIApplicativeZero; RawIAlternative)
    open import Effect.Monad using (RawMonad; RawMonadZero; RawMonadPlus)
    open import Function.Identity.Effectful as Id using (Identity)
    open import Effect.Monad.Indexed
      using (RawIMonad; RawIMonadZero; RawIMonadPlus)
    open import Data.Product.Base using (_×_; _,_; uncurry)
    open import Data.Unit.Polymorphic using (⊤)
    open import Function.Base using (const; _∘_)
    open import Level using (Level; _⊔_; suc)

    private
      variable
        i f : Level
        I : Set i

    ------------------------------------------------------------------------
    -- Indexed state

    IStateT : (I → Set f) → (Set f → Set f) → IFun I f
    IStateT S M i j A = S i → M (A × S j)

    ------------------------------------------------------------------------
    -- Indexed state applicative

    StateTIApplicative : ∀ (S : I → Set f) {M} →
                         RawMonad M → RawIApplicative (IStateT S M)
    StateTIApplicative S Mon = record
      { pure = λ a s → pure (a , s)
      ; _⊛_  = λ f t s → do
         (f′ , s′)  ← f s
         (t′ , s′′) ← t s′
         pure (f′ t′ , s′′)
      } where open RawMonad Mon

    StateTIApplicativeZero : ∀ (S : I → Set f) {M} →
                             RawMonadZero M → RawIApplicativeZero (IStateT S M)
    StateTIApplicativeZero S Mon = record
      { applicative = StateTIApplicative S rawMonad
      ; ∅           = const ∅
      } where open RawMonadZero Mon

    StateTIAlternative : ∀ (S : I → Set f) {M} →
                         RawMonadPlus M → RawIAlternative (IStateT S M)
    StateTIAlternative S Mon = record
      { applicativeZero = StateTIApplicativeZero S rawMonadZero
      ; _∣_             = λ m n s → m s ∣ n s
      } where open RawMonadPlus Mon

    ------------------------------------------------------------------------
    -- Indexed state monad

    StateTIMonad : ∀ (S : I → Set f) {M} → RawMonad M → RawIMonad (IStateT S M)
    StateTIMonad S Mon = record
      { return = λ x s → pure (x , s)
      ; _>>=_  = λ m f s → m s >>= uncurry f
      } where open RawMonad Mon

    StateTIMonadZero : ∀ (S : I → Set f) {M} →
                       RawMonadZero M → RawIMonadZero (IStateT S M)
    StateTIMonadZero S Mon = record
      { monad           = StateTIMonad S (RawMonadZero.rawMonad Mon)
      ; applicativeZero = StateTIApplicativeZero S Mon
      } where open RawMonadZero Mon

    StateTIMonadPlus : ∀ (S : I → Set f) {M} →
                       RawMonadPlus M → RawIMonadPlus (IStateT S M)
    StateTIMonadPlus S Mon = record
      { monad       = StateTIMonad S rawMonad
      ; alternative = StateTIAlternative S Mon
      } where open RawMonadPlus Mon

    ------------------------------------------------------------------------
    -- State monad operations

    record RawIMonadState {I : Set i} (S : I → Set f)
                          (M : IFun I f) : Set (i ⊔ suc f) where
      field
        monad : RawIMonad M
        get   : ∀ {i} → M i i (S i)
        put   : ∀ {i j} → S j → M i j ⊤

      open RawIMonad monad public

      modify : ∀ {i j} → (S i → S j) → M i j ⊤
      modify f = get >>= put ∘ f

    StateTIMonadState : ∀ {i f} {I : Set i} (S : I → Set f) {M} →
                        RawMonad M → RawIMonadState S (IStateT S M)
    StateTIMonadState S Mon = record
      { monad = StateTIMonad S Mon
      ; get   = λ s   → pure (s , s)
      ; put   = λ s _ → pure (_ , s)
      }
      where open RawMonad Mon