Returning a mutable reference to a field of a mutably referenced struct

Question

I am still quite new to the advanced topics in rust, but for context, I am trying to implement a generic quadtree in rust.
With the method find_mut(&mut self,x,y) I want to traverse the quadtree to find the lowermost subtree containing that coordinate and return a mutable reference of it.

Quadtree Struct

pub struct QuadTree<T> {
    x: i32,
    y: i32,
    dx: i32,
    dy: i32,
    leaf: bool,
    subtrees: [Option<Box<Self>>; 4],
    value: Option<T>,
}

Methods and functions

fn find_mut(&mut self, x: i32, y: i32) -> Result<&mut Self, &mut Self> {
    let mut current = self;
    loop {
        // If we've arrived at a leaf return it as Ok
        if current.leaf { // First borrow occurs here for some reason
            return Ok(current);
        }
        // Getting the subtree that contains our coordinate
        match current.mut_subtree_at(x, y) {
            // Go a level deeper
            Some(child) => current = child,
            // Return an Err containing the lowest subtree that is sadly not a leaf
            None => return Err(current),
        }
    }
}

fn subtree_id<'a>(&'a self, x: i32, y: i32) -> usize {
    let mut child_id = 0;
    if x >= self.x {
        child_id += 1;
    }
    if y >= self.y {
        child_id += 2;
    }
    child_id
}

#[inline(always)]
fn mut_subtree_at(&mut self, x: i32, y: i32) -> Option<&mut Self> {
    self.subtrees[self.subtree_id(x, y)].as_deref_mut()
}

Error

error[E0499]: cannot borrow `*current` as mutable more than once at a time
   --> src/quadtree.rs:128:36
    |
115 |     fn find_mut(&mut self, x: i32, y: i32) -> Result<&mut Self, &mut Self> {
    |                 - let's call the lifetime of this reference `'1`
...
121 |                 return Ok(current);
    |                        ----------- returning this value requires that `*current` is borrowed for `'1`
...
124 |             match current.mut_subtree_at(x, y) {
    |                   ---------------------------- first mutable borrow occurs here
...
128 |                 None => return Err(current),
    |                                    ^^^^^^^ second mutable borrow occurs here

How would you approach this problem. I am missing something about the way borrowing mutable references and lifetimes work?

Answer 1

While not ideal, here is a recursive version that compiles, but requires querying mut_subtree_at twice:

pub struct QuadTree<T> {
    x: i32,
    y: i32,
    dx: i32,
    dy: i32,
    leaf: bool,
    subtrees: [Option<Box<Self>>; 4],
    value: Option<T>,
}

impl<T> QuadTree<T> {
    fn find_mut(&mut self, x: i32, y: i32) -> Result<&mut Self, &mut Self> {
        if self.leaf {
            Ok(self)
        } else {
            if self.mut_subtree_at(x, y).is_some() {
                // Go a level deeper
                self.mut_subtree_at(x, y).unwrap().find_mut(x, y)
            } else {
                // Return an Err containing the lowest subtree that is sadly not a leaf
                Err(self)
            }
        }
    }

    fn subtree_id<'a>(&'a self, x: i32, y: i32) -> usize {
        let mut child_id = 0;
        if x >= self.x {
            child_id += 1;
        }
        if y >= self.y {
            child_id += 2;
        }
        child_id
    }

    #[inline(always)]
    fn mut_subtree_at(&mut self, x: i32, y: i32) -> Option<&mut Self> {
        self.subtrees[self.subtree_id(x, y)].as_deref_mut()
    }
}

To my understanding this works because .is_some() returns a bool, which is an owned value and therefore Rust can prove that at the time of Err(self) no reference is held to self any more.

It seems that the same principle also works for the iterative solution:

pub struct QuadTree<T> {
    x: i32,
    y: i32,
    dx: i32,
    dy: i32,
    leaf: bool,
    subtrees: [Option<Box<Self>>; 4],
    value: Option<T>,
}

impl<T> QuadTree<T> {
    fn find_mut(&mut self, x: i32, y: i32) -> Result<&mut Self, &mut Self> {
        let mut current = self;
        loop {
            // If we've arrived at a leaf return it as Ok
            if current.leaf {
                return Ok(current);
            }
            // Getting the subtree that contains our coordinate
            if current.mut_subtree_at(x, y).is_some() {
                // Go a level deeper
                current = current.mut_subtree_at(x, y).unwrap()
            } else {
                // Return an Err containing the lowest subtree that is sadly not a leaf
                return Err(current);
            }
        }
    }

    fn subtree_id<'a>(&'a self, x: i32, y: i32) -> usize {
        let mut child_id = 0;
        if x >= self.x {
            child_id += 1;
        }
        if y >= self.y {
            child_id += 2;
        }
        child_id
    }

    #[inline(always)]
    fn mut_subtree_at(&mut self, x: i32, y: i32) -> Option<&mut Self> {
        self.subtrees[self.subtree_id(x, y)].as_deref_mut()
    }
}

---

For more performance (meaning, if you don't want to call mut_subtree_at twice), you would have to override the borrow checker as this is obviously a false positive.

Luckily, there is the polonius-the-crab crate that is written for specifically the problem you ran into and hides the unsafe code in a safe and reliable manner.

Here is my first working version using polonius-the-crab:

use ::polonius_the_crab::prelude::*;
use polonius_the_crab::WithLifetime;

pub struct QuadTree<T> {
    x: i32,
    y: i32,
    dx: i32,
    dy: i32,
    leaf: bool,
    subtrees: [Option<Box<Self>>; 4],
    value: Option<T>,
}

impl<T> QuadTree<T> {
    fn find_mut(&mut self, x: i32, y: i32) -> Result<&mut Self, &mut Self> {
        type SelfRef<K> = dyn for<'lt> WithLifetime<'lt, T = &'lt mut QuadTree<K>>;

        let mut current = self;
        loop {
            // If we've arrived at a leaf return it as Ok
            if current.leaf {
                return Ok(current);
            }
            // Getting the subtree that contains our coordinate
            match polonius::<SelfRef<T>, _, _, _>(current, |current| {
                current.mut_subtree_at(x, y).ok_or(())
            }) {
                Ok(child) => {
                    // Go a level deeper
                    current = child;
                }
                Err((current, ())) => {
                    // Return an Err containing the lowest subtree that is sadly not a leaf
                    return Err(current);
                }
            }
        }
    }

    fn subtree_id<'a>(&'a self, x: i32, y: i32) -> usize {
        let mut child_id = 0;
        if x >= self.x {
            child_id += 1;
        }
        if y >= self.y {
            child_id += 2;
        }
        child_id
    }

    #[inline(always)]
    fn mut_subtree_at(&mut self, x: i32, y: i32) -> Option<&mut Self> {
        self.subtrees[self.subtree_id(x, y)].as_deref_mut()
    }
}

or this version, which is a little easier to understand as it uses polonius' macros:

use ::polonius_the_crab::prelude::*;

pub struct QuadTree<T> {
    x: i32,
    y: i32,
    dx: i32,
    dy: i32,
    leaf: bool,
    subtrees: [Option<Box<Self>>; 4],
    value: Option<T>,
}

impl<T> QuadTree<T> {
    pub fn find_mut(&mut self, x: i32, y: i32) -> Result<&mut Self, &mut Self> {
        let mut current = self;
        while !current.leaf {
            // Getting the subtree that contains our coordinate.
            // If no subtree exists, return Err(current).
            current = current.mut_subtree_at(x, y)?;
        }
        // If we are at a leaf node with the coordinate, success!
        Ok(current)
    }

    fn subtree_id<'a>(&'a self, x: i32, y: i32) -> usize {
        let mut child_id = 0;
        if x >= self.x {
            child_id += 1;
        }
        if y >= self.y {
            child_id += 2;
        }
        child_id
    }

    #[inline(always)]
    fn mut_subtree_at(&mut self, x: i32, y: i32) -> Result<&mut Self, &mut Self> {
        let mut current = self;

        polonius!(
            |current| -> Result<&'polonius mut Self, &'polonius mut Self> {
                if let Some(child) = current.subtrees[current.subtree_id(x, y)].as_deref_mut() {
                    polonius_return!(Ok(child))
                }
            }
        );

        // Return the ownership of `self` back through the `Err()` value.
        // This in conjunction with the `polonius!()` macro resolves the
        // ownership problem.
        Err(current)
    }
}