Changing a node in a tree in Rust

Question

I\'m trying to write a function that, given a tree structure, returns a copy of that tree but with a node changed at a particular index. Here is what I have so far:

Answer 1

This code compiles:

#[derive(Clone)]
pub enum Node {
    Value(u32),
    Branch(u32, Box, Box),
}

fn main() {
    let root = Node::Branch(1, Box::new(Node::Value(2)), Box::new(Node::Value(3)));
    zero_node(&root, 2);
}

pub fn zero_node (tree: &Node, node_index: u8) -> Node {

    let mut new_tree = tree.clone();

    fn zero_rec (node : &mut Node, node_count : u8, node_index : u8) -> u8 {
        if node_index == node_count {
            match node {
                &mut Node::Value(ref mut val) => { *val = 0; },
                &mut Node::Branch(ref mut val, _, _) => { *val = 0; }
            }
            node_count
        } else {
            match node {
                &mut Node::Value(_) => {1},
                &mut Node::Branch(_, ref mut left, ref mut right) => {
                    let count_left = zero_rec(&mut **left, node_count + 1, node_index);
                    let count_right = zero_rec(&mut **right, node_count + 1 + count_left, node_index);
                    count_left + count_right + 1
                }
            }
        }
    }

    zero_rec(&mut new_tree, 0, node_index);

    new_tree

}

The changes I made were:

• &new_tree → &mut new_tree and &**left → &mut **left etc.: the way to create an &mut T reference is with the &mut operator (i.e. the mut is necessary). This fixes the cannot borrow immutable borrowed content as mutable error by passing a mutable reference rather than an immutable
• changing the node_index == node_count branch to directly mutate the values, rather than trying to overwrite in place. This fixes the cannot move out of borrowed content error by just not doing any moves at all.

The overwriting could actually be achieved with careful use of std::mem::replace, to swap in a new value (e.g. Value(0) since that's cheap to create) to the left and right references. The replace functions returns the value that existed before, i.e. exactly the things inside left and right that you need to create the new branch. This change to the relevant match arm looks a bit like:

&mut Node::Branch(_, ref mut left, ref mut right) => { 
    let l = mem::replace(left, Box::new(Node::Value(0)));
    let r = mem::replace(right, Box::new(Node::Value(0)));
    *node = Node::Branch(0, l , r); 
}

(Having added use std::mem; to the top of the file.)

However it hits a new error:

:25:9: 25:39 error: cannot assign to `*node` because it is borrowed
:25                   *node = Node::Branch(0, l , r); 
                            ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
:22:26: 22:38 note: borrow of `*node` occurs here
:22               &mut Node::Branch(_, ref mut left, ref mut right) => { 
                                             ^~~~~~~~~~~~

The left and right values are pointers deep into the old contents of node, so, as far as the compiler knows (at the moment), overwriting node will invalidate those pointers which would cause any further code using them to be broken (of course, we can see that neither are used more, but the compiler doesn't pay attention to things like that yet). Fortunately there's an easy fix: both match arms are setting node to a new value, so we can use the match to compute the new value and then set node to it after doing the computation:

*node = match node {
    &mut Node::Value(_) => Node::Value(0),
    &mut Node::Branch(_, ref mut left, ref mut right) => { 
        let l = mem::replace(left, Box::new(Node::Value(0)));
        let r = mem::replace(right, Box::new(Node::Value(0)));
        Node::Branch(0, l , r)
    }
};

(NB. the order of operations is a little strange, that's the same as let new_val = match node { ... }; *node = new_val;.)

However, that is more expensive than doing it as I wrote above, since it has to allocate 2 new boxes for the new Branch, while the one that modifies in-place doesn't have to do this.

---

A slightly "nicer" version might be (comments inline):

#[derive(Clone, Show)]
pub enum Node {
    Value(u32),
    Branch(u32, Box, Box),
}

fn main() {
    let root = Node::Branch(1, Box::new(Node::Value(2)), Box::new(Node::Value(3)));
    let root = zero_node(root, 2);

    println!("{:?}", root);
}

// Taking `tree` by value (i.e. not by reference, &) possibly saves on
// `clone`s: the user of `zero_node can transfer ownership (with no
// deep cloning) if they no longer need their tree.
//
// Alternatively, it is more flexible for the caller if it takes 
// `&mut Node` and returns () as it avoids them having to be careful 
// to avoid moving out of borrowed data.
pub fn zero_node (mut tree: Node, node_index: u8) -> Node {

    fn zero_rec (node : &mut Node, node_count : u8, node_index : u8) -> u8 {
        if node_index == node_count {
            // dereferencing once avoids having to repeat &mut a lot
            match *node {
                // it is legal to match on multiple patterns, if they bind the same
                // names with the same types
                Node::Value(ref mut val) | 
                    Node::Branch(ref mut val, _, _) => { *val = 0; },
            }
            node_count
        } else {
            match *node {
                Node::Value(_) => 1,
                Node::Branch(_, ref mut left, ref mut right) => {
                    let count_left = zero_rec(&mut **left, node_count + 1, node_index);
                    let count_right = zero_rec(&mut **right, node_count + 1 + count_left, node_index);
                    count_left + count_right + 1
                }
            }
        }
    }

    zero_rec(&mut tree, 0, node_index);

    tree

}