Is there a way to iterate over a mutable tree to get a random node?

Question

I am trying to update a node of a tree structure. A node which is to be updated is selected randomly. To sample a node in the tree using the Reservoir Sampling algorithm, I

Answer 1

No, it is not safe to write an iterator of the mutable references to the nodes of a tree.

Assume we have this tree structure:

         +-+
    +----+ +----+
    |    +-+    |
    |           |
    |           |
 +--v-+      +--v--+
 | 50 |      | 100 |
 +----+      +-----+

If such an iterator existed, we could call it like this:

let mut all_nodes: Vec<&mut Node> = tree.iter_mut().collect();

Assume that the parent node ends up in index 0, the left node in index 1, and the right node in index 2.

let (head, tail) = all_nodes.split_at_mut(1);

let x = match &mut head[0] {
    Branch(ref mut l, _) => l,
    Leaf(_) => unreachable!(),
};

let y = &mut tail[1];

Now x and y are mutable aliases to each other. We have violated a fundamental Rust requirement in completely safe code. That's why such an iterator is not possible.

---

You could implement an iterator of mutable references to the values in the tree:

impl<'a> Iterator for IterMut<'a> {
    type Item = &'a mut u8;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let node = self.stack.pop()?;

            match node {
                Node::Branch(a, b) => {
                    self.stack.push(b);
                    self.stack.push(a);
                }
                Node::Leaf(l) => return Some(l),
            }
        }
    }
}

This is safe because there's no way to go from one mutable reference to a value to another one. You can then build your random selection on top of that:

{
    let rando = match rand::seq::sample_iter(&mut rand::thread_rng(), tree.iter_mut(), 1) {
        Ok(mut v) => v.pop().unwrap(),
        Err(_) => panic!("Not enough elements"),
    };

    *rando += 1;
}