Apr 30, 2019 • Research, Rust • Edits • Permalink

Stacked Borrows 2

Recently, I have significantly updated 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. I assume some familiarity with the prior version and will not explain everything from scratch.

Stacked Borrows: Table of Contents

Before we start, let me give a brief overview over the posts on Stacked Borrows that I have written so far. I didn’t plan this out in advance, so things are a bit more messy than I would like.

The problem

The problem I wanted to solve with Stacked Borrows 2 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 is not harmful to grant anyone read access. However, as @arielby noted, 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:

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 `*a`. Frozen locations can be read by any raw pointer.
    let _val = unsafe { *y; }; // Hence, this legal in Stacked Borrows.
    *a = 2; // But we might want to drop the earlier `*a = 1` because it gets overwritten!
    _val
}

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 read *a! This violates the idea that a is a unique pointer: we’d like to remove the first write (*a = 1) because it gets overwritten later and there is no read through a in the mean time. The issue is, there is an aliasing read through y, so removing the seemingly redundant write would change the return value of foo from 1 to 0. 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:

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,
}

There is no longer a separate “frozen” state; keeping shared references read-only is now handled by the borrow stack itself:

pub struct Stack {
    borrows: Vec<Item>
}

The tag is also simpler than it was before: there are no longer separate tags for mutable and shared references.

pub type PtrId = u64;
pub enum Tag {
    Tagged(PtrId),
    Untagged,
}

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. This is the item that grants pointers tagged tag the permission to perform the given access, hence the name (the same concept used to be called “matching item” in Stacked Borrows 1). 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 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 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.)

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:

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:

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!
    let _val = unsafe { *y };

    *a = 2;

    _val
}

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 Untagged 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.)

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 reads from 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 like y 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

fn aliasing(refcell: &RefCell<i32>, inner: &mut i32)

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. 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, but some advanced forms of two-phase borrows that used to be allowed by the borrow checker could not be handled. 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:

fn two_phase_overlapping1() {
    let mut x = vec![];
    let p = &x;
    x.push(p.len());
}

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. 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?

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:

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;
} }

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:

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;
} }

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 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, 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!

Posted on Ralf's Ramblings on Apr 30, 2019.
Comments? Drop me a mail or leave a note in the forum!