An update monad as defined by https://doi.org/10.4230/LIPIcs.TYPES.2013.1

Update.lean

module

class Monoid (ω : Type u) where
  unit : ω
  mul : ω → ω → ω
  unit_mul (w : ω) : mul unit w = w
  mul_unit (w : ω) : mul w unit = w
  mul_assoc (w₁ w₂ w₃ : ω) :
    mul (mul w₁ w₂) w₃ = mul w₁ (mul w₂ w₃)

class RightAct (ω σ : Type u) extends Monoid ω where
  act : ω → σ → σ
  act_unit (s : σ) : act unit s = s
  act_mul (w₁ w₂ : ω) (s : σ) :
    act (mul w₁ w₂) s = act w₂ (act w₁ s)

attribute [simp] Monoid.unit_mul Monoid.mul_unit
attribute [simp] RightAct.act_unit RightAct.act_mul

def UpdateM (ω σ α : Type u) :=
  σ → α × ω

namespace UpdateM

def run (x : UpdateM ω σ α) (s : σ) : α × ω :=
  x s

protected def pure [RightAct ω σ] (a : α) : UpdateM ω σ α :=
  fun _ => (a, Monoid.unit)

protected def bind [RightAct ω σ] (x : UpdateM ω σ α) (f : α → UpdateM ω σ β) : UpdateM ω σ β :=
  fun s =>
    let (a, w₁) := x.run s
    let (b, w₂) := (f a).run (RightAct.act w₁ s)
    (b, Monoid.mul w₁ w₂)

instance [RightAct ω σ] : Monad (UpdateM ω σ) where
  pure := UpdateM.pure
  bind := UpdateM.bind

protected def read [RightAct ω σ] : UpdateM ω σ σ :=
  fun s => (s, Monoid.unit)

instance [RightAct ω σ] : MonadReaderOf σ (UpdateM ω σ) where
  read := UpdateM.read

def tell [RightAct ω σ] (w : ω) : UpdateM ω σ Unit :=
  fun _ => ((), w)

@[simp]
theorem run_pure [RightAct ω σ] (a : α) (s : σ) :
    (Pure.pure a : UpdateM ω σ α).run s = (a, Monoid.unit) := rfl

@[simp]
theorem run_read [RightAct ω σ] (s : σ) :
    (MonadReader.read : UpdateM ω σ σ).run s = (s, Monoid.unit) := rfl

@[simp]
theorem run_tell [RightAct ω σ] (w : ω) (s : σ) :
    (tell w).run s = ((), w) := rfl

@[simp]
theorem run_bind [RightAct ω σ] (x : UpdateM ω σ α) (f : α → UpdateM ω σ β) (s : σ) :
    (x >>= f).run s =
      let (a, w₁) := x.run s
      let (b, w₂) := (f a).run (RightAct.act w₁ s)
      (b, Monoid.mul w₁ w₂) := rfl

@[simp]
theorem run_map [RightAct ω σ] (f : α → β) (x : UpdateM ω σ α) (s : σ) :
    (f <$> x).run s =
      let (a, w) := x.run s
      (f a, w) := by
  change (x >>= Pure.pure ∘ f).run _ = _
  simp

@[ext]
theorem ext [RightAct ω σ] {x y : UpdateM ω σ α} (h : ∀ (s : σ), x.run s = y.run s) : x = y := by
  funext s
  exact h s

theorem bind_assoc [RightAct ω σ]
    (x : UpdateM ω σ α) (f : α → UpdateM ω σ β) (g : β → UpdateM ω σ γ) :
    x >>= f >>= g = x >>= fun x => f x >>= g := by
  ext1 s
  simp
  rcases x.run s with ⟨a, w₁⟩
  rcases (f a).run (RightAct.act w₁ s) with ⟨b, w₂⟩
  rcases (g b).run (RightAct.act w₂ (RightAct.act w₁ s))
  simp [Monoid.mul_assoc]

instance [RightAct ω σ] : LawfulMonad (UpdateM ω σ) := .mk'
  (id_map := by
    intros α x
    ext1
    simp)
  (pure_bind := by
    intros α β a f
    ext1
    simp)
  (bind_assoc := bind_assoc)

end UpdateM

UpdateCache.lean

module

class Monoid (ω : Type u) where
  unit : ω
  mul : ω → ω → ω
  unit_mul (w : ω) : mul unit w = w
  mul_unit (w : ω) : mul w unit = w
  mul_assoc (w₁ w₂ w₃ : ω) :
    mul (mul w₁ w₂) w₃ = mul w₁ (mul w₂ w₃)

class RightAct (ω σ : Type u) extends Monoid ω where
  act : ω → σ → σ
  act_unit (s : σ) : act unit s = s
  act_mul (w₁ w₂ : ω) (s : σ) :
    act (mul w₁ w₂) s = act w₂ (act w₁ s)

attribute [simp] Monoid.unit_mul Monoid.mul_unit
attribute [simp] RightAct.act_unit RightAct.act_mul

def UpdateM (ω σ : Type u) [RightAct ω σ] (α : Type u) :=
  (w₁ : ω) → (s₁ : σ) →
    α × (w₂ : ω) × { s₂ // ∃ w, w₂ = Monoid.mul w₁ w ∧ s₂ = RightAct.act w s₁ }

namespace UpdateM

def run [RightAct ω σ] (x : UpdateM ω σ α) (w : ω) (s : σ) : α × ω × σ :=
  let ⟨a, w, ⟨s, _⟩⟩ := x w s
  (a, w, s)

protected def pure [RightAct ω σ] (a : α) : UpdateM ω σ α :=
  fun w s => ⟨a, w, s, Monoid.unit, by simp⟩

protected def bind [RightAct ω σ] (x : UpdateM ω σ α) (f : α → UpdateM ω σ β) : UpdateM ω σ β :=
  fun w₀ s₀ =>
    let ⟨a, w₁, s₁, hw₁⟩ := x w₀ s₀
    let ⟨b, w₂, s₂, hw₂⟩ := f a w₁ s₁
    have h := by
      rcases hw₁ with ⟨w₁', hw₁, hs₁⟩
      rcases hw₂ with ⟨w₂', hw₂, hs₂⟩
      exists Monoid.mul w₁' w₂'
      simp [hs₁, hs₂, hw₁, hw₂, Monoid.mul_assoc]
    ⟨b, w₂, s₂, h⟩

instance [RightAct ω σ] : Monad (UpdateM ω σ) where
  pure := UpdateM.pure
  bind := UpdateM.bind

protected def read [RightAct ω σ] : UpdateM ω σ σ :=
  fun w s => ⟨s, w, s, Monoid.unit, by simp⟩

instance [RightAct ω σ] : MonadReaderOf σ (UpdateM ω σ) where
  read := UpdateM.read

def tell [RightAct ω σ] (w : ω) : UpdateM ω σ Unit :=
  fun w₀ s₀ => ⟨(), Monoid.mul w₀ w, RightAct.act w s₀, w, by simp⟩

@[simp]
theorem run_pure [RightAct ω σ] (a : α) (w : ω) (s : σ) :
    (Pure.pure a : UpdateM ω σ α).run w s = (a, w, s) := rfl

@[simp]
theorem run_read [RightAct ω σ] (w : ω) (s : σ) :
    (MonadReader.read : UpdateM ω σ σ).run w s = (s, w, s) := rfl

@[simp]
theorem run_tell [RightAct ω σ] (w' w : ω) (s : σ) :
    (tell w').run w s = ((), Monoid.mul w w', RightAct.act w' s) := rfl

theorem run_spec [RightAct ω σ] (x : UpdateM ω σ α) (w₀ : ω) (s₀ : σ) :
    ∃ a w, x.run w₀ s₀ = (a, Monoid.mul w₀ w, RightAct.act w s₀) := by
  rcases hx : x w₀ s₀ with ⟨a, w₁, s₁, w, rfl, rfl⟩
  exists a, w
  rw [run, hx]

theorem delta_of_run_eq [RightAct ω σ] {x : UpdateM ω σ α}
    (h : x.run w₀ s₀ = (a, w₁, s₁)) :
    ∃ w, w₁ = Monoid.mul w₀ w ∧ s₁ = RightAct.act w s₀ := by
  rcases run_spec x w₀ s₀ with ⟨_, w, hspec⟩
  rcases h ▸ hspec with ⟨_, rfl, rfl⟩
  exists w

@[simp]
theorem run_bind [RightAct ω σ] (x : UpdateM ω σ α) (f : α → UpdateM ω σ β) (w₀ : ω) (s₀ : σ) :
    (x >>= f).run w₀ s₀ =
      let (a, w₁, s₁) := x.run w₀ s₀
      (f a).run w₁ s₁ := rfl

@[simp]
theorem run_map [RightAct ω σ] (f : α → β) (x : UpdateM ω σ α) (w : ω) (s : σ) :
    (f <$> x).run w s =
      let (a, w', s') := x.run w s
      (f a, w', s') := by
  change (x >>= Pure.pure ∘ f).run w s = _
  simp

@[ext]
theorem ext [RightAct ω σ] {x y : UpdateM ω σ α}
    (h : ∀ (w : ω) (s : σ), x.run w s = y.run w s) : x = y := by
  funext w s
  specialize h w s
  rcases hx : x w s with ⟨a, wx, sx, hwx⟩
  rcases hy : y w s with ⟨b, wy, sy, hwy⟩
  rw [run, run, hx, hy] at h
  cases h
  rfl

theorem bind_assoc [RightAct ω σ]
    (x : UpdateM ω σ α) (f : α → UpdateM ω σ β) (g : β → UpdateM ω σ γ) :
    x >>= f >>= g = x >>= fun x => f x >>= g := by
  ext1 w s
  simp

instance [RightAct ω σ] : LawfulMonad (UpdateM ω σ) := .mk'
  (id_map := by
    intros α x
    ext1 w s
    simp)
  (pure_bind := by
    intros α β a f
    ext1 w s
    simp)
  (bind_assoc := bind_assoc)

end UpdateM