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?

This is a known borrow-checker limitation. https://github.com/rust-lang/rust/issues/54663 — PitaJ, Aug 09 '22 at 20:41
Any ideas on how I would get around this limitation and still get the functionality of finding my subtree and returning a mutable reference to it? — RibsGrowBack, Aug 10 '22 at 08:05
This might help: https://docs.rs/polonius-the-crab/latest/polonius_the_crab/ — PitaJ, Aug 10 '22 at 13:02
That's a weirdly formatted compiler error, how did you compile it? — Finomnis, Aug 10 '22 at 16:36
Ok lol, I just got it :D please don't copy the errors from your IDE, but instead copy the output of `cargo check`. — Finomnis, Aug 10 '22 at 16:44
@PitaJ While that might help I think it would be nice to figure out how to do it differently. But thank you — RibsGrowBack, Aug 10 '22 at 18:42
@Finomnis Those seemed awfully verbose and full of whitespace but here you go — RibsGrowBack, Aug 10 '22 at 18:45
@RibsGrowBack Yes they are verbose, but at least we can understand where they point at. Line numbers are sadly completely meaningless if you don't post the actual code, but separate it in several pieces. — Finomnis, Aug 10 '22 at 19:04

Finomnis · Accepted Answer · 2022-08-10T21:23:04.430

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)
    }
}

I also learned a lot from it, so thanks for the nice case study ;) — Finomnis, Aug 10 '22 at 21:12

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

Quadtree Struct

Methods and functions

Error

1 Answers1