/-
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura, Sebastian Ullrich

The State monad transformer using IO references.
-/
prelude
import Init.System.IO
import Init.Control.State

def StateRefT' (ω : Type) (σ : Type) (m : Type → Type) (α : Type) : Type := ReaderT (ST.Ref ω σ) m α

/- Recall that `StateRefT` is a macro that infers `ω` from the `m`. -/

@[inline] def StateRefT'.run {ω σ : Type} {m : Type → Type} [Monad m] [MonadLiftT (ST ω) m] {α : Type} (x : StateRefT' ω σ m α) (s : σ) : m (α × σ) := do
  let ref ← ST.mkRef s
  let a ← x ref
  let s ← ref.get
  pure (a, s)

@[inline] def StateRefT'.run' {ω σ : Type} {m : Type → Type} [Monad m] [MonadLiftT (ST ω) m] {α : Type} (x : StateRefT' ω σ m α) (s : σ) : m α := do
  let (a, _) ← x.run s
  pure a

namespace StateRefT'
variable {ω σ : Type} {m : Type → Type} {α : Type}

@[inline] protected def lift (x : m α) : StateRefT' ω σ m α :=
  fun _ => x

instance [Monad m] : Monad (StateRefT' ω σ m) := inferInstanceAs (Monad (ReaderT _ _))
instance : MonadLift m (StateRefT' ω σ m) := ⟨StateRefT'.lift⟩
instance (σ m) [Monad m] : MonadFunctor m (StateRefT' ω σ m) := inferInstanceAs (MonadFunctor m (ReaderT _ _))

@[inline] protected def get [Monad m] [MonadLiftT (ST ω) m] : StateRefT' ω σ m σ :=
  fun ref => ref.get

@[inline] protected def set [Monad m] [MonadLiftT (ST ω) m] (s : σ) : StateRefT' ω σ m PUnit :=
  fun ref => ref.set s

@[inline] protected def modifyGet [Monad m] [MonadLiftT (ST ω) m] (f : σ → α × σ) : StateRefT' ω σ m α :=
  fun ref => ref.modifyGet f

instance [MonadLiftT (ST ω) m] [Monad m] : MonadStateOf σ (StateRefT' ω σ m) := {
  get       := StateRefT'.get
  set       := StateRefT'.set
  modifyGet := StateRefT'.modifyGet
}

instance (ε) [MonadExceptOf ε m] : MonadExceptOf ε (StateRefT' ω σ m) := {
  throw    := StateRefT'.lift ∘ throwThe ε
  tryCatch := fun x c s => tryCatchThe ε (x s) (fun e => c e s)
}

end StateRefT'

instance (ω σ : Type) (m : Type → Type) : MonadControl m (StateRefT' ω σ m) :=
  inferInstanceAs (MonadControl m (ReaderT _ _))

instance {m : Type → Type} {ω σ : Type} [MonadFinally m] [Monad m] : MonadFinally (StateRefT' ω σ m) :=
  inferInstanceAs (MonadFinally (ReaderT _ _))