Suboptimal compilation of `bind` in `IO` monad

Stream: lean4

[Jovan Gerbscheid]

I was looking at some output of set_option trace.compiler.ir.result true, and noticed that withReducible was compiled into a more complicated term than the equivalent withReader ({· with config.transparency := .reducible}). Minimizing this example leads to this suboptimally compiled function:

def foo (x : IO α) : IO α := do
  let a ← x
  pure a

which should be just the identity function, but gets compiled to

def foo._rarg (x_1 : obj) (x_2 : obj) : obj :=
  let x_3 : obj := app x_1 x_2;
  case x_3 : obj of
  EStateM.Result.ok →
    let x_4 : u8 := isShared x_3;
    case x_4 : u8 of
    Bool.false →
      ret x_3
    Bool.true →
      let x_5 : obj := proj[0] x_3;
      let x_6 : obj := proj[1] x_3;
      inc x_6;
      inc x_5;
      dec x_3;
      let x_7 : obj := ctor_0[EStateM.Result.ok] x_5 x_6;
      ret x_7
  EStateM.Result.error →
    let x_8 : u8 := isShared x_3;
    case x_8 : u8 of
    Bool.false →
      ret x_3
    Bool.true →
      let x_9 : obj := proj[0] x_3;
      let x_10 : obj := proj[1] x_3;
      inc x_10;
      inc x_9;
      dec x_3;
      let x_11 : obj := ctor_1[EStateM.Result.error] x_9 x_10;
      ret x_11

After inling the functions that are inlined by the compiler, we get the definition

def foo (x : IO α) : IO α :=
  fun s => match x s with
    | EStateM.Result.ok a s => EStateM.Result.ok a s
    | EStateM.Result.error e s => EStateM.Result.error e s

Now we can see that the problem is that the match expression is acting like the identity function, but this isn't seen by the compiler. If we replace either of | EStateM.Result.ok a s => EStateM.Result.ok a s or | EStateM.Result.error e s => EStateM.Result.error e s with | x => x, then the corresponding match branch gets compiled more efficiently, and if we do this for both branches, foo is compiled as the identity function.

This optimization wouldn't just help to reduce away a combination of bind and pure. It would improve the compilation of every single bind, because looking at the definition of bind:

def bind (x : EStateM ε σ α) (f : α → EStateM ε σ β) : EStateM ε σ β := fun s =>
  match x s with
  | Result.ok a s    => f a s
  | Result.error e s => Result.error e s

we see that we still have | Result.error e s => Result.error e s, and this part gets compiled to

EStateM.Result.error →
  dec x_2;
  let x_8 : u8 := isShared x_4;
  case x_8 : u8 of
  Bool.false →
    ret x_4
  Bool.true →
    let x_9 : obj := proj[0] x_4;
    let x_10 : obj := proj[1] x_4;
    inc x_10;
    inc x_9;
    dec x_4;
    let x_11 : obj := ctor_1[EStateM.Result.error] x_9 x_10;
    ret x_11

which instead should just be

EStateM.Result.error →
  dec x_2;
  ret x_4

So I think that implementing this would make the compiled size of functions in IO, MetaM, etc. significantly smaller.

Specifically, what the compiler simplifier should do is, for every match branch, check if the result is exactly the thing that is being matched against (or a constructor of the same shape with the same fields), and in that case make this branch simply return that thing that is being matched against.

[Henrik Böving]

Many issues like this are known with the old compiler and tracked in the core issue tracker. They will eventually get fixed with the new compiler but there is no plan to invest time into the old one anymore. If you want you can add your thing onto the heap of TODOs

[Mauricio Collares]

Item 7 in the year 2 roadmap says "Bandwidth-permitting, we will also deliver a new code generator", if you want a rough timeline

[Henrik Böving]

Emphasis on bandwith-permitting^^

[Eric Wieser]

Does the "new compiler" already exist in any form?

[Henrik Böving]

Its been 70% done a couple of years ago mhm

[Henrik Böving]

70% done to the point of "it would work end to end" that is

[Eric Wieser]

Is this "70% done and in some unused file" or "70% done in a (private?) branch"?

[Henrik Böving]

Lean.Compiler.LCNF