Inferred Mutability: Reachability-based reference lifetime/mutability tracking

[duckki]

Summary of idea:

• The borrow checker shall track the reachability from one reference to another reference within the function (whenever feasible).
• Inter-procedurally, the type signature may associate the lifetime of returned reference directly to a function argument.
• This idea can be added on top of the basic Rust-style lifetime tracking (Lifetimes + Provenance proposals #338).

Expected benefits:

• Functions that don't mutate argument no longer have to take mutable references (only to return a mutable reference).
• Easier to avoid implementing both immutable and mutable getters.
• Immutable references can be upgraded to mutable ones when it's safe to do so.

(Here, I'm using Rust syntax as examples. But, it should translate to Mojo naturally.)

Problem statement

The mutability of return reference must be fixed regardless of caller's context.

Example:

struct Example {
    // This example struct is simple, but imagine a complex data structure.
    data1: i32,
    data2: i32,
}

impl Example {
    pub fn get_some_ref_mutable( &mut self, key: i32 ) -> & mut i32 {
        // This code is simple, but imagine a complex algorithm retreiving a reference.
        if key == 42 {
            &mut self.data1
        } else {
            &mut self.data2
        }
    }
}

The function get_some_ref_mutable does NOT mutate self, but self has to be declared as &mut. It just needs to return a mutable subdata reference.

Now, suppose we want the same functionality, but want to work with immutable references. One has to implement a clone of the same function like below:

    // immutable version
    pub fn get_some_ref( &self, key: i32 ) -> &i32 {
        if key == 42 {
            & self.data1
        } else {
            & self.data2
        }
    }

Moreover, the implementations can't share the other function's code, even if they are practically the same code:

    pub fn get_some_ref( &self, key: i32 ) -> &i32 {
        // You can't re-use `get_some_ref_mutable`.
        return self.get_some_ref_mutable( key ); // error: `self` is not mutable.
    }

The fundamental issue in Rust is that the mutability of returned reference is not really a concern of the function returning it. We better let the caller to resolve the mutability of returned value based on the context.

Proposal

In addition to the scope-based lifetime (like 'a, 'static), we can allow argument-based lifetime to indicate subdata relationship between the argument and the returned reference (say 'argname).

Example:

    pub fn get_some_ref_proposal( &self, key: i32 ) -> &'self i32 {
        if key == 42 {
            & self.data1
        } else {
            & self.data2
        }
    }

The return type is &'self i32 where self is the argument, indicating the returned reference is a subdata of self.

Also, 'self annotation can be elided based on static analysis. If that wasn't possible for some reason, we can fall back to scope-based lifetime tracking. For example, the fn longest from the lifetime proposal will have to fall back to the scope-based tracking, since the subdata-relationship is non-deterministic.

Motivating use-case - Upgrading immutable reference to mutable:

    pub fn example_use_case( &mut self, key: i32) {
        let ref1 = self.get_some_ref_proposal( key );
        let mut ref2 = ref1; // upgrade to a mutable reference
        *ref2 = 5;
    }
    // Note: The code could be shortened as `let mut ref1 = ...; *ref1 = 5`.

It's safe because ref1 is the only immutable borrow of self, let mut ref2 = ref1; is the last use of ref1 and self is mutable. To make this work, the borrow checker must know that ref1 is a subdata reference of self across the function call.

Rust does not track whether a reference is a subdata of another, as stated in the Rustonomicon. I propose to do exactly that. My hunch is that it will work for many use cases. When dataflow can't conclude the relationship, we can fall back to the Rust-style lifetime tracking.

[nmsmith]

As I understand it, the crux of your proposal is that the provenance of a reference should be declared by explicitly listing which variable(s) it was derived from. This allows a function to be "generic" over the mutability of the reference it returns.

I agree it would be wonderful if Mojo had this capability. In fact, it's more-or-less required if Mojo wants to avoid replicating the complexity and boilerplate associated with Rust's borrowing system. A lot of Rust APIs are duplicated twice—once for each mutability choice. For example, Rust's equivalent of Python's __getitem__ protocol is split over two traits: Index and IndexMut, which have to be separately implemented by each data type, leading to a lot of code duplication and boilerplate.

I believe the fundamental problem is that in Rust, mutability is associated with the type of an expression/value, whereas in order to be generic over mutability, you really want mutability to be associated with a name. In other words, instead of saying "this value is immutable", we want to be able to say "through this name, no mutations will be performed or observed". Naturally then, when returning a reference from a function, you don't specify its mutability directly. Instead, you specify the set of memory locations that the reference might point to, and then the caller can naturally figure out whether it is possible to perform mutations upon it.

This is all very subtle and I'm still trying to distill precisely what the right model is.

In the original discussion thread on lifetimes for Mojo, I proposed a solution that is vaguely similar to @duckki's proposal above. It was just a rough idea and requires more consideration, but I think this is all worth investigating with care and precision. We might be able to end up with a semantics for borrowing that is simpler and more flexible than Rust's.

[duckki]

Today I had a chance to learn about Rust's Polonius project as someone pointed it out on Discord . It has very similar mechanics as my proposal (except for the mutability-generics). This 2019 talk on YouTube was helpful to get the idea.

Polonius's lifetime tracking is not just based on "location" (on CFG node). It sounded more like a dataflow analysis of loans propagating from borrow sites to variables. So, it can answer the more fine-grained question of "Does variable A borrows loan L at a given location?", instead of "Should variable A be live at a given location?".

BTW, the speaker of the talk also suggested a similar reachability tracking that I proposed above (at 31:51 mark).

struct Message {
    buffer: Vec<String>,
    slice: &'buffer [u8], // "borrowed from the field `buffer`"
}

It was just a future possibility back then and I'm not sure what's actually implemented today.

Going one step further, Polonius allows us to find the origin of loans. I think we can take advantage of that for enabling mutability upgrade.

[duckki]

Here's simpler example with a reasoning behind this idea.

Motivation:

fn motivation() {
    let mut v = Example { data1: 1, data2: 2 };
    let x1 = &v;
    let x2 = &x1.data1;
    *x2 = 4; // not allowed
}

We can work around it by executing an identical mutable access:

fn workaround() {
    let mut v = Example { data1: 1, data2: 2 };
    let x1 = &v;
    let x2 = &x1.data1;
    // *x2 = 4; // not allowed
    let x3 = &mut v.data1; // x2 and x3 are the same reference.
    *x3 = 4;
}

But, that's silly since we are recalculating the same reference.

Logically, I think we need two new lifetime/ownership rules:

Rule 1: Allow the only live immutable borrow to become mutable, if the origin of borrow is mutable.
Rule 2: Allow lifetimes of two immutable references to invert, if they share the same origin transitively.

The Rule 1 means that we can pretend the original borrow was a mutable one as long as that's safe.

The Rule 2 needs more explanation. In the example above, v lends to x1 and x1.data1 lends to x2. So, the usual lifetime rules enforce that v > x1 > x2. However, when the original loan structure is v --> x1 -> x1.data1 --> x2 (both loans immutable), we can safely flatten the loan structure like v --> x1 and v.data1 --> x2, which breaks the lifetime constraint between x1 and x2.

Putting together. In the example above, the rule 2 applies, allowing x1 to die first. Then the rule 1 applies, upgrading x2 to become mutable.

[duckki]

I've posted this idea on Rust community (here).

It's clear that mutability upgrade won't work without establishing "dependency" relationship between args and returned values. Rust's lifetime is merely "lives as long as" relation and won't be strong enough to infer exclusivity.

Also, in Rust, if a struct has an immutable field, any subdata reference to the struct can't be mutability-upgraded, since that may violate the field's immutability. However, this won't be an issue in Mojo.

Here's my blog post on the details of this issue and the proof why mutability upgrade is safe.

[lattner]

it is still /suuuuuper/ early, but Mojo 0.7 just shipped the earliest support for safe references, and yes it does support parametric mutability. The support in this release isn't usable much at all, but you can express some things like this:

# Immutable
struct S[lifetime: ImmLifetime]:
    var s: Reference[String, __mlir_attr.`0: i1`, lifetime] # zero is immut

or

# Mutable
struct S[lifetime: MutLifetime]:
    var s: Reference[String, __mlir_attr.`1: i1`, lifetime] # one is mut

or

# Either
struct S[is_mutable: __mlir_type.i1, lifetime: AnyLifetime[is_mutable].type]:
    var s: Reference[String, is_mutable, lifetime]

We have a number of things that are coming that will clean this up, including parametric aliases (fixing the AnyLifetime[is_mutable].type thing and eliminating the need for ImmLifetime vs MutLifetime type aliases). Support for inferring parameters from other parameters will eliminate the redundant need for is_mutable and the “mlir i1” nonsense (among many other things, but at least it is a start.

Thank you for advocating for parametric mutability!