RFC: 'Functional but in place' compiler should recognize when the return object is equal to the destructed object

The functional but in place compiler, as described in the "Counting Immutable Beans" paper, looks for input objects to be reused for constructing an output object. At runtime, the reference count of the input object is checked to see whether the input object is shared. If it is shared, the output object has to be created from scratch. Otherwise, the memory of the input object can be reused for the output object.

Proposal

In the case that the output object is exactly equal to the input object, this reference count check should not be performed, and the input object should be returned no matter what.

Motivating example:

The implementation of bind for the IO monad is

/-- The `bind` operation of the `EStateM` monad. -/
@[always_inline, inline]
protected 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

Notice that the output object Result.error e s is exactly equal to x s. The IR of | Result.error e s => Result.error e s looks like this:

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

And I propose that it should instead be

EStateM.Result.error →
    dec x_2;
    ret x_4

Since the IO monad is widely used, and the bind function is inlined for every ← in a do block, this redundant IR makes the IR for many functions a lot bigger than it has to be.

In #7301, I implemented the more optimal implementation of bind manually by using unsafeCast:

stdlib bench:
stdlib size — lines C -3.2 %
build instructions -0,22 %

mathlib bench:
build instructions -0.11% (which is an order of magnitude above noise)
lint instructions -0.31%

Another example I found is the function Lean.RBNode.balLeft, where the IR shows that it tries to match on argument r unnecessarily, and then tries to reuse r depending on whether it is shared. But this example is easy to fix by splitting the match into 2 matches like in Lean.RBNode.balRight

Community Feedback

I mentioned this about half a year ago as well: https://leanprover.zulipchat.com/#narrow/channel/270676-lean4/topic/Suboptimal.20compilation.20of.20.60bind.60.20in.20.60IO.60.20monad/near/468812428

Impact

Add 👍 to issues you consider important. If others benefit from the changes in this proposal being added, please ask them to add 👍 to it.