Help with mutable borrows more than one time

help
1

I need to insert a node in a tree in a way roughly like below:

pub struct Node {
    data: Vec<u64>,
    nodes: Vec<Box<Node>>,
}

pub struct Tree {
    root: Box<Node>,
}

impl Tree {
    pub fn insert(&mut self, data: Vec<u64>) {
        let mut current = &mut self.root;
        let mut index = 0;

        loop {
            if let Some(child) = current.nodes.get_mut(index) {
                index += 1;
                current = child;
            } else {
                let node = Node {
                    data,
                    nodes: Vec::default(),
                };

                current.nodes.push(Box::new(node));
                return;
            }
        }
    }
}

but I get the error:

error[E0499]: cannot borrow `current.nodes` as mutable more than once at a time                                                                                                                                                                                                                                                 
  --> src/tree.rs:25:17                                                                                                                                                                                                                                                                                                         
   |                                                                                                                                                                                                                                                                                                                            
16 |             if let Some(child) = current.nodes.get_mut(index) {                                                                                                                                                                                                                                                            
   |                                  ------------- first mutable borrow occurs here                                                                                                                                                                                                                                            
...                                                                                                                                                                                                                                                                                                                             
25 |                 current.nodes.push(Box::new(node));                                                                                                                                                                                                                                                                        
   |                 ^^^^^^^^^^^^^                                                                                                                                                                                                                                                                                              
   |                 |                                                                                                                                                                                                                                                                                                          
   |                 second mutable borrow occurs here                                                                                                                                                                                                                                                                          
   |                 first borrow later used here

How can I solve this efficiently?

2

this is another classic example of polonius.

the "easy" workaround is to use raw pointer and unsafe:

    pub fn insert(&mut self, data: Vec<u64>) {
-       let mut current = &mut self.root;
+       let mut current = &raw mut self.root;
        let mut index = 0;

        loop {
-           if let Some(child) = current.nodes.get_mut(index) {
+           if let Some(child) = unsafe { (*current).nodes.get_mut(index) } {
                index += 1;
                current = child;
            } else {
                let node = Node {
                    data,
                    nodes: Vec::default(),
                };
-               current.nodes.push(Box::new(node));
+               unsafe { (*current).nodes.push(Box::new(node)) };
                return;
            }
        }
    }

this should work, but it's ugly and unsafe, you should check out the polonius-the-crab crate instead.

but for this particular example, you can use another workaround by converting the loop into tail recursion:

impl Tree {
	pub fn insert(&mut self, data: Vec<u64>) {
		insert_recursive(&mut self.root, 0, data);
	}
}

fn insert_recursive(current: &mut Box<Node>, index: usize, data: Vec<u64>) {
	match current.nodes.get_mut(index) {
		Some(child) => insert_recursive(child, index + 1, data),
		None => {
			let node = Node {
				data,
				nodes: Vec::default(),
			};
			current.nodes.push(Box::new(node))
		}
	}
}
1 Like
3

In this case, you can also rewrite it by simplifying a little:

    pub fn insert(&mut self, data: Vec<u64>) {
        let mut current = &mut self.root;
        let mut index = 0;

        while current.nodes.len() > index {
            current = current.nodes.get_mut(index).unwrap();
            index = index + 1;
        }
        let node = Node {
            data,
            nodes: Vec::default(),
        };
        current.nodes.push(Box::new(node));
    }

@nerditation I thought tail recursion optimization wasn't guaranteed in Rust; has it changed?

1 Like
4

Thanks. This helps a lot. I didn't know about polonius and I'll look it up. I actually wanted an iterative (and safe) solution to avoid stack problems with deep trees. Is rust tail call optimized by now? I read about it a while ago and thought that tail optimization is not guaranteed.

5

@Redglyph Sorry, my example code is highly reduced. I need to work with the child node in the Some(child) branch and the moment I change your code example to

    pub fn insert(&mut self, data: Vec<u64>) {
        let mut current = &mut self.root;
        let mut index = 0;

        while current.nodes.len() > index {
            if let Some(child) = current.nodes.get_mut(index) {
                // work with the child node
                current = child;
                index = index + 1;
            }
        }

        let node = Node {
            data,
            nodes: Vec::default(),
        };

        current.nodes.push(Box::new(node));
    }

I run into the same error as before.

6

Heh. In that case, you don't benefit from the full simplification, but you can still work around the problem by not re-using child to assign the next value of current. You can still use child in the loop, though. The repetition of get_mut(index) isn't the most elegant, but that's what it is: a work-around for the current borrow checker:

    pub fn insert(&mut self, data: Vec<u64>) {
        let mut current = &mut self.root;
        let mut index = 0;

        while let Some(child) = current.nodes.get_mut(index) {
            // work with the child node
            current = current.nodes.get_mut(index).unwrap();
            index = index + 1;
        }

        let node = Node {
            data,
            nodes: Vec::default(),
        };

        current.nodes.push(Box::new(node));
    }

EDIT: From a quick look in the Compiler Explorer, it looks like the loop is well optimized and doesn't show that redundancy in the produced code.

7

Ok. I see your point now. This would work. However, the problem with it is that in my production code the current.nodes.get_mut(index) function call is more complicated and I search for a specific node in current.nodes and the method call repetition is not just an index operation and it looks like the compiler doesn't optimize that in the same way as it does with just get_mut. I actually didn't want to but what do you think about a hybrid between unsafe and your solution:

pub fn insert(&mut self, data: Vec<u64>) {
        let mut current = &mut self.root;
        let mut index = 0;

        while let Some(child) = current.nodes.get_mut(index) {
            // work with the child node
            let raw_child = &raw mut *child;
            current = unsafe { &mut (*raw_child) };
            index = index + 1;
        }

        let node = Node {
            data,
            nodes: Vec::default(),
        };

        current.nodes.push(Box::new(node));
    }

I don't think it's ub but I may be wrong.

8

You should clearly state what you simplify when it's in a critical position for the current problem. :wink:

I don't think what you wrote could lead to undefined behaviour, but that's something you could (and should) verify with Miri in your unit tests. It's sad to have to use unsafe code for that, though it's probably fair enough given the circumstances. If you do, make sure to add a "SAFETY" comment justifying why it's actually safe.

It's hard to tell if there's another way since your code differ significantly, but from what you're saying, there's no efficient work around the loop dependency problem. EDIT: for example, you might be able to come up with alternative tests that are relatively simpler than duplicating a part of the code:

        while current.nodes.len() > index {
            let child = current.nodes.get_mut(index).unwrap();
            // work with the child node
            current = child;
            index = index + 1;
        }

The recursive code is a tempting alternative, as is often the case when you iterate mutably through a nonlinear graph. But I think Rust should have a safeguard that guarantees the optimization, since it's critical. I'm not even sure it's optimized in non-release mode.

For example, Kotlin has a dedicated keyword for that pattern: tailrec. If the compiler can't optimize the tail recursion, it generates an error, which is much safer than having to introduce unit tests that check the stack pointer.

9

It's already in an RFC, fortunately, and even in a more general language construct:

1 Like
10

this snippet as shown here is sound, but it's not easy to reason about in more complex scenario, and it's even harder to make change to the code in the future, that's why polonius-the-crab exists.

the core problem is the conditional control flow, namely, the borrow in the if let clause extends the lifetime to the whole function body, making it impossible to use in the else clause or even after the if-let-else.

here's how you can replace the if-let-else with polonius-the-crab without unsafe:

type MutRefBoxNode = ForLt!(&mut Box<Node>);
loop {
	match polonius_the_crab::polonius::<_, _, MutRefBoxNode>(current, |current| {
		if let Some(child) = current.nodes.get_mut(index) {
			PoloniusResult::Borrowing(child)
		} else {
			PoloniusResult::Owned(())
		}
	}) {
		// this was the `if let` branch of the original code
		PoloniusResult::Borrowing(child) => {
			index += 1;
			current = child;
		}
		// this was the `else` branch
 		// `input_borrow` is the real magic
		PoloniusResult::Owned {
			input_borrow: current, 
			value: _,
		} => {
			current.nodes.push(todo!());
			return;
		}
	}
}

this is somewhat cumbersome, but it's the best what we can get for now.

even if you don't use polonius-the-crab, you can roll your own customized polonius for this particular case, that is, encapsulate this unsafe logic to a minimal scope so it's easy to review and refactor, the signature probably looks like this:

impl Node {
	// this is essentially a tailored `polonius` equivalence
	// without the generics and hkt stuff
	fn child_or_parent(&mut self, index: usize) -> Either<&mut Self, &mut Self> {
		if let Some(child) = self.nodes.get_mut(index) {
			// SAFETY: polonius, blahblahblah ...
			Either::Left(unsafe { &mut *&raw mut *child })
		} else {
			Either::Right(self)
		}
	}
}

I used Either in this example as replacement for PoloniusResult, you can also use a custom enum, e.g. with variants of NodeRef::Child and NodeRef::Parent, and then you can use this to replace the original if-let-else:

pub fn insert(&mut self, data: Vec<u64>) {
	let mut current: &mut Node = &mut self.root;
	let mut index = 0;
	loop {
		match current.child_or_parent(index) {
			Either::Left(child) => {
				index += 1;
				current = child;
			}
			Either::Right(current) => {
				current.nodes.push(todo!());
				return;
			}
		}
	}
}
1 Like
11

I finally managed to read through polonius-the-crab and I also think it is the best and most efficient way to solve my particular problem, at least for the moment. Thank you for your code examples which helped me a lot.

Imho, tail recursion would be the most elegant way if tail optimization were guaranteed. I found a crate https://crates.io/crates/tailcall which guarantees tail optimization until it is superseded by the become keyword.

1 Like