Simple tree addition implementation

Question

I am trying to learn Rust and I feel like a dummy every step of the way.

I went through the tutorial on making lists in Rust. I am trying to apply the knowledge to implement a simple BST.

The code looks pretty horrible though. Any suggestions on how it can be improved?

pub struct Tree<T: Ord> {
    root: Link<T>,
}

type Link<T> = Option<Box<Node<T>>>;

pub struct Node<T:Ord> {
    elem: T,
    left: Link<T>,
    right: Link<T>
}

impl<T: Ord> Tree<T> {
    pub fn new() -> Self {
        Tree {root: None}
    }

    pub fn add(&mut self, elem: T) {
        let mut current;
        match self.root {
            None => {
                self.root = Some(Box::new(Node{elem: elem, left: None, right: None}));
                return;
            },
            Some(_) => {
                current = self.root.as_mut();
            }
        }
        loop {
            if elem < current.as_ref().unwrap().elem {
                if current.as_ref().unwrap().left.is_none() {
                    current.unwrap().left = Some(Box::new(Node{elem: elem, left: None, right: None}));
                    break;
                } else {
                    current = current.unwrap().left.as_mut();
                }
            } else {
                if current.as_ref().unwrap().right.is_none() {
                    current.unwrap().right = Some(Box::new(Node{elem: elem, left: None, right: None}));
                    break;
                } else {
                    current = current.unwrap().right.as_mut();
                }
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super:: Tree;

    #[test]
    fn basics() {
        let mut tree = Tree::new();
        tree.add(5);
        tree.add(3);
        tree.add(2);
        tree.add(4);

        assert_eq!(tree.root.as_ref().unwrap().elem, 5);
        assert_eq!(tree.root.as_ref().unwrap().left.as_ref().unwrap().elem, 3);
        assert_eq!(tree.root.as_ref().unwrap().left.as_ref().unwrap().left.as_ref().unwrap().elem, 2);
        assert_eq!(tree.root.as_ref().unwrap().left.as_ref().unwrap().right.as_ref().unwrap().elem, 4);
    }
}

Answer 1

There are a few things you can do to reduce repetition in your code. The first would be encapsulating your new Node<T> into a method.

impl<T> Node<T> {
    fn new(elem: T) -> Self {
        Self {
            elem,
            left: None,
            right: None,
        }
    }
}

Then all the constructions of Node in add can be replaced with Node::new(elem).

---

Another thing that can reduce verbosity is having current be an actual (boxed) node rather than an option. Right now you maintain the invariant that current is the Some variant of an Option<&mut Box<Node<T>>> (Link<T>). By instead keeping the &mut Box<Node<T>>, it'll allow you to reduce the number of unwraps. Conceptually, this means that we're expressing the invariant using the type system.

For example, to initialize current,

let mut current;
match self.root {
    None => {
        self.root = Some(Box::new(Node::new(elem)));
        return;
    }
    Some(ref mut node) => {
        current = node;
    }
}

(matching on Some(ref mut node) makes node a mutable reference - there are some other ways to get the same effect, but this is the clearest).

Having this allows you to change the later blocks to

if elem < current.elem {
    if current.left.is_none() {
        current.left = Some(Box::new(Node::new(elem)));
        break;
    } else {
        current = current.left.as_mut().unwrap();
    }
}

---

My last major suggestion is to use matching to get rid of all the unwrapping. Blocks like

if current.left.is_none() {
    current.left = Some(Box::new(Node::new(elem)));
    break;
} else {
    current = current.left.as_mut().unwrap();
}

are just the poor man's match statement. They can be replaced by

match current.left {
    None => {
        current.left = Some(Box::new(Node::new(elem)));
        break;
    }
    Some(ref mut node) => {
        current = node;
    }
}

(again, ref mut node means we're capturing a mutable reference).

This could more succinctly be expressed using if let.

if let Some(ref mut node) = current.left {
    current = node;
} else {
    current.left = Some(Box::new(Node::new(elem)));
    break;
}

---

Now for some extra things. First, as I mentioned in the comments, the code as written handles duplicates by placing them to the right. This isn't necessarily what you want, but it's easy to change. The trait Ord provides a cmp method that has three possible outputs: Less, Greater or Equal. By doing a match on the output of cmp, we can handle all three.

match elem.cmp(&current.elem) {
    std::cmp::Ordering::Less => {
        //...
    }
    std::cmp::Ordering::Equal => break,
    std::cmp::Ordering::Greater => {
        //...
    }
}

If you don't want to type out the whole names, you could do use std::cmp::Ordering::*; inside the fn add block (or wherever else, but it's best to keep the scope small).

---

This suggestion would be a pretty major rework of your code, but it leads to a much more elegant solution. As I said before, you maintain the invariant that current is a Some variant. I suggested that you could keep a &mut Box<Node<T>> instead so that the invariant is expressed in the type of current. What you can do instead is not try to keep that invariant at all, but instead check whether current is Some or None at the start of the loop. Then you wouldn't need the initial test on self.root.

fn add(&mut self, elem: T) {
    let mut current = &mut self.root;
    loop {
        match current {
            Some(node) => {
                // check how elem compares to node.elem
                // assign node.left/right to current as appropriate (or break)
            }
            None => {
                // set the data behind `current` to a new node.
                // i.e. `*current = ...`
            }
        }
    }

---

Finally, just a fairly minor thing. Don't forget to run cargo fmt and cargo clippy. cargo fmt will make your code more readable to other Rust programmers and keep things in a consistent style. cargo clippy will give you a few helpful tips for making your code more idiomatic and sometimes help avoid logic errors.