How to implement an iterator of mutable references to the values in the right edges of a Binary Search Tree?

Question

I implemented a simple Binary Search Tree in Rust (following CIS 198, it's great), and for learning I'm doing iterators that just run through the right edges.

I could not implement an iterator that gives mutable references. I tried a lot of ways, but none were accepted by Rust compiler. The code I need help is the one below (while I made a gist with the complete code here):

#[derive(Debug)]
pub struct Tree<T>(Option<Box<Node<T>>>);

#[derive(Debug)]
pub struct Node<T> {
    elem: T,
    left: Tree<T>,
    right: Tree<T>,
}

// MUTABLE BORROW STRUCT
pub struct IterMut<'a, T: 'a> {
    next: &'a mut Tree<T>,
}

// MUTABLE BORROW NEXT (I'M STUCK HERE, NOTHING WORKS)
impl<'a, T> Iterator for IterMut<'a, T> {
    type Item = &'a mut T;
    fn next(&mut self) -> Option<Self::Item> {
        // 1 try: cannot infer lifetime
        self.next.0.as_mut().map(|node| {
            self.next = &mut node.right;
            &mut node.elem
        })

        // 2 try: node.right, node.elem does not live long enough
        self.next.0.take().map(|node| {
            self.next = &mut node.right;
            &mut node.elem
        })
    }
}

Answer 1

You need to change the type of the field IterMut::next to Option<&'a mut Node<T>>:

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

impl<'a, T> Iterator for IterMut<'a, T> {
    type Item = &'a mut T;
    fn next(&mut self) -> Option<Self::Item> {
        self.next.take().map(|node| {
            self.next = node.right.0.as_mut().map(|node| &mut **node);
            &mut node.elem
        })

    }
}

You can find more useful information about the implementation of the mutable iterator for recursive data structures in the IterMut chapter of Learning Rust With Entirely Too Many Linked Lists.

Answer 2

I think you cannot split self into 2 mutable objects (one for the Item, one for self itself) without using some unsafe code.