Trying to understand GATs (the lending iterator example)

Hello! Could someone please explain how the flexibility around lifetimes here has helped us write this nice implementation below? Would really appreciate an explanation covering all the lifetime details! Thanks :smiley:

use ::core::mem;

pub trait LendingIterator {
    type Item<'this>
    where
        Self: 'this;
    fn next(&mut self) -> Option<Self::Item<'_>>;
}

pub fn windows_mut<T, const WINDOW_SIZE: usize>(slice: &mut [T]) -> WindowsMut<'_, T, WINDOW_SIZE> {
    assert_ne!(WINDOW_SIZE, 0);
    WindowsMut { slice, first: true }
}

pub struct WindowsMut<'a, T, const WINDOW_SIZE: usize> {
    slice: &'a mut [T],
    first: bool,
}

impl<'a, T, const WINDOW_SIZE: usize> LendingIterator for WindowsMut<'a, T, WINDOW_SIZE> {
    type Item<'this>
        = &'this mut [T; WINDOW_SIZE]
    where
        'a: 'this;

    fn next(&mut self) -> Option<Self::Item<'_>> {
        if !self.first {
            self.slice = &mut mem::take(&mut self.slice)[1..];
        }
        self.first = false;

        Some(self.slice.get_mut(..WINDOW_SIZE)?.try_into().unwrap())
    }
}

fn main() {
    let mut slice = [1usize, 2, 3, 4];
    let mut wm: WindowsMut<'_, usize, 1> = windows_mut(&mut slice);

    while let Some(i) = wm.next() {
        *i = [0; 1];
    }
}

Also, this doesn't compile...could someone please explain why? Thanks!

impl<'a, T, const WINDOW_SIZE: usize> Iterator for WindowsMut<'a, T, WINDOW_SIZE> {
    type Item = &'a mut [T; WINDOW_SIZE];

    fn next(&mut self) -> Option<Self::Item> {
        if !self.first {
            self.slice = &mut mem::take(&mut self.slice)[1..];
        }
        self.first = false;

        Some(self.slice.get_mut(..WINDOW_SIZE)?.try_into().unwrap())
    }
}

The above lending iterator's next signature is short for this:

fn next<'x>(&'x mut self) -> Option<&'x mut [T; WINDOW_SIZE]>

The regular iterator's next signature that doesn't compile is short for this:

fn next<'x>(&'x mut self) -> Option<&'a mut [T; WINDOW_SIZE]>

The first function connects the return's borrow with the &mut self borrow created when this function is called. The second function connects the return with the borrow inside WindowsMut, which must exist before and after the WindowsMut containing it.

Both functions are giving permission to the caller to hold onto the return value for some amount of time: 'x for the first and 'a for the second. 'x is valid until self is used again, which will probably happen when calling next a second time. 'a is valid until the slice parameter of the windows_mut function is used again, which is a much larger scope.

The second function breaks down because it can't produce a slice with the 'a lifetime. get_mut operates on &mut [T] (any lifetime), but we actually have &'x mut &'a mut [T]. Converting this by reborrowing, which happens automatically on method calls, can only give us &'x mut [T]. This is also why the line above uses take instead of just self.slice = &mut self.slice[1..]. take is operating on &'x mut &'a mut [T] and produces the inner lifetime 'a, whereas only indexing would work the same way as get_mut and produces the outer lifetime 'x.

Now, the first function is a perfectly normal function to write in Rust. You could even have it in a trait. You could even have a trait that's generic over the type inside the reference:

trait RefMutIterator {
    type Item;
    fn next(&mut self) -> Option<&'_ mut Self::Item>;
}

However, this limits you to returning types that are fully contained in a plain mutable reference. You couldn't return something like (&mut K, &mut V) or Struct<'_> using that lifetime. That's what GATs enable. They let the trait specify a lifetime, and the implementer to decide where in the type the lifetime goes.

The other question is why would you need lifetimes in iterator items? The mutable windows iterator is one reason you'd want to do this: since the borrow expires before the next call to next, you can produce another mutable reference to the same elements as the previous call. Another reason is when the items are owned by the iterator itself. For example, if you have a compressed string, and want to iterate over uncompressed lines, you could implement that while reusing your allocation:

use std::io::Cursor;

pub struct CompressedLines {
    string: Cursor<Vec<u8>>,
    line: Vec<u8>,
}

fn decompress_line(_string: &mut Cursor<Vec<u8>>, _buffer: &mut Vec<u8>) {
    todo!()
}

impl LendingIterator for CompressedLines {
    type Item<'a> = &'a [u8];

    fn next(&mut self) -> Option<Self::Item<'_>> {
        self.line.clear();
        decompress_line(&mut self.string, &mut self.line);
        if self.line.is_empty() { None } else { Some(&self.line) }
    }
}

thanks a lot for providing an elaborate explanation :slight_smile:

could you please explain this? the borrow is created after calling windows_mut, how can it exist before WindosMut?

the sliding windows can overlap each other, so if you could call Iterator::next() multiple times, you would create &mut references to the same elements, which is UB.

I see I see.

fn main() {
    let mut slice = [1usize, 2, 3, 4];
    let mut wm: WindowsMut<'_, usize, 1> = windows_mut(&mut slice);

    while let Some(i) = (&mut wm).next() {
        *i = [0; 1];
    }
}

If I have understood it correctly: In the case of lending iterator as the lifetime of the mutable slice (given out by the next function) is tied to &mut self the borrow of the given out slice actually ends before each call to next as each new call takes in a reborrow of &mut self. is that correct?

but in the case of normal Iterator we sort of extended 'a by saying that all the mutable slices given out by the next function should live as long as 'a (which is the lifetime inside WindowsMut) and therefore slices could have overlapping mutable borrows. is that correct?

also, a general question on NLL:

    let mut slice = [1, 2, 3, 4, 5];
    let m1 = &mut slice;
    let m2 = &mut slice;
    // println!("{:?}", m1);

m1 creates a new reborrow, m2 creates a new reborrow; as long as two borrows (mutable-mutabe, mutable-immutable) do not interfere (which could be done by uncommenting the line) we are okay, right?

if that's correct then could you please tell why this is okay?

fn main() {
    let mut slice = [1usize, 2, 3, 4];
    let mut wm: WindowsMut<'_, usize, 1> = windows_mut(&mut slice);

    while let Some(i) = (&mut wm).next() {
        *i = [0; 1];
    }

    println!("{:?}", wm);
}

we are using the same mutable borrow created above but this doesn't error out....why?

yes.

today's iterators alllow the yielded items to be stored all at the same time, e.g. you can collect the items into a container such as Vec. this requires the items to be either shared, or disjointly exclusive.

lending iterators give out items one at a time, since they borrows the lending iterators themselves exclusively.

It took a bit of time to finish writing this reply, so it's somewhat redundant with the others.


Let's start with the non-lending iterator.

pub trait Iterator {
    type Item;
    fn next<'this>(&'this mut self) -> Option<Self::Item>;
}

Item is defined outside of the context of 'this. It's impossible for 'this to be part of the Item type in any implementation. That means that it's impossible for it to contain a (lifetime-tracked) borrow derived from &'this mut self in the return of next. And it has to be defined everywhere that the implementation is valid.

The language exploits this in various ways. For example:

fn example<I: Iterator<Item: std::fmt::Debug>>(mut iter: I) -> I::Item {
    // Exclusively borrow `iter` to get an item.
    let one = iter.next();
    // And again... this exclusive borrow does not invalidate `one`.
    let two = iter.next();
    // So they're both usable here.
    println!("{:?}", (&one, &two));
    // And we can even return an item even though the iterator will
    // be destructed before this function returns.
    one.unwrap()
}

If Items were able to contain a borrow derived from &mut self, this couldn't compile. (Calling next or destructing iter would invalidate any items you previously got from calling next.)


We can have some exclusively-borrowing iterators with the standard Iterator, so why can't we have WindowsMut? Consider a window size of 2.

let arr = [0, 1, 2];
let mut iter = windows_mut::<_, 2>(&mut arr);
let one = iter.next().unwrap(); // &mut arr[0..2] ([0, 1])
let two = iter.next().unwrap(); // &mut arr[1..3] ([1, 2])

If one and two were able to be usable at the same time, we would have two active &muts to the same memory (where 1 is stored). This is a violation of one of Rust's foundational axioms -- &mut are exclusive -- and would thus be UB.

Therefore, whatever the specific error ends up being, a WindowsMut implementation with the standard Iterator trait must not be allowed to safely compile.[1]

With a lending iterator, when the Item<'_> ends up containing a borrow derived from &mut self, you can never hold on to more than one item at a time. Calling next again requires a fresh exclusive borrow of the iterator, which -- unlike Iterator -- invalidates any outstanding borrowing items.

(The exclusively-borrowing iterators that do work with Iterator never hand out references to overlapping memory.)


@drewtato covered this well -- you can't get a &'long mut T out of a &'short mut &'long mut T.

The attempt at making WindowsMut work with the standard Iterator trait is actually an illustration of why this has to be enforced. If it compiled, callers could get overlapping &mut _s!

At the type system level, &'short mut &'long mut T implies 'long: 'short.[2] This means that uses of 'short keep the 'long borrow alive... but not the other way around. If you could get a &'long mut T out from behind a &'short mut &'long mut T, the 'short can expire without invalidating your new &'long mut T -- and without invalidating the original &'long mut T. That is UB if the two &'long mut T overlap.

In the mutable slice iterator implementation I linked, we have to jump through some hoops to ensure there is no overlapping memory. But a proper WindowsMut implementation requires overlapping memory, so there is no sound workaround using Iterator.

A little more detail
        Some(self.slice.get_mut(..WINDOW_SIZE)?.try_into().unwrap())

To get a &'a mut [..] from get_mut, you would need to call get_mut on the slice when it was not behind &mut self, to avoid a &'short mut &'long mut _ situation.

        let slice = match self.first {
            true => mem::take(&mut self.slice),
            false => &mut mem::take(&mut self.slice)[1..],
        };
        self.first = false;
        let window: &'a mut [T; WINDOW_SIZE] = slice.get_mut(..WINDOW_SIZE)?.try_into().unwrap();
        // self.slice = slice;
        Some(window)

But there's no way to properly restore self.slice since the portion you want to keep overlaps with the portion you want to return.


Yes.

The returned slice keeps *self -- the WindowsMut itself -- exclusively borrowed. You can't call next again until it's gone.

All the slices given out have the same borrow duration, and they do not keep the WindowsMut itself borrowed. Perhaps think of it like you had a Vec<&'a mut [T]> and you popped one out. The Vec itself doesn't stay borrowed after the pop even if you keep using the item (you can pop again while still using the item). It just gave you a borrow from elsewhere that it was holding onto.

You know[3] the slices do not overlap, because they are exclusive (&mut _) and you can use them all at the same time (e.g. you can always collect a standard Iterator).

Do you mean something like this, which also compiles?

    let mut slice = [1usize, 2, 3, 4];
    let mut wm: WindowsMut<'_, usize, 1> = windows_mut(&mut slice);
    let mr = &mut wm; // <--- "the same mutable borrow"

    while let Some(i) = mr.next() {
        *i = [0; 1];
    }

    println!("{:?}", mr);
    println!("{:?}", wm);

Every call to mr.next() reborrows *mr for some duration that ends before the next call. If we unroll the loop a bit:

    let mr = &mut wm; //-------------------------+
    //                                           |
    if let Some(i) = mr.next() { //-----+        :
        *i = [0; 1]; //                 |        :
        //----no more uses of i---------+        :
    } //                                         |
    if let Some(i) = mr.next() { //-----+        :
        *i = [0; 1]; //                 |        :
        //----no more uses of i---------+        :
    } //                                         |
    //                                           |
    println!("{:?}", mr); //                     |
    //----no more uses of mr---------------------+

It's harder to visualize the loop, but basically the compiler sees that there's a point in the control flow at the end of the loop body where the borrow doesn't need to be active any more, so it doesn't propagate to the subsequent call to next.

If you make the outer lifetime of mr nameable ('r), you can indirectly see that you must not be getting back &'r mut [T]s.

fn ex<'r, 'a>(mr: &'r mut WindowsMut<'a, usize, 1>) {
    while let Some(i) = mr.next() {
        // Uncomment this line to see an error
        // let i: &'r mut _ = i;
        *i = [0; 1];
    }

    println!("{:?}", mr);
}

  1. And an unsafe implementation would be unsound. â†Šī¸Ž

  2. and T: 'long â†Šī¸Ž

  3. assuming the iterator is sound â†Šī¸Ž

Understood. thanks a lot!

thanks a lot for this detailed explanation!

yeah, I meant this :sweat_smile:

isn't this (mr) overlapping with the new mutable borrows we are giving out at each call to next?
or the compiler can somehow figure out that each call to next gives out a memory location that no longer exists behind the mutable reference mr?

Yes, but mr isn't usable while any reborrow through mr is active. (In the quoted code, we create a reborrow to call next and the return value we assign to i keeps the reborrow active until the last use of i.)

If it didn't work that way, reborrows would always have to be the same duration as the original borrow, which would mean for example this would fail.

Although everyone says "no overlapping with &mut", it's really "no overlapping with an active &mut".