Weird issue with do notation/reader monad

Stream: general

[Aaron Eline]

I'm not sure why these should be semantically different, but they are? Unsure if I'm just really misunderstanding things. Could someone clarify why this works the way it does or if I'm holding it wrong?

abbrev M (α : Type)  : Type := ReaderT Env (Result TypeError) α

def locally (f : Env → Env) (k : M α) : M α :=  do
  let env ← read
  k (f env)

def locally' (f : Env → Env) (k : M α) : M α :=  do
  let env ← read
  -- This does not work
  -- Results in bindings getting lost
  ReaderT.run k (f env)

image.png
(This is the bind that locally' desugars

[Malvin Gattinger]

I wanted to try out the code on live.lean-lang.org but cannot get it to compile. Do you have an #mwe with imports?

[Aaron Eline]

Here is a mwe:

abbrev M (α : Type)  : Type := ReaderT Nat Option α

def locally (f : Nat → Nat) (k : M α) : M α :=  do
  let env ← read
  k (f env)

def locally' (f : Nat → Nat) (k : M α) : M α :=  do
  let env ← read
  -- This does not work
  -- Results in bindings getting lost
  ReaderT.run k (f env)

def some_computation (x : Nat) : M Nat := do
  let z ← read;
  return (x + z)

def higher : M Nat := do
  locally (λ x => x + 1) (some_computation 2)

def higher' : M Nat := do
  locally' (λ x => x + 1) (some_computation 2)

def top : Option Nat :=
  ReaderT.run higher 0

def top' : Option Nat :=
  ReaderT.run higher' 0

#eval top
#eval top'

[Eric Wieser]

Using #print you can see you got

def locally' : {α : Type} → (Nat → Nat) → M α → M α :=
fun {α} f k => do
  let env ← read
  (liftM k).run (f env)

where that liftM is probably causing trouble

[Aaron Liu]

liftM k has type ReaderT Nat M α

[Aaron Liu]

so you've added on an additional Nat reader state which is then immediately discarded by liftM

[Aaron Eline]

Sure, but _why_, is the do notation behaving weird here? ReaderT.run k (f e) and k (f e) reduce to the same thing, so why is using inside do changing things?

[Aaron Liu]

k (f env) has type Option α so it inserts a coercion and you get liftM (k (f env))

[Aaron Eline]

k (f env) is the one that produces the correct results?

[Aaron Liu]

ReaderT.run k (f e) has an expected type M α so it elaborates the argument k with expected type ReaderT Nat M α

[Aaron Liu]

so the coercion is inserted inside

[Aaron Liu]

this doesn't happen with k (f env) because the type of k doesn't have any holes in it so having an expected type doesn't change your elaboration

[Aaron Liu]

so it works if you elaborate ReaderT.run k (f e) without an expected type

[Harry Goldstein]

Is the claim here that this is a type inference issue, then? Somehow inference is assuming ReaderT.run is running a different reader monad, which is inserted via coercion, rather than running the reader that it seems like it should be running? This seems like pretty unintuitive behavior

[Eric Wieser]

If you write

def locally' (f : Nat → Nat) (k : M α) : M α :=  do
  let env ← read
  (ReaderT.run k (f env) :)

then it does what you expect

[Eric Wieser]

Otherwise the coercion gets inserted in the wrong place

[Eric Wieser]

At any rate, you're looking for docs#withTheReader I think

[Harry Goldstein]

I love that if you just write a little smiley face it works correctly! :)

Do we think there's any hope of making this behavior more intuitive / less error-prone? I don't think the code that Aaron wrote was particularly arcane — an average user could totally run into this and write some pretty hard-to-find bugs. Would a change to the implicit arguments of run help? Or maybe some kind of warning on do elaboration? I realize that this behavior is consistent with the rest of Lean, so it's not a bug in that sense, but I would call this behavior "incorrect" from the perspective of most reasonable end-users

[Eric Wieser]

One thing that I think would help is changing ReaderT (and stateT etc) to a structure along the lines of

structure ReaderT (R m A) where
  mk :: run (r : R) : m A

[Aaron Eline]

Chiming in that I did in fact find this behavior baffling, as someone who's relatively used to writing monadic code in ocaml/haskell. Took me a bit to find the workaround I did.

[Eric Wieser]

Maybe @Sebastian Ullrich knows if that change would result in worse compilation output

[Sebastian Ullrich]

It shouldn't, no

[Robin Arnez]

It would?

structure ReaderTStruct (ρ : Type u) (m : Type u → Type v) (α : Type u) where
  mk :: run (r : ρ) : m α

set_option trace.compiler.ir.result true
/--
trace: [Compiler.IR] [result]
    def testStruct._lam_0 (x_1 : @& tobj) : tobj :=
      inc x_1;
      ret x_1
    def testStruct._lam_0._boxed (x_1 : tobj) : tobj :=
      let x_2 : tobj := testStruct._lam_0 x_1;
      dec x_1;
      ret x_2
[Compiler.IR] [result]
    def testStruct._closed_0 : obj :=
      let x_1 : obj := pap testStruct._lam_0._boxed;
      ret x_1
    def testStruct : obj :=
      let x_1 : obj := testStruct._closed_0;
      ret x_1
-/
#guard_msgs in
def testStruct : ReaderTStruct Nat Id Nat :=
  .mk fun x => x

/--
trace: [Compiler.IR] [result]
    def test (x_1 : @& tobj) : tobj :=
      inc x_1;
      ret x_1
    def test._boxed (x_1 : tobj) : tobj :=
      let x_2 : tobj := test x_1;
      dec x_1;
      ret x_2
-/
#guard_msgs in
def test : ReaderT Nat Id Nat :=
  .mk fun x => x

ReaderTStruct forces the compiler to allocate a closure

[Robin Arnez]

At least, with the current state of the compiler

[Aaron Liu]

Eric Wieser said:

One thing that I think would help is changing ReaderT (and stateT etc) to a structure along the lines of

structure ReaderT (R m A) where
  mk :: run (r : R) : m A

I'm not sure this would help?