Retrieve a mutable reference to a tree value

Question

Disclosure: This is a homework problem for my Rust programming class. I am also really trying to understand/learn this material.

So I have a tree similar to the following (contain change):

struct Tree<V> {
  val: Option<V>,
  children: Vec<Tree<V>>
}

My task is to search the tree from the root and retrieve a mutable reference to the val. Here is the function signature:

fn get_mut(&mut self, key: &str) -> Option<&mut V>

The function itself it pretty simple. I take the self parameter and set it to a current_node variable and search until I find the correct node. I already have this part done and tested (I can see it searching through the tree with debug statements). Since I was passed a mutable version of myself, this is current_node:

let mut current_node: &mut Tree<V> = self;

In order to keep the mutable reference to current_node, when searching through the children of each tree/subtree, I use iter_mut() on the Vec:

for child in current_node.children.iter_mut() {
  ...
}

Once I have found the correct node, I pull out the value from it and return:

let val_opt: Option<&mut V> = current_node.val.as_mut();
return match val_opt {
  Some(val) => {
    return Some(val)
  },
  None => None
}

However, I have an error on the children loop:

cannot borrow `current_node.children` as mutable more than once at a time. mutable borrow starts here in previous iteration of loop.

let's call the lifetime of this reference `'1` (start of the function)

mutable borrow starts here in previous iteration of loop (child loop)

returning this value requires that `current_node.children` is borrowed for `'1` (on the line where I return Some(val) at the very end of the function)

I'm trying to understand why this occurs and how I can overcome it. From basic searching, it seems like it has to do with lifetimes but I'm a little lost on what I have to do to remedy it. Also, I cannot change the Tree<V> struct or the function signature.

Minimal implementations:

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Tree<V> {
    val: Option<V>,
    children: Vec<(char, Tree<V>)>,
}

fn str_split_first(x: &str) -> Option<(char, &str)> {
    let mut chars = x.chars();
    let first = chars.next()?;
    Some((first, chars.as_str()))
}

impl<V> Tree<V> {
    pub fn get_mut(&mut self, key: &str) -> Option<&mut V> {
        let mut current_node: &mut Tree<V> = self;
        let mut key_index: usize = 0;

        loop {
            match str_split_first(&key[key_index..]) {
                Some((c, _)) => {
                    let mut found: bool = false;

                    'child_loop: for (character, child) in current_node.children.iter_mut() {
                        if c == *character {
                            current_node = child;
                            found = true;
                            break 'child_loop;
                        }
                    }

                    // If nothing was found, this key does not exist.
                    if found == false {
                        return None;
                    }

                    key_index += 1;
                }
                _ => break,
            }
        }

        let val_opt: Option<&mut V> = current_node.val.as_mut();
        return match val_opt {
            Some(val) => return Some(val),
            None => None,
        };
    }
}

Playground link

Answer 1

The problem here is that you are modifying current_node while iterating over its children. The Rust compiler is trying to protect you from a iterator invalidation, but in this case it is actually correct to do so, because the iteration stops immediately after changing value of current_node.

Unfortunately, the compiler is not smart enough to see it yet. A simple workaround is to store a new value in some temporary variable and update current_node outside of the loop. Like so:

let mut found = None;

for (character, child) in current_node.children.iter_mut() {
    if c == *character {
        found = Some(child);
        break;
    }
}

current_node = match found {
    Some(child) => child,
    None => return None,
};

Link to the playground