Partial function crashes when hidden behind a type class

[hgoldstein95]

Prerequisites

Please put an X between the brackets as you perform the following steps:

• Check that your issue is not already filed:
  https://github.com/leanprover/lean4/issues
• Reduce the issue to a minimal, self-contained, reproducible test case.
  Avoid dependencies to Mathlib or Batteries.
• Test your test case against the latest nightly release, for example on
  https://live.lean-lang.org/#project=lean-nightly
  (You can also use the settings there to switch to “Lean nightly”)

Description

Sorry for the cryptic title, I'm not 100% sure how to explain this. I've been playing around with a DSL embedded using a tagless-final style with typeclasses, and I managed to crash Lean with what seems like it should be a pretty innocuous piece of code. Here's the MWE (or at least as minimal as I could come up with):

class RandomChoice (m : Type → Type) where
  choose : (lo hi : Nat) → (h : lo ≤ hi) → m Nat

def RandomChoice.pick [Monad m] [RandomChoice m] (x y : m α) := do
  if (← choose 0 1 (by simp)) == 0 then x else y

open RandomChoice

def Gen (α : Type) := ∀ {m : Type → Type} [Monad m] [RandomChoice m], m α

instance : RandomChoice IO where
  choose lo hi _ := IO.rand lo hi

partial def Nat.arbitrary : Gen Nat := do
  pick
    (pure 0)
    (do
      let n ← Nat.arbitrary
      pure (n + 1))

def Gen.runIO (g : Gen α) : IO α := g

-- Works fine.
#eval (Nat.arbitrary : IO Nat)

-- This breaks!
#eval Gen.runIO Nat.arbitrary

Context

I encountered this when working with @nomeata on some partial_fixpoint stuff. Originally I thought this must be related to that, but it turns out it works with partial too. The type-class embedding I'm doing is pretty central to the approach, so unfortunately I don't have a good way to work around this problem.

Steps to Reproduce

See the example above. (Tested on nightly.)

Expected behavior: The second #eval behaves the same as the first.

Actual behavior: The second #eval crashes Lean.

Versions

Tested on the Lean online editor on nightly.

[hargoniX]

The reason for the issue is that Nat.arbitrary gets optimized into the following code if it does not get specialized (which your, rather elaborate, hiding of the monad arguments ensures):

[Compiler.saveMono] size: 1
    def Nat.arbitrary m.1 inst.2 inst.3 : lcAny :=
      let _x.4 := Nat.arbitrary._redArg inst.2 inst.3;
      return _x.4
[Compiler.saveMono] size: 3
    def Nat.arbitrary._redArg._lam_0 toPure n : lcAny :=
      let _x.1 := 1;
      let _x.2 := Nat.add n _x.1;
      let _x.3 := toPure ◾ _x.2;
      return _x.3
[Compiler.saveMono] size: 10
    def Nat.arbitrary._redArg inst.1 inst.2 : lcAny :=
      cases inst.1 : lcAny
      | Monad.mk toApplicative toBind =>
        cases toApplicative : lcAny
        | Applicative.mk toFunctor toPure toSeq toSeqLeft toSeqRight =>
          let _f.3 := Nat.arbitrary._redArg._lam_0 toPure;
          let _x.4 := 0;
          let _x.5 := toPure ◾ _x.4;
          let _x.6 := Nat.arbitrary._redArg inst.1 inst.2;
          let _x.7 := toBind ◾ ◾ _x.6 _f.3;
          let _x.8 := RandomChoice.pick._redArg inst.1 inst.2 _x.5 _x.7;
          return _x.8

At this point we lost because _x.6 causes an infinite loop. I guess this kind of transformation would be acceptable in a lazy language because x_6 would only be evaluated further down the line in RandomChoice but it breaks here because it is eagerly evaluated.

A minified version of this issue would be:

set_option trace.Compiler.saveMono true

@[nospecialize, noinline]
def choice [Monad m] (flag : Bool) (a b : m α) : m α := if flag then a else b

@[nospecialize]
partial def loop [Monad m] : m Nat :=
  choice false loop (pure 0)

def test : IO Nat := loop

You can fix this by making choice take explicitly delayed arguments like a : Unit -> m α, I'm not quite sure what the compiler can do here to make the code work in general.

[hgoldstein95]

Ah yes, I used to write a lot of Haskell, so my brain sometimes defaults lazy semantics.

Although now that I think about it, isn't it confusing that (Nat.arbitrary : IO Nat) doesn't diverge? I think part of the reason I was surprised by this behavior was that I haven't yet been caught by it when working directly in IO. The language reference says that #eval evaluates terms in CBV, I think that's not quite true in this case. Should these programs have the same behavior, even if that means both of them diverging?

[nomeata]

Also, doesn't the problem also happen if you add a unit parameter to arbitrary? (On the phone right now)

[hargoniX]

Although now that I think about it, isn't it confusing that (Nat.arbitrary : IO Nat) doesn't diverge? I think part of the reason I was surprised by this behavior was that I haven't yet been caught by it when working directly in IO. The language reference says that #eval evaluates terms in CBV, I think that's not quite true in this case. Should these programs have the same behavior, even if that means both of them diverging?

The behavior doesn't occur because Lean specializes on IO in that situation and manages to fix up your code to actually be lazy. As I mentioned in my original message I do understand why the compiler does the things it does in both situations and I'm not sure which one of the two is the right one. That being said the current behavior is at least not surprising (to the initiated 😄) and already accounted for in other parts of our monad generic APIs.

Also, doesn't the problem also happen if you add a unit parameter to arbitrary? (On the phone right now)

The Unit parameter has to be placed like this:

set_option trace.Compiler.saveMono true

@[nospecialize, noinline]
def choice [Monad m] (flag : Bool) (a b : Unit → m α) : m α := if flag then a () else b ()

@[nospecialize]
partial def loop [Monad m] : m Nat :=
  choice false (fun _ => loop) (fun _ => pure 0)

def test : IO Nat := loop

This is also what's done in e.g. Seq.seq

[nomeata]

@hgoldstein95, if all your practical problems can be worked around by adding unit parameters, does that unblock your work?

[hgoldstein95]

Oh yes, thunking the arguments to pick fixes my issue. Thank you for the help! Sorry, I should have said that explicitly.

I was just continuing the discussion because users of my eventual library may be confused to find that code they write in my library diverges, even though there's seemingly equivalent code written directly in IO that works just fine. My two cents is that the "fixing up" that Lean is doing to make the code lazy may be more confusing than helpful in the long-run, although now I've nerd-sniped myself because as a PL+HCI researcher it occurs to me that this is a thing that might be worth actually studying...

Anyway, for now if this is intended behavior then I can just document that in my library and hope users don't get too stuck.

[nomeata]

This is related to GHC's “state hack”, I believe.