csimp can violate linear access of data

Stream: lean4

[Jeremy Avigad]

When you found that csimp did a swap that resulted in a shared reference, what did you do to fix it?

[Mario Carneiro]

The issue was that simp was basically inlining portions of the result of the function, but without actually deleting the call to the function and resulting in two versions of the value lying around. I marked the function as @[inline] so simp didn't have to deal with it

[Mario Carneiro]

I'll make a MWE

[Mario Carneiro]

structure Foo where
  a : String
  b : Unit
  c : Unit

def withA (f : String → String) (s : Foo) : Foo :=
  { s with a := f s.a }

def doStuffLinearly : String → String :=
dbgTraceIfShared "oh no"

set_option trace.compiler.ir.result true in
def bar (s : Foo) : Foo := do
  let s' := withA doStuffLinearly s
  -- let s' := dbgTraceIfShared "this would fail" s'
  pure { s' with c := () }

#eval show IO Unit from do
  let n ← IO.rand 0 1 -- to make sure lean doesn't put s in a global variable
  let s : Foo := ⟨toString n, (), ()⟩
  let s := dbgTraceIfShared "ok here" s
  let ⟨a, b, c⟩ ← bar s
  b

[Mario Carneiro]

The interesting part here is bar. It performs two linear operations on the state Foo, first it modifies a with doStuffLinearly, then we get the resulting value s' and resets c in it. However, the generated IR makes it clear that what lean is actually compiling is

def bar (s : Foo) : Foo := do
  let s' := withA doStuffLinearly s
  pure { a := s'.a, b := s.b, c := () }

where field b is taken from the previous value of s. At first I had these named the same thing and I thought it was some kind of name confusion, but what's actually happening is that lean looks at the definition of withA and determines that s'.b = s.b so the replacement is permitted. Which is true, but causes a reference to s to hang around when withA doStuffLinearly s is called, meaning that doStuffLinearly will not have linear access to s.a.

[Mario Carneiro]

def bar (x_1 : obj) : obj :=
  let x_2 : obj := bar._closed_1;
  inc x_1;
  let x_3 : obj := withA x_2 x_1;
  let x_4 : obj := proj[0] x_3;
  inc x_4;
  dec x_3;
  let x_5 : u8 := isShared x_1;
  case x_5 : u8 of
  Bool.false →
    let x_6 : obj := proj[2] x_1;
    dec x_6;
    let x_7 : obj := proj[0] x_1;
    dec x_7;
    let x_8 : obj := ctor_0[PUnit.unit];
    set x_1[2] := x_8;
    set x_1[0] := x_4;
    ret x_1
  Bool.true →
    let x_9 : obj := proj[1] x_1;
    inc x_9;
    dec x_1;
    let x_10 : obj := ctor_0[PUnit.unit];
    let x_11 : obj := ctor_0[Foo.mk] x_4 x_9 x_10;
    ret x_11

Here x_1 is s, and x_3 is s'. You can see references to x_1 appearing later in the function after the withA call, as well as inc x_1 before calling withA.

[Mario Carneiro]

Putting @[inline] on withA causes it to actually be inlined, instead of just partially inlined like this, meaning that all the destructuring happens only at the beginning of the function and the individual projections are used linearly.

[Jeremy Avigad]

Thanks, Mario. The example is clear, and the explanations are especially interesting and helpful.