Dynamically Identifying Linearity Bugs in Lean 4

One of the key ingredients for writing performant Lean code is its functional but in-place (FBIP) idiom. While Lean does have automatic memory management, it does not use a normal garbage collector, unlike its functional relatives. Instead, every allocation is managed via reference counting. This allows Lean programs to determine at execution time whether an allocation is unique and update it in place instead of performing a copy-on-write.

The most important FBIP data structure in Lean is Array. As long as the reference count of an Array is equal to 1, it will be updated in place and have the same asymptotics as a std::vector. Otherwise, it acts as a copy-on-write data structure. On top of Array Lean builds many other data types, such as the imperative Std.HashMap and Std.HashSet. Thanks to the in-place updates, all of these data structures can maintain the asymptotics of imperative data structures, as long as they are referenced uniquely. On the flip side, once they do fall back to their copy-on-write behavior, the asymptotics usually degrade by an O(n) factor or worse. Because the property of being uniquely referenced is closely related to linear type systems, the Lean community calls these types "linear data structures", and performance issues due to violating this property are "linearity bugs".

Linearity Bugs

Linearity bugs usually occur when a linear data structure gets kept alive slightly too long in order to be eligible for an in-place update. For example, in the following code we would like xs to be updated in place:

def pushIt (xs : Array Nat) (n : Nat) : Array Nat × Nat :=
  let ys := xs.push n
  (ys, xs.size)

However, the fact that xs is still used after the creation of ys means that its reference count will be 2 instead of 1 inside Array.push. Because of this, Array.push will first create a copy of xs and then push onto it. To avoid this copy, we have to write the following:

def pushIt (xs : Array Nat) (n : Nat) : Array Nat × Nat :=
  let sz := xs.size
  let ys := xs.push n
  (ys, sz)

While the non-linearity in this example is easy to spot, there are many other ways to mess up linearity. The most common one occurs in combination with monadic state:

abbrev M := StateRefT (Array Nat) IO

def pushIt (n : Nat) : M Unit := do
  let xs ← get
  let ys := xs.push n
  set ys

In this example, xs is also used non-linearly because the state monad still maintains a reference to it while Array.push is being executed. In order to achieve a linear update, we must either manually clear the state before doing a push or use the modify operation, which ensures that the state monad doesn't hold a reference to the Array as it runs:

abbrev M := StateRefT (Array Nat) IO

def pushIt (n : Nat) : M Unit := do
  modify fun xs => xs.push n

Detecting Linearity Bugs

In the past, there has been some work towards building a type system extension for Lean that can enforce linearity constraints, in a similar vein to Rust's type system. Unfortunately, this line of work is not yet at the point where it can be used in production, so linearity bugs cannot be detected in a static fashion today.

Instead of detecting linearity issues statically, we can attempt to identify them dynamically at execution time. This kind of approach has already proven very successful for detecting things like memory safety issues in C and C++ using tools like AddressSanitizer or ThreadSanitizer. However, not all copy-on-write uses of an Array are linearity bugs. For example, the built-in hash array mapped trie type Lean.PersistentHashMap makes use of the copy-on-write behavior of Array. For this reason, we cannot declare every copy-on-write on an Array as a bug.

Given that only the author of a program knows which data structures are expected to be used linearly and non-linearly, we can also expect them to tell us about this. For this purpose, Lean v4.35.0-rc1 introduced markLinear functions on all built-in data structures that are supposed to be used linearly. If a value has been markLinear-ed and the environment variable LEAN_ABORT_ON_NONLINEAR is set, Lean will crash on the first copy-on-write access to this value. This already makes it possible to detect the above linearity issue:

abbrev M := StateRefT (Array Nat) IO

def pushIt (n : Nat) : M Unit := do
  let xs ← get
  let ys := xs.push n
  set ys

def main (args : List String) : IO Unit := do
  let some size := String.toNat? args[0]! | throw <| .userError "Bad size"
  let some init := String.toNat? args[1]! | throw <| .userError "Bad init"
  let some value := String.toNat? args[2]! | throw <| .userError "Bad value"
  let arr := (Array.replicate size init).markLinear
  let (arr, _) ← StateRefT'.run (pushIt value) arr
  IO.println s!"{arr}"

$ LEAN_ABORT_ON_NONLINEAR=1 lake exe linear 1 2 3
INTERNAL PANIC: array marked by `Array.markLinear` was used non-linearly

One important distinction between most sanitizers and the implementation of markLinear is that calling markLinear comes at basically no execution time cost to the program. For this reason, developers can liberally apply markLinear across their codebase without any drawbacks. Then once it comes time to look for a linearity bug, users can enable LEAN_ABORT_ON_NONLINEAR and profit from all the already installed annotations.

While this can help us identify that a linearity issue is present, e.g. during large-scale runs of test suites, it does not yet help identify where the linearity issue occurs in the code. For identifying the precise locations, it helps to first compile our code with debug symbols enabled. This can be done with a simple modification of lakefile.toml and a rebuild:

name = "linear"
version = "0.1.0"
defaultTargets = ["linear"]

[[lean_exe]]
name = "linear"
root = "Main"
+ moreLeancArgs = ["-g3"]

Then we can run the program with a debugger and break on lean_internal_panic to catch the abort:

$ LEAN_ABORT_ON_NONLINEAR=1 gdb --args ./.lake/build/bin/linear 1 2 3
(gdb) b lean_internal_panic
Breakpoint 1 at 0x1943c0
(gdb) run
Starting program: /linear/.lake/build/bin/linear 1 2 3
Thread 3 "linear" hit Breakpoint 1, 0x00005555556e83c0 in lean_internal_panic ()
(gdb) bt
#0  0x00005555556e83c0 in lean_internal_panic ()
#1  0x00005555556e7141 in lean_copy_expand_array_nonlinear ()
#2  0x00005555556e718c in lean_array_push ()
#3  0x00005555555f7e01 in lp_linear_pushIt (v_n_1_=<optimized out>, v_a_2_=<optimized out>) at /linear/.lake/build/ir/Main.c:56
#4  _lean_main (v_args_40_=<optimized out>) at /linear/.lake/build/ir/Main.c:200
#5  0x00005555556d7d27 in std::__1::__function::__func<lean_run_main::$_0, void ()>::operator()[abi:nqe230100]() ()
#6  0x00005555556d857a in lean::lthread::imp::_main(void*) ()
#7  0x00007ffff7e14d19 in start_thread (arg=<optimized out>) at pthread_create.c:454
#8  0x00007ffff7e9864c in __GI___clone3 () at ../sysdeps/unix/sysv/linux/x86_64/clone3.S:78

As we can see, we are now stuck in the lean runtime, but a piece of our own code, pushIt, is on the call stack as well; let's inspect it:

(gdb) up 3
#3  0x00005555555f7e01 in lp_linear_pushIt (v_n_1_=<optimized out>, v_a_2_=<optimized out>) at /linear/.lake/build/ir/Main.c:56
56	v___x_5_ = lean_array_push(v___x_4_, v_n_1_);
(gdb) list
51	LEAN_EXPORT lean_object* lp_linear_pushIt(lean_object* v_n_1_, lean_object* v_a_2_){
52	_start:
53	{
54	lean_object* v___x_4_; lean_object* v___x_5_; lean_object* v___x_6_; lean_object* v___x_7_; lean_object* v___x_8_;
55	v___x_4_ = lean_st_ref_get(v_a_2_);
56	v___x_5_ = lean_array_push(v___x_4_, v_n_1_);
57	v___x_6_ = lean_st_ref_swap(v_a_2_, v___x_5_);
58	lean_dec(v___x_6_);
59	v___x_7_ = lean_box(0);
60	v___x_8_ = lean_alloc_ctor(0, 1, 0);

Unfortunately, no source-level Lean debugger exists yet, so this, together with the views we can get into the compilation pipeline, will have to suffice for identifying the location of the non-linearity:

abbrev M := StateRefT (Array Nat) IO

set_option trace.Compiler.saveBase true in
set_option trace.Compiler.saveMono true in
set_option trace.Compiler.saveImpure true in
def pushIt (n : Nat) : M Unit := do
  let xs ← get
  let ys := xs.push n
  set ys

[Compiler.saveBase] size: 8
    def pushIt n @&a a.1 : EST.Out IO.Error lcAny PUnit :=
      let _x.2 := @ST.Prim.Ref.get _ _ a a.1;
      cases _x.2 : EST.Out IO.Error lcAny PUnit
      | ST.Out.mk val.3 state.4 =>
        let _x.5 := @Array.push _ val.3 n;
        let _x.6 := @ST.Prim.Ref.set _ _ a _x.5 state.4;
        cases _x.6 : EST.Out IO.Error lcAny PUnit
        | ST.Out.mk val.7 state.8 =>
          let _x.9 := @EST.Out.ok _ _ _ val.7 state.8;
          return _x.9
[Compiler.saveMono] size: 8
    def pushIt n @&a a.1 : EST.Out IO.Error lcAny PUnit :=
      let _x.2 := @ST.Prim.Ref.get ◾ ◾ a a.1;
      cases _x.2 : EST.Out IO.Error lcAny PUnit
      | ST.Out.mk val.3 state.4 =>
        let _x.5 := Array.push ◾ val.3 n;
        let _x.6 := @ST.Prim.Ref.set ◾ ◾ a _x.5 state.4;
        cases _x.6 : EST.Out IO.Error lcAny PUnit
        | ST.Out.mk val.7 state.8 =>
          let _x.9 := @EST.Out.ok ◾ ◾ ◾ val.7 state.8;
          return _x.9
[Compiler.saveImpure] size: 4
    def pushIt n @&a a.1 : obj :=
      let _x.2 := ST.Prim.Ref.get ◾ ◾ a ◾;
      let _x.3 := Array.push ◾ _x.2 n;
      let _x.4 := ST.Prim.Ref.set ◾ ◾ a _x.3 ◾;
      let _x.5 := ctor_0[EST.Out.ok] _x.4;
      return _x.5

If we squint a bit, we can see that the reported location of the non-linearity at lean_array_push corresponds to our xs.push n call from the original code, which is indeed correct.

Instead of aborting, we can also employ conditional breakpoints to break on the copy-on-write function for arrays whenever a marked array passes through it:

(gdb) break *lean_copy_expand_array_nonlinear if (((lean_object*)$rdi)->m_other & 0x80) != 0
# $rdi is the register in which lean_copy_expand_array_nonlinear receives the array
# `m_other & 0x80` is where the lean runtime stores the `markLinear` flag
Breakpoint 1 at 0x193100
(gdb) run
Thread 3 "linear" hit Breakpoint 1, 0x00005555556e7100 in lean_copy_expand_array_nonlinear ()
(gdb) bt
#0  0x00005555556e7100 in lean_copy_expand_array_nonlinear ()
#1  0x00005555556e718c in lean_array_push ()
#2  0x00005555555f7e01 in lp_linear_pushIt (v_n_1_=<optimized out>, v_a_2_=<optimized out>) at /linear/.lake/build/ir/Main.c:56
#3  _lean_main (v_args_40_=<optimized out>) at /linear/.lake/build/ir/Main.c:200
#4  0x00005555556d7d27 in std::__1::__function::__func<lean_run_main::$_0, void ()>::operator()[abi:nqe230100]() ()
#5  0x00005555556d857a in lean::lthread::imp::_main(void*) ()
#6  0x00007ffff7e14d19 in start_thread (arg=<optimized out>) at pthread_create.c:454
#7  0x00007ffff7e9864c in __GI___clone3 () at ../sysdeps/unix/sysv/linux/x86_64/clone3.S:78

If LEAN_ABORT_ON_NONLINEAR is not set, continuing the program at this point will break again at the next linearity issue and so on, allowing us to collect statistics about all the issues present in a program. This can, for example, be used to prioritize linearity issues in preexisting large code bases.

Conclusion

In addition to Array.markLinear the standard library ships the same function on String, Std.HashSet, Std.HashMap, Std.DHashMap, ByteArray, FloatArray, and Vector which are by far the most common linear data structures in use today. I hope that these functions together with the above debugging recipes will allow us to detect, identify and fix linearity issues in a more principled way as opposed to just looking at the generated code until we realize why they are present in the future. It has already started helping Lean itself.