State monad in F#

Stateful.fs

[<Struct>]
type Stateful<'state, 'result> =
  Stateful of ('state -> 'state * 'result)

type Delayed<'state, 'result> = unit -> Stateful<'state, 'result>

module State = 

  let update f =
    Stateful (fun state -> f state, ())

  let just x = 
    Stateful (fun state -> state, x)

  let bind (m : Stateful<'state, 't>) (f : 't -> Stateful<'state, 'u>) : Stateful<'state, 'u> = 
    Stateful (fun s1 -> 
      let (Stateful m) = m
      let s2, x = m s1
      let (Stateful f) = f x
      f s2)

  let doWhile (guard : unit -> bool) (body : Delayed<'state, unit>) = 
    let rec loop m = 
      if guard ()
      then
        bind m (fun () -> loop (body ()))
      else
        m

    loop (just ())

  let run state (m : Stateful<'state, 't>) = 
    let (Stateful m) = m
    m state

[<AutoOpen>]
module ComputationExpression =

  open State
  
  type StatefulBuilder() = 
    member this.Bind(m, f) = 
      bind m f

    member this.Return(x) =
      just x

    member this.Zero() =
      just ()

    member this.Delay(f : Delayed<'state, 'result>) = 
      f

    member this.Run(f : Delayed<'state, 'result>) =
      f ()

    member this.While(guard, body) = 
      doWhile guard body
    

  let state = StatefulBuilder()

state {
  return 3
} 
|> State.run ()
|> printfn "%A"

state {
  do! State.update (fun x -> x + 1)
  do! State.update (fun x -> x * 2)
} 
|> State.run 3
|> printfn "%A"

state {
  let mutable i = 0
  
  while i < 3 do
    do! State.update (fun state -> i :: state)
    i <- i + 1
} 
|> State.run []
|> printfn "%A"