[web.git] / personal / _posts / 2019-04-30-stacked-borrows-2.md

---
title: "Stacked Borrows 2"
categories: rust research
---

Recently, I have [significantly updated](https://github.com/rust-lang/miri/pull/695) Stacked Borrows in order to fix some issues with the handling of shared references that were uncovered in the previous version.
In this post, I will describe what the new version looks like and how it differs from [Stacked Borrows 1]({% post_url 2018-11-16-stacked-borrows-implementation %}).
I assume some familiarity with the prior version and will not explain everything from scratch.

<!-- MORE -->

## The problem

The problem I wanted to solve was that the first version of Stacked Borrows only performed very little tracking of shared references.
My thinking was, if the location is read-only anyway, then it does not harm to grant anyone read access.
However, [as @arielby noted](https://github.com/rust-lang/unsafe-code-guidelines/issues/87), this leads to loss of optimization potential in cases where a function receives a mutable reference (which is supposed to have no aliases) and then creates a shared reference from it:
{% highlight rust %}
fn main() {
    let x = &mut 0u32;
    let p = x as *mut u32;
    foo(x, p);
}

fn foo(a: &mut u32, y: *mut u32) -> u32 {
    *a = 1;
    let _b = &*a; // This freezes `x`. Frozen locations can be read by any raw pointer.
    unsafe { *y = 2; } // Hence, this legal in Stacked Borrows.
    return *a; // But we want to optimize this to always return 1!
}
{% endhighlight %}
Stacked Borrows 1 allowed any raw pointer to read a frozen location.
Basically, once a location is frozen we do not keep track of which pointer is allowed to read; we just allow all pointers to read.
However, that means that `&*a` above (reborrowing a mutable reference as a shared reference) has the unintended side-effect of permitting raw pointers to access `*a`!
This violates the idea that `a` is a unique pointer: we'd like to assume that `foo` can only ever return `1`, because `y` cannot possibly alias with `a`.
Hence calling `foo` like we do here with `a` and `y` pointing to the same thing must be undefined behavior---and yet, Stacked Borrows 1 considered the example above to be legal.

## The solution

To fix this, I had to replace the mechanism of "freezing" by something else.
Remember that in Stacked Borrows 1, when a shared reference was created, we stored the timestamp at which that happened and also recorded in memory that the location this reference points to is frozen since now.
Whenever a shared reference gets used, we check that the location is frozen *at least* since our reference got created.
This is in contrast to mutable references, where we require the exact unique ID of that reference to be present in the borrow stack.

To rule out cases like the example above, instead of freezing a location and allowing *all* shared reference created henceforth to access this location, we keep track precisely of *which* shared references are allowed access.
The per-location borrow stack now consists of items that grant a particular *permission* to a pointer identified by a particular *tag*:
{% highlight rust %}
pub struct Item {
    /// The pointers the permission is granted to.
    tag: Tag,
    /// The permission this item grants.
    perm: Permission,
}

pub enum Permission {
    /// Grants unique mutable access.
    Unique,
    /// Grants shared mutable access.
    SharedReadWrite,
    /// Greants shared read-only access.
    SharedReadOnly,
}
{% endhighlight %}
There is no longer a separate "frozen" state; keeping shared references read-only is now handled by the borrow stack itself:
{% highlight rust %}
pub struct Stack {
    borrows: Vec<Item>
}
{% endhighlight %}
The *tag* is also simpler than it was before: there are no longer separate tags for mutable and shared references.
{% highlight rust %}
pub type PtrId = NonZeroU64;
pub enum Tag {
    Tagged(PtrId),
    Untagged,
}
{% endhighlight %}
Just like before, a new `Tag` gets picked on every retagging; in particular whenever a reference gets created with `&mut <expr>`/`& <expr>` and when a reference gets cast to a raw pointer.
Other than that (like on pointer arithmetic), the tag just gets propagated to keep track of where this pointer comes from.
However, unlike before, the only difference between mutable and shared references is the permissions that are associated with that tag.

### Memory accesses

But before we go into creation of new references, let us look at how permissions in the borrow stack affect what happens at a memory access is legal.
So assume that some pointer tagged `tag` is used to either read from or write to memory.
For each affected location, we go through two steps: first we try to find the *granting item*, then we remove *incompatible items*.

**Finding the granting item.**
To find the granting item, we traverse the borrow stack top-to-bottom and search for an item that has the same tag.
If this is a write access, we go on looking until we find an item with the right tag whose permission is `SharedReadWrite` or `Unique`---we do not consider items with `SharedReadOnly` permission to grant a write access.
Once we found the granting item, we remember its position in the stack; that will be important for the second step.
If no granting item can be found, the access causes undefined behavior.

For example, if the stack is
```
[ Item { tag: Tagged(0), perm: Unique },
  Item { tag: Tagged(1), perm: SharedReadOnly} ]
```
then a read access with tag `Tagged(1)` is granted by the item at index 1; a read/write access with tag `Tagged(0)` is granted by the item at index 0, but a write access with tag `Tagged(1)` is not granted by any item (and thus not allowed).
Just to keep things shorter and easier to read, in the following I will use a short-hand syntax for writing down items, so the above stack would look as follows:
```
[ (0: Unique), (1: SharedReadOnly) ]
```

To consider a second example, if the stack is
```
[ (0: Unique), (Untagged: SharedReadWrite), (Untagged: SharedReadOnly) ]
```
then a read access with an `Untagged` pointer is granted by the item at index 2, but a write access with an `Untagged` pointer is only granted by the item at index 1.

**Removing incompatible items.**
In the second step, we traverse all the items *above* the granting item in the stack, and see if they are compatible with this access.
This realizes the idea (already present in the original Stacked Borrows) that using one pointer disallows future uses of another pointer.
For example, if the current stack is
```
[ (0: Unique), (1: Unique) ]
```
then doing any kind of access with a `Tagged(0)` pointer should remove the `1: Unique` item from the stack.
This matches the part of Stacked Borrows 1 where items got popped off the stack until the granting item is at the top.
We say that the granting `Unique` permission is *incompatible* with the `Unique` permission of the item higher up the stack, and hence the latter must be removed.

However, in the new model, we don't always remove *all* items that are above the granting item:
```
[ (0: Unique), (1: SharedReadOnly), (2: SharedReadOnly) ]
```
In a situation like this, there are two pointers that may be used for reading, `Tagged(1)` and `Tagged(2)`.
Using either of them should not have an impact on the other one, and the fact that their items are in a particular order on the stack has no impact.
In other words, a granting `SharedReadOnly` permission is *compatible* with other `SharedReadOnly` permissions, and hence when using a `SharedReadOnly` permission to grant an access, other `SharedReadOnly` permissions above it are maintained.

Sometimes, being compatible depends on the kind of access that was performed:
in our last example stack, if we *write* with a `Tagged(0)` pointer, that should remove the `SharedReadOnly` tags because writing to a unique pointer should make sure that that pointer is at the top of the stack.
In other words, on a write, a `Unique` permission is incompatible with everything.
However, if we just *read* with a `Tagged(0)` pointer, that's fine!
This is to allow [safe code like this](https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=d9aa504b1be72a0f55fb5a744cba69ff) that interleaves reads from aliasing mutable and shared references.
To enable this, on a read, a `Unique` permission is only incompatible with other `Unique` permissions.

The following table fully documents which permissions are compatible with which other permissions on a read or write access, respectively:

|↓ **Granting item** \ **Compatible with** → | `Unique`  | `SharedReadWrite`   | `SharedReadOnly`  |
|---:|---|---|---|
|`Unique`                                    | no | only reads | only reads  |
|`SharedReadWrite`                           | no | yes  | only reads |
|`SharedReadOnly`                            | no | no  | only reads |

So, for example, if the granting item has `SharedReadWrite` permissions, then all `Unique` above it in the stack will be removed, and moreover on a write, all `SharedReadOnly` above it will be removed.
This makes sure that the pointers for which those `SharedReadOnly` permission grants read access cannot be used any more: those pointers assumed the location was read-only, so we must make sure that after a write, they become invalid.

The name "stack" is a bit of a misnomer at this point, because we do *not* follow a stack discipline when removing items: just because some item got removed due to incompatibility, does not mean all items above it will also get removed.
I will discuss [later][no-stack] why a strict stack discipline does not work.
But I felt it would be even more confusing to rename the entire thing at this point, and the order of elements in the "stack" still *does* matter, so I kept the name.

In Stacked Borrows 2, there is no more action occurring when a pointer gets dereferenced: the tags only matter for the actual access.
This simplification is possible because the rules for accesses are now more complicated than before, so there is no need for extra checks on the dereference operation.
(In particular, previously the dereference checks relied on knowing the Rust type at which the access happens for proper handling of interior mutability; but now all that information is encoded using `SharedReadWrite`/`SharedReadOnly` tags as we will see. Not relying on types any more here also [fixed an annoying issue around `UnsafeCell`](https://github.com/rust-lang/miri/issues/615).)

### Retagging: when permissions get added

The framework of permissions in a stack with a notion of (in)compatibility that we have seen so far allows us to express some ideas like:

* "A `Unique` pointer will get invalid when any of its parent pointers get used."<br> That is realized by having "no" in the entire first column of our table: every access removes all `Unique` permissions above the granting item.
* "A `SharedReadOnly` pointer will get invalid when the location gets written to."<br> This gets realized by having "only reads" in the entire last column of the table and not granting write accesses with a `SharedReadOnly` permission. In addition to that, however, we *also* need to ensure that in the stack we never have a `Unique` or `SharedReadWrite` *on top* of a `SharedReadOnly`! In a nonsense stack like `[ (0: SharedReadOnly), (1: SharedReadWrite) ]` we could write with a `Tagged(1)` pointer without invalidating the `Tagged(0)` pointers.

Now the question is, how to we *use* this "language"?
We have to define which items and permissions get added to the borrow stack.
Just like before, this happens during retagging.
Before we discuss retagging in general, let us look at our motivating example and see how retagging works there in Stacked Borrows 2, and how this program is defined to cause UB:
{% highlight rust %}
fn main() {
    let x = &mut 0u32;
    Retag(x); // Tagged(0)
    // Stack: [ (0: Unique) ]

    let p = x as *mut u32;
    Retag([raw] p); // Untagged
    // Stack: [ (0: Unique), (Untagged: SharedReadWrite) ]

    foo(x, p);
}

fn foo(a: &mut u32, y: *mut u32) -> u32 {
    Retag(a); // Tagged(1)
    // Stack: [ (0: Unique), (1: Unique) ]

    *a = 1;

    let _b = &*a;
    Retag(_b); // Tagged(2)
    // Stack: [ (0: Unique), (1: Unique), (2: SharedReadOnly) ]

    // UB: y is Untagged, and there is no granting item in the stack!
    unsafe { *y = 2; }

    return *a;
}
{% endhighlight %}
Initially, `x` with tag `Tagged(0)` is the only reference, and the stack says that this is the only pointer with any kind of permission.
Next, we cast `x` to a raw pointer.
The raw retagging of `p` turns `p` into an `Untagged` pointer, and adds a new item granting thusly tagged pointers `SharedReadWrite` permission.
(Really, in the MIR it will say `&mut *x as *mut u32`, so there will be an additional `Unique` permission for the temporary mutable reference, but that makes no difference and I hope [we will change that eventually](https://github.com/rust-lang/rfcs/pull/2582).)

Then `foo` gets called, which starts with the usual retagging of all reference arguments.
`a` is originally `Tagged(0)`, and retagging a mutable reference acts like an access (just like in Stacked Borrows 1), so the first thing that happens is that all items above the granting item that are incompatible with a write access for a `Unique` permission get removed from the stack.
In our case, this means that the `Untagged: SharedReadWrite` gets removed.
Then, `a` gets the new tag `Tagged(1)` and `1: Unique` gets pushed on top of the stack.
Nothing interesting happens when writing to `a`.
When `_b` gets created, it gets assigned the new tag `Tagged(2)`, and a new item `2: SharedReadOnly` is pushed to the stack.
As we can see, shared and mutable references no longer differ in the tag they carry; the only difference is what kind of permission they get granted.
(I think this is a nice improvement from Stacked Borrows 1, where shared references had a different kind of tag. In particular, transmutes between mutable references and shared pointers no longer need any kind of special treatment.)
Finally, we come to the interesting point: the program writes to `y`.
That pointer is `Untagged`, and there is no item granting any access to such pointers, and thus the access is UB!
This is in contrast to Stacked Borrows 1, where instead of `2: SharedReadOnly` we set a special "frozen" marker on this location that would, as a side-effect, also grant untagged pointers read-only access.

In general, during retagging, we start with some original tag used by our "parent" pointer, and we have a fresh tag `N` that will be used for the new pointer.
We have to decide which permission to grant so this tag, and where in the "stack" to insert it.
(We will not always insert new items at the top---again, as we will see following a strict stack discipline does not work.)
All of this happens for every location that the reference covers (as determined by its type).

**Retagging a mutable reference.**
The simplest case of retagging is handling mutable references:
just like in Stacked Borrows 1, this starts by performing the actions of a write access with the parent tag.
`Unique` permissions are incompatible with anything on a write, so all items above the granting item get removed.
Then we add `N: Unique` on top of the stack to grant the new tag unique access to this location.

This encodes the idea that mutable references must be used in a well-nested way---if you just consider mutable references, Stacked Borrows 2 still follows a stack discipline.

**Retagging a shared reference.**
When retagging a shared reference, we have to be mindful of `UnsafeCell`.

Outside of `UnsafeCell`, we start by performing the actions of a read access with the parent tag.
Then we push `N: SharedReadOnly` on top of the borrow stack.
This way, we grant the new pointer read-only access but make sure that it gets invalidated on the next write through any aliasing pointer (because all write-granting items are below us, and thus we test for compatibility when they get used).
There might be items left between the granting item and the one we just added, but that's okay: if any of them gets used for reading, that will be compatible with the `SharedReadOnly` according to our table above; and if any of them gets used for writing then it is important that our `SharedReadOnly` gets removed.

When we are inside an `UnsafeCell`, we will *not* perform the actions of a memory access!
Interior mutability allows all sorts of crazy aliasing, and in particular, one can call a function with signature
{% highlight rust %}
fn aliasing(refcell: &RefCell<i32>, inner: &mut i32)
{% endhighlight %}
such that `inner` points *into* `refcell`, i.e., the two pointers overlap!
Retagging `refcell` must not remove the `Unique` item associated with `inner`.

So, when retagging inside an `UnsafeCell`, we *find* the write-granting item for the parent's tag, and then we add `N: SharedReadWrite` just above it.
This grants write access, as is clearly needed for interior mutability.
We cannot add the new item at the top of the stack because that would add it *on top of* the item granting `inner`.
Any write to `inner` would remove our item from the stack!
Instead, we make our item sit just on top of its parent, reflecting the way one pointer got derived from the other.

**Retagging a raw pointer.**
Retagging for raw pointers happens only immediately after a reference-to-raw-pointer cast.
Unlike in Stacked Borrows 1, retagging depends on whether this is a `*const T` or `*mut T` pointer---this is a departure from the principle that these two types are basically the same (except for variance).
I have recently learned that the borrow checker actually handles `*mut` and `*const*` casts differently; also see [this long comment](https://github.com/rust-lang/rust/issues/56604#issuecomment-477954315).
Given that the idea of Stacked Borrows is to start with what the borrow checker does and extrapolate to a dynamic model that also encompasses raw pointers, I felt that it makes sense for now to mirror this behavior in Stacked Borrows.
This is certainly not a final decision though, and I feel we should eventually have a discussion whether we should make the borrow checker and Stacked Borrows both treat `*const T` and `*mut T` the same.

When casting to a `*mut T`, we basically behave like in the above case for inside an `UnsafeCell` behind a shared reference: we find the write-granting item for our parent's tag, and we add `Untagged: SharedReadWrite` just on top of it.
The way compatibility is defined for `SharedReadWrite`, there can be many such items next to each other on the stack, and using any one of them will not affect the others.
However, writing with a `Unique` permission further up the stack *will* invalidate all of them, reflecting the idea that when writing to a mutable reference, all raw pointers previously created from this reference become invalid.

When casting to a `*const T`, we behave just like retagging for a shared reference `&T` (including being mindful of `UnsafeCell`).
There isn't really anything that a `*const T` can do that a shared reference cannot (both in terms of aliasing and mutation), so modeling them the same way makes sense.

### Two-phase borrows

Stacked Borrows 1 [had some support for two-phase borrows]({% post_url 2018-12-26-stacked-borrows-barriers %}), but some advanced forms of two-phase borrows that used to be allowed by the borrow checker [could not be handled](https://github.com/rust-lang/rust/issues/56254).
With the additional flexibility of Stacked Borrows 2 and its departure from a strict stack discipline, it is possible to accept at least the known examples of this pattern that were previously rejected:
{% highlight rust %}
fn two_phase_overlapping1() {
    let mut x = vec![];
    let p = &x;
    x.push(p.len());
}
{% endhighlight %}
Previously, when the implicit reborrow of `x` in `x.push` got executed, that would remove the item for `p` from the stack.
This happens as part of the implicit write access that occurs when a mutable reference gets retagged.
With Stacked Borrows 2, we can do something else:
when retagging a two-phase mutable borrow, we do *not* perform the actions of a write access.
Instead, we just find the write-granting item for our parent's tag, and then add `N: Unique` just above it.
This has the consequence that on the first write access to this new pointer, everything on top of it will be removed, but until then the existing item for `p` remains on the stack and can be used just as before.

Just like with Stacked Borrows 1, we then proceed by doing a *shared* reborrow of the parent's tag from `N`.
This way, the parent pointer can be used in the ways a shared reference can (including writes, if there is interior mutability) without invalidating `N`.
This is somewhat strange because we then have "parent tag - new tag - parent tag" on the stack in that order, so we no longer properly reflect the way pointers got derived from each other.
More analysis will be needed to figure out the consequences of this.

We should also try to fully understand the effect this "weak" form of a mutable reborrow (without a virtual write access) has on the optimizations that can be performed.
Most of the cases of Stacked Borrows violations I found in the standard library are caused by the fact that even just *creating* a mutable reference asserts that it is unique and invalidates aliasing pointers, so if we could weaken that we would remove one of the most common causes of UB caused by Stacked Borrows.
On the other hand, this means that the compiler will have a harder time reordering uses of mutable references with function calls, because there are fewer cases where a mutable reference is actually assumed to be unique.

### Barriers are dead, long live protectors

With the departure from a strict stack discipline, I also had to re-think the concept of [barriers]({% post_url 2018-12-26-stacked-borrows-barriers %}).
The name was anyway terrible, so I replaced barriers by *protectors*: an `Item` actually consists not only of a `Tag` and a `Permission`, but also optionally of a "protector" represented as a `CallId` referencing some function call (i.e., some stack frame).
As long as that function call is running, the item with the protector may not be removed from the stack (or else we have UB).
This has pretty much the same effects as barriers did in Stacked Borrows 1.

## Why not a strict stack discipline?
[no-stack]: #why-not-a-strict-stack-discipline

The biggest surprise for me in designing Stacked Borrows 2 was that I was not able to enforce a strict stack discipline any more.
For retagging, that is only partially surprising; really in Stacked Borrows 1 we already added barriers below the "frozen" marker sitting next to a stack, and the `aliasing` example with the `RefCell` mentioned above only worked due to a hack that relied on reusing existing items in the middle of the stack instead of pushing a new one.
However, I had initially hoped that the rules for memory accesses would be slightly different:
my plan was that after finding the granting item, we would seek upwards through the stack and find the first incompatible item, and then remove everything starting there.
That would be a stack discipline; we would basically pop items until all items above the granting item are compatible.

Unfortunately, that would reject many reasonable programs, such as:
{% highlight rust %}
fn test(x: &mut [i32]) -> i32 { unsafe {
  let raw = x.as_mut_ptr(); // implicitly: as_mut_ptr(&mut *x)
  let _val = x[1];
  return *raw;
} }
{% endhighlight %}
The issue with this example is that when calling `as_mut_ptr`, we reborrow `x`.
So after the first line the stack would look something like (using `x` as notation for `x`'s tag)
```
[ ..., (x: Unique), (_: Unique), (Untagged: SharedReadWrite) ]
```
In other words, there would be some `Unique` item *between* the items for `x` and `raw`.
When reading from `x` in the second line, we determine that this `Unique` tag is not compatible.
This is important; we have to ensure that any access from a parent pointer invalidates `Unique` pointers---that's what makes them unique!
However, if we follow a stack discipline, that means we have to pop the `Untagged: SharedReadWrite` *and* the `_: Unique` off the stack, making the third line UB because the raw pointer got invalidated.

`test` seems like a perfectly reasonable function, and in fact this pattern is used in the standard library.
So in order to allow such code, accesses in Stacked Borrows 2 do *not* follow a stack discipline:
we remove all incompatible items above the granting item and ignore any interleaved compatible items.
As a consequence, line 2 in the above program removes `_: Unique` but keeps `Untagged: SharedReadWrite`, so line 3 is okay.

However, this means we also accept the following, even if we eventually do fully precise tracking of raw pointers:
{% highlight rust %}
fn main() { unsafe {
  let x = &mut 0;
  let raw1 = x as *mut _;
  // Stack: [ (x: Unique), (raw1: SharedReadWrite) ]

  let tmp = &mut *raw1;
  let raw2 = tmp as *mut _;
  // Stack: [ (x: Unique), (raw1: SharedReadWrite), (tmp: Unique), (raw2: SharedReadWrite) ]

  *raw1 = 1; // This will invalidate tmp, but not raw2.
  // Stack: [ (x: Unique), (raw1: SharedReadWrite), (raw2: SharedReadWrite) ]

  let _val = *raw2;
} }
{% endhighlight %}
Naively I would assume that if we tell LLVM that `tmp` is a unique pointer, it will conclude that `raw2` cannot alias with `raw1` as that was not derived from `tmp`, and hence LLVM might conclude that `_val` must be 0.
This means the program above would have to have UB.
However, in the current framework I do not see a way to make it UB without also making the reasonable example from the beginning of this section UB.
So maybe we can find a way to express to LLVM precisely what we mean, or maybe we can only exploit some of these properties in home-grown optimizations, or maybe we find a way to have UB in the second example but not in the first.
(Like, maybe we do the stack-like behavior on writes but keep the more lenient current behavior on reads? That seems annoying to implement though. Maybe a stack is just the wrong data structure, and we should use something more tree-like.)

## Conclusion

Stacked Borrows 2 shares with Stacked Borrows 1 just the general structure: pointers are tagged, and there is a per-location stack indicating which tags are allowed to perform which kinds of operations on this location.
In the new version, shared references are tagged the same way as mutable references (just with an ID to distinguish multiple references pointing to the same location); the stack keeps track of which IDs are read-only (shared) pointers and which are unique (mutable) pointers.
The only operations affected by Stacked Borrows 2 are memory accesses and the `Retag` instructions explicitly represented in the MIR; there is no longer any action on a pointer dereference.
This balances the extra complexity of the new access rules (the new implementation is actually a dozen lines shorter than the old, despite long comments).
The "stack" is unfortunately not used in a completely stack-like fashion though.

I leave it as an exercise to the reader to convince yourself that the [key properties of Stacked Borrows 1]({% post_url 2018-11-16-stacked-borrows-implementation %}#5-key-properties) still hold, and that uniqueness of mutable references is maintained even if we do a shared reborrow, as long as we keep that shared reference to ourselves.

The new model gives some wiggle-room in both the notion of which permissions are compatible with others and where exactly which permissions are added in the "stack" in a reborrow.
We could use that wiggle-room to experiment with a bunch of closely related models to see how they affect which code gets accepted and analyze their impact on optimizations.

I am quite excited by this new model!
It puts us into a good position to do more precise tracking for raw pointers, similar to what already happens for shared references (and something I was worried about in the original model).
That will be needed for compatibility with LLVM.
However, [there are still some known issues](https://github.com/rust-lang/unsafe-code-guidelines/issues?q=is%3Aissue+is%3Aopen+label%3Atopic-stacked-borrows), and also the fact that we do not actually use the "stack" as a stack indicates that maybe we should use a different structure there (potentially something more like a tree).
This is definitely not the final word, but I think it is a step in the right direction, and I am curious to see how it works out.
As usual, if you have any questions or comments, please join the discussion in the forums!