map Array while mutating state

Stream: lean4

[Scott Morrison]

I want to map a monadic function over an Array, while mutating some additional state. What's the idiomatic way to do this?

def F (f : α → β → MetaM (α × γ)) (a : α) (bs : Array β) : MetaM (α × Array γ) := do
  let mut z := a
  let gs ← bs.mapM (fun b => do
    let (g, a') ← f z b
    z ← a'
    return g)
  return (z, gs)

Fails with

`z` cannot be mutated, only variables declared using `let mut` can be mutated. If you did not intent to mutate but define `z`, consider using `let z` instead.

[Scott Morrison]

I know how to do it "the old way" with StateT:

def G (f : α → β → MetaM (α × γ)) : β → StateT α MetaM γ := fun b => do
  let (a, g) ← f (←get) b
  set a
  return g

def F (f : α → β → MetaM (α × γ)) (a : α) (bs : Array β) : MetaM (Array γ × α) := do
  StateT.run (bs.mapM (G f)) a

#eval F (fun n m => return (n+1, n*m)) 0 #[1,2,3,4] -- prints (#[0, 2, 6, 12], 4)

(I changed α × Array γ to Array γ × α for convenience here.)

[Adam Topaz]

Is z \l g what's intended? (z has type \a)

[Adam Topaz]

Anyway, I got something to work with a for loop

def F (f : α → β → MetaM (α × γ)) (a : α) (bs : Array β) : MetaM (α × Array γ) := do
  let mut z := a
  let mut gs : Array γ := #[]
  for b in bs do
    let (a',g) ← f z b
    z := a'
    gs := gs.push g
  return (z, gs)

#eval F (fun n m => return (n+1, n*m)) 0 #[1,2,3,4]

[Tomas Skrivan]

I encountered identical problem with a similar high order function like mapM but which cannot be easily replaced with a for loop.

Is it possible for let mut to introduce the required state monad?

[Mario Carneiro]

let mut does introduce the required state monad. The hard part is that do is basically a closed class, the elaborator has to know all the things that are going on in the block and in particular it can't lift its state over arbitrary monadic combinators. It has specific support for bundling and unbundling its state to do the try .. catch, try .. finally and for .. in .. do combinators, which all have specific syntax in the do block.

[Mario Carneiro]

the usual way you will see this being done in lean 4 is what Adam wrote, just don't use mapM at all and push into an array in a loop

[Mario Carneiro]

it's not that much worse than what mapM does already; the main optimization it has over pushing in a loop is that the initial array is not #[] but Array.mkEmpty bs.size