Trying to understand StateRefT, ST.Ref, STWorld

Stream: lean4

[Richard L Ford]

For a new comer to Lean 4 (like me) , the monad structure of Lean 4 seems rather intimidating (though elegant).
In reading the source code, I see that many of the monads are defined using StateRefT. I learned that
this is a macro defined like this (in Parser/Term.lean):

@[builtinTermParser] def stateRefT   := leading_parser "StateRefT" >> macroArg >> macroLastArg

This is then elaborated (in Elab/BuiltinNotation.lean) like this:

@[builtinTermElab stateRefT] def elabStateRefT : TermElab := fun stx _ => do
  let σ ← elabType stx[1]
  let mut mStx := stx[2]
  if mStx.getKind == ``Lean.Parser.Term.macroDollarArg then
    mStx := mStx[1]
  let m ← elabTerm mStx (← mkArrow (mkSort levelOne) (mkSort levelOne))
  let ω ← mkFreshExprMVar (mkSort levelOne)
  let stWorld ← mkAppM ``STWorld #[ω, m]
  discard <| mkInstMVar stWorld
  mkAppM ``StateRefT' #[ω, σ, m]

So a StateRefT σ m expands into a StateRefT' ω σ m.
Here, STWorld is used to infer ω, based on previous instantiations of STWorld? STWorld and some instances are given
in ST.lean as follows:

-- Auxiliary class for inferring the "state" of `EST` and `ST` monads
class STWorld (σ : outParam Type) (m : Type → Type)

instance {σ m n} [MonadLift m n] [STWorld σ m] : STWorld σ n := ⟨⟩
instance {ε σ} : STWorld σ (EST ε σ) := ⟨⟩

StateRefT' is defined like this (in Init/Control/StateRef.lean):

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

I understand that ReaderT adds a context that can be accessed with <- read. So it seems this
is adding a context that is an ST.Ref. These are defined in Init/System/ST.lean:

structure Ref (σ : Type) (α : Type) : Type where
  ref : RefPointed.type
  h   : Nonempty α

One of my questions is why there is the (σ : Type) type variable in this definition since it is not used? In looking at generated code I learned that an ST.Ref is implemented (in C) as an

typedef struct {
    lean_object   m_header;
    lean_object * m_value;
} lean_ref_object;

So it is effectively a pointer to something of type α.
Operations on these ST.Refs are defined in stages. First (e.g.) ST.Prim.Ref.get is defined like this:

@[extern "lean_st_ref_get"]
constant Ref.get {σ α} (r : @& Ref σ α) : ST σ α := inhabitedFromRef r

I asked myself, why is σ involved here, and why does it have an ST σ α result, since the C that implements it does not use any state?
Is it because you want to use it from other functions in the ST monad?
Finally, outside the Prim namespace we have these "lifted" definitions, that define ST.Ref.get, for example:

section
variable {σ : Type} {m : Type → Type} [Monad m] [MonadLiftT (ST σ) m]
@[inline] def Ref.get {α : Type} (r : Ref σ α) : m α := liftM <| Prim.Ref.get r

So is the purpose of this lifting so that the ST.Ref functions can be used as actions in any monad?

A final question: It seems that the implementation mutates the ST.Ref's in place so their use seems to violate the purity of functional paradigm. Is that because we know these are only being used once so it is OK?

Sorry to be so long-winded. I just thought if I wrote something like this it might clarify my own thinking, and hopefully someone will be able to answer some of these questions.

[Sebastian Ullrich]

These are very good questions. If you're comfortable with Haskell, it's probably best to read an explanation of their corresponding ST monad, e.g. https://stackoverflow.com/a/12468757/161659. Their s is our σ.

[Leonardo de Moura]

There are different questions here. The design of Ref and ST is based on a similar approach used in Haskell. Haskell docs will explain in great detail the purpose of σ. The PointedType and Nonempty in the Lean version are used to make sure we are not compromising soundness. For example, we use them to justify that ST σ (Ref σ α) is not an empty type when we declare mkRef.
The StateRefT macro is not really needed, we added it just to improve the compilation times in the past. Not sure we even need it anymore. To understand the design you can pretend StateRefT' is being used instead.