Is the stack OK? – @errmayank

After Reflecting on A Bad Stack it’s time to improve it. I’ve finished the second chapter An Ok Stack and here are the things we did:

Improved our existing types and made them all pointy.
Added peek functionality.
Made our List iterable by implementing the Iterator trait.

Improving our existing types

First, we converted our Link enum:

enum Link {
    Empty,
    More(Box<Node>),
}

into an Option type and used the newtype pattern to make it cleaner:

type Link = Option<Box<Node>>;

The Option type also comes with a bunch of useful methods, like take() and map() which we used to update our pop() method that now looks like:

pub fn pop(&mut self) -> Option<i32> {
    self.head.take().map(|node| {
        self.head = node.next;
        node.element
    })
}

take() method takes the value out of self.head, leaving it as None.
map() runs the closure if the node is Some, otherwise it just returns None. Inside the closure we can access variables from the outer scope allowing us to set self.head to the next node in the list.

Making our types all pointy

struct List<T> {
    head: Link<T>,
}

type Link<T> = Option<Box<Node<T>>>;

struct Node<T> {
    element: T,
    next: Link<T>,
}

We also updated all our impls and now our list is completely generic and can be used for any arbitrary type T instead of just being limited to i32.

For peek we need to return a reference to the element but if we try to directly use map() it takes ownership of the value inside the Option, which will move the Box<Node> out of self.head and we don’t want that.

Instead we used as_ref(), which converts Option<Box<Node>> into Option<&Box<Node>>. This creates a new Option containing a reference to the Box, rather than the Box itself.

pub fn peek(&self) -> Option<&T> {
    self.head.as_ref().map(|node| &node.element)
}

Iterating through…

To make our List iterable, we need to implement the Iterator trait which looks like this:

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

In Rust, there are mainly three forms of iteration for any collection:

into_iter() which iterates over T and is the easiest to implement.
iter() iterates over &T and requires understanding of lifetimes.
iter_mut() iterates over &mut T and is almost identical to iter(), at least at our level coz compiler does all the heavy lifting.

IntoIter

The into_iter() method consumes the List and returns owned values.

To implement this we used the existing pop method, which will be called on each iteration. The .0 syntax is how we access fields in a tuple struct.

struct IntoIter<T>(List<T>);

impl<T> List<T> {
    pub fn into_iter(self) -> IntoIter<T> {
        IntoIter(self)
    }
}

impl<T> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.pop()
    }
}

Iter

The main idea here is that we need to hold a pointer to the current node which we want to yield next.

For this we need to tell the compiler that all our borrows are valid, so we need to add lifetime annotations to our functions and types.

First we added lifetime to the Iter struct, which is generic over some lifetime 'a.

struct Iter<'a, T> {
    next: Option<&'a Node<T>>,
}

Now we need to specify a lifetime for the borrow that creates the iterator. When we call iter(), we borrow the list with &'a self and the returned iterator contains references with the same lifetime 'a. This ties the iterator’s lifetime to the list’s borrow.

Since the references in the iterator must remain valid for the entire lifetime 'a, the list must stay borrowed as long as the iterator exists.

As for the next(), we need to go from a boxed node to a reference to the node inside the box.

impl<T> List<T> {
    pub fn iter<'a>(&'a self) -> Iter<'a, T> {
        Iter {
            next: self.head.as_ref().map(|node| &**node),
        }
    }
}

What’s happening here?

Our self.head is of type Option<Box<Node<T>>>.
We convert it to a reference using as_ref(), now it becomes Option<&Box<Node<T>>>.

Then we need to unbox it, which is handled inside the map() closure. Breaking down the weird &**node syntax:

node:     &Box<Node<T>>     Reference to the Box
*node:    Box<Node<T>>      The Box value itself
**node:   Node<T>           Actual Node on the heap
&**node:  &Node<T>          Reference directly to the Node

The same could also be achieved by as_deref() method, which pretty much does the same thing under the hood.

Iter {
    next: self.head.as_deref(),
}

Last, we need to add lifetime annotations to the Iterator trait implementation since Iter has one. We’ll use as_deref() to make our lives easier.

impl<'a, T> Iterator for Iter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        self.next.map(|node| {
            self.next = node.next.as_deref();
            &node.element
        })
    }
}

IterMut

Implementing IterMut is quite similar to Iter.

struct IterMut<'a, T> {
    next: Option<&'a mut Node<T>>,
}

We’ll use anonymous lifetime annotation here, same could be done for Iter as well. Rust’s lifetime elision rules automatically connect the lifetime of &mut self to the lifetime needed by IterMut<'_, T>

impl<T> List<T> {
    pub fn iter_mut(&mut self) -> IterMut<'_, T> {
        IterMut {
            next: self.head.as_deref_mut(),
        }
    }
}

Now we need to implement the Iterator trait.

But using map() directly, like we did in Iter, won’t work here coz shared references implement the Copy trait, which means that it is safe to copy the reference by just copying its bits.

And when we implemented the Iterator trait for Iter, we did self.next.map(|node| ...) which copied the reference of self.next to node.

The original self.next still had the reference and you could use both original and the copy. But you can’t copy a mutable reference coz that would create two mutable references to the same location which is not allowed.

So instead:

We use take() to move the value out of self.next, which gets replaced with None, leaving no lingering references.
Now node owns that reference exclusively.
We can create a new mutable reference to the next node without conflict since each node gets borrowed only once.

impl<'a, T> Iterator for IterMut<'a, T> {
    type Item = &'a mut T;

    fn next<'b>(&'b mut self) -> Option<&'a mut T> {
        self.next.take().map(|node| {
            self.next = node.next.as_deref_mut();
            &mut node.element
        })
    }
}

Adding explicit lifetime annotations makes it clear that each time next() is called, we borrow the iterator self for lifetime 'b just during the call and return a mutable reference with lifetime 'a.

Borrow of self ends when the call finishes, but the returned reference continues to exist, having the same lifetime as the original borrow of the list which is 'a.

Rust understands that we have exclusive access to the mutable reference and there’s no way to “go back”, so we can safely get a mutable reference to every element of the list.