Based on the material discussed in https://www.youtube.com/watch?v=9v4_FQm-b4I

State.agda

module State where

open import Data.Unit
open import Data.Product
open import Function
open import Relation.Binary.PropositionalEquality

_⟶_ : {I : Set} (X Y : I → Set) → Set
X ⟶ Y = ∀ {i} → X i → Y i

κ : {I : Set} → Set → (I → Set)
κ A = λ _ → A

record StateSig (S : Set) : Set₁ where
  field
    St      : (S → Set) → S → Set
    return  : ∀ {X} → X ⟶ St X
    compose : ∀ {X Y Z} → (X ⟶ St Y) → (Y ⟶ St Z) → (X ⟶ St Z)
    put     : ∀ {s} → (ŝ : S) → St (λ s′ → s′ ≡ ŝ) s
    get     : ∀ {s} → St (λ s′ → Σ[ š ∈ S ] (s′ ≡ s × š ≡ s)) s

  bind : ∀ {X Y} → (X ⟶ St Y) → St X ⟶ St Y
  bind = λ f m → compose (λ x → x) f m

  syntax bind f m = m >>= f

StateM : (S : Set) → StateSig S
StateM S = record
  { St      = λ P s → Σ[ s′ ∈ S ] P s′
  ; return  = ,_
  ; compose = λ f g Xs → g (proj₂ $ f Xs)
  ; put     = λ ŝ → ŝ , refl
  ; get     = λ {s} → s , s , refl , refl
  }

open import Data.Nat

module _ (s : StateSig ℕ) (open StateSig s) where

  increment : ∀ {k} → St (λ n → n ≡ suc k) k
  increment =
    get          >>= λ { (k′ , eq , eq') →
    put (suc k′) >>= λ eq'' →
    return (trans eq'' (cong suc eq')) }

  decrement : ∀ {k l} → k ≡ suc l → St (λ n → n ≡ l) k
  decrement refl =
    get           >>= λ { (k' , eq , eq') →
    put (pred k') >>= λ eq'' →
    return (trans eq'' (cong pred eq')) }

  st2 : St (λ n → n ≡ 0) 2
  st2 = decrement refl >>= λ eq →
        decrement eq   >>= λ eq →
        return eq

  st3 : St (λ n → n ≡ 0) 3
  st3 = decrement refl >>= λ eq →
        decrement eq   >>= λ eq →
        decrement eq   >>= λ eq →
        return eq