Tree implementation in Rust

Question

I am a beginner to Rust, started learning because of Advent of Code 2022. I solved Day 7's problem by implementing a tree data structure and traversing it. TLDR, the problem provides a log of terminal inputs and outputs, consisting of either ls or cd <dir> commands, such as

$ cd /
$ ls
dir a
1 b
$ cd a
$ ls a
2 a
3 b

And here is my solution:

use std::cell::RefCell;
use std::default::Default;
use std::fmt::Debug;
use std::ops::Add;
use std::rc::Rc;

const ITER_ERR: &str = "err: iterator is empty";
const PARSE_ERR: &str = "err: can't parse int";
const SPLIT_ERR: &str = "err: can't split string";

struct Solver {
    content: String,
}

trait Solvable {
    fn solve(self) -> (usize, usize);
}

// Blanket implementation
trait NodeValTrait<T>: Add<Output = T> + Default + Copy + Debug {}

impl<T> NodeValTrait<T> for T where T: Add<Output = T> + Default + Copy + Debug {}

#[derive(Debug)]
struct TreeNode<T: NodeValTrait<T>> {
    val: Option<T>,
    name: String,
    children: Vec<Rc<RefCell<TreeNode<T>>>>,
}

impl<T: NodeValTrait<T>> TreeNode<T> {
    fn new(s: impl ToString) -> TreeNode<T> {
        TreeNode {
            val: None,
            name: s.to_string(),
            children: Vec::new(),
        }
    }

    fn val(&mut self, val: T) {
        self.val = Some(val);
    }

    fn push(&mut self, node: TreeNode<T>) {
        self.children.push(Rc::new(RefCell::new(node)));
    }

    fn insert_and_find(&mut self, name: impl ToString) -> Rc<RefCell<TreeNode<T>>> {
        let name = name.to_string();
        for child in &self.children {
            if child.borrow().name == name {
                return Rc::clone(child);
            }
        }
        let child = TreeNode::new(name);
        self.push(child);
        Rc::clone(&self.children[&self.children.len() - 1])
    }

    fn sum(&self) -> T {
        let mut res: T = Default::default();
        if let Some(val) = &self.val {
            res = res + *val;
        }
        for child in &self.children {
            res = res + child.borrow().sum();
        }
        res
    }

    fn dfs<F, T2>(&self, f: F) -> Vec<T2>
    where
        F: Fn(&TreeNode<T>) -> T2,
    {
        let mut res = Vec::new();
        res.push(f(self));
        for child in &self.children {
            // workaround for infinite recursion
            let f = &f as &dyn Fn(&TreeNode<T>) -> T2;
            res.extend(child.borrow_mut().dfs(f));
        }
        res
    }
}

impl Solvable for Solver {
    fn solve(self) -> (usize, usize) {
        let root: Rc<RefCell<TreeNode<usize>>> = Rc::new(RefCell::new(TreeNode::new("")));
        let mut cur_path = Vec::new();
        cur_path.push(Rc::clone(&root));

        // each $ represents new command input / output group
        for cmds in self.content.split("$ ") {
            let cur = &cur_path[cur_path.len() - 1];
            let lines: Vec<&str> = cmds.trim().lines().collect();
            if lines.is_empty() {
                continue;
            }

            let cmd: Vec<&str> = lines[0].split(" ").collect();
            if cmd[0] == "ls" {
                // handles `ls` commands
                for &line in &lines[1..] {
                    let (prefix, dir) = line.split_once(" ").expect(SPLIT_ERR);
                    let child = cur.borrow_mut().insert_and_find(dir);
                    match prefix {
                        "dir" => {
                            cur.borrow_mut().insert_and_find(dir);
                        }
                        _ => {
                            let prefix = prefix.parse().expect(PARSE_ERR);
                            child.borrow_mut().val(prefix);
                        }
                    }
                }
            } else {
                // handles `cd` commands
                match cmd[1] {
                    ".." => {
                        cur_path.pop();
                    }
                    dir => {
                        let child = cur.borrow_mut().insert_and_find(dir);
                        cur_path.push(child);
                    }
                }
            }
        }

        let func = |node: &TreeNode<usize>| (node.val.is_some(), node.sum());
        let binding = root.borrow_mut().dfs(func);
        let vals = binding
            .iter()
            .filter(|(is_file, _)| !is_file)
            .map(|(_, sum)| sum);

        // part 1
        let part1 = vals.clone().filter(|&sum| *sum <= 100000).sum();

        // part 2
        let freeup = root.borrow_mut().sum() - 40000000;
        let part2 = vals
            .clone()
            .filter(|&sum| *sum >= freeup)
            .min()
            .expect(ITER_ERR);

        (part1, *part2)
    }
}

fn main() {
    let solver = Solver {
        content: include_str!("../input").to_string(),
    };
    let (part1, part2) = solver.solve();
    println!("Part 1: {part1:?}");
    println!("Part 2: {part2:?}");
}

If there are any suggestions on how to simplify this code, or how to make it more Rust-like, then it would be great.

• In particular, the implementation for dfs just seems... weird, especially the "hack" with &f as &dyn ... to avoid infinite recursion at compile time.

• Also, is it possible to modify dfs such that it returns say Vec<(&TreeNode<T>, T2)> or similar? I tried doing it naively by replace res.push(f(self)) with res.push((self, f(self))), but it fails horribly with tons of lifetime problems.

• Another question - is it encouraged to use match and iter.method().method().method() chains over other options like if let? It seems that especially when using match, the indentation level goes quite deep.

• Finally, are there other options to Rc<RefCell<...>>? It is working great but keeping track of the types are annoying.

Answer 1

In particular, the implementation for dfs just seems... weird, especially the "hack" with &f as &dyn ... to avoid infinite recursion at compile time.

Let's take a look:

fn dfs<F, T2>(&self, f: F) -> Vec<T2>
where
    F: Fn(&TreeNode<T>) -> T2,
{
    let mut res = Vec::new();
    res.push(f(self));
    for child in &self.children {
        // workaround for infinite recursion
        let f = &f as &dyn Fn(&TreeNode<T>) -> T2;
        res.extend(child.borrow_mut().dfs(f));
    }
    res
}

The fundamental problem here is that dfs takes ownership of f, but you don't want to pass ownership of f when calling it recursively. The easiest solution is to never take ownership. Specifically, take a borrow of f:

fn dfs<F, T2>(&self, f: &F) -> Vec<T2>
where
    F: Fn(&TreeNode<T>) -> T2,
{
    let mut res = Vec::new();
    res.push(f(self));
    for child in &self.children {
        res.extend(child.borrow_mut().dfs(f));
    }
    res
}

This will require updating calling sites to pass a borrow.

However, there is another problem with this implementation. It will create and destroy a lot of Vecs. It'd be more efficient to put everything directly into a single target vec. I'd split it into two functions:

fn dfs_impl<F, T2>(&self, res: &mut Vec<T2>, f: &F)
where
    F: Fn(&TreeNode<T>) -> T2,
{
    res.push(f(self));
    for child in &self.children {
        child.borrow_mut().dfs_impl(res, f);
    }
}

fn dfs<F, T2>(&self, f: F) -> Vec<T2>
where
    F: Fn(&TreeNode<T>) -> T2,
{
    let mut res = Vec::new();
    self.dfs_impl(&mut res, &f);
    res
}

Another question - is it encouraged to use match and iter.method().method().method() chains over other options like if let? It seems that especially when using match, the indentation level goes quite deep.

In cases where there is a binary decisions (if this else that) then I recommend using if let over a match. Matches makes sense if you have several different paths to choose from. I also usually prefer if let over methods but that's more situations specific.

Also, is it possible to modify dfs such that it returns say Vec<(&TreeNode, T2)> or similar? I tried doing it naively by replace res.push(f(self)) with res.push((self, f(self))), but it fails horribly with tons of lifetime problems.

As it stands, you shouldn't. Your code works with Rc<RefCell< and so you should return that and not simple borrows. You can do this by defining static methods that operate on Rc<RefCell<.

fn dfs_impl<F, T2>(me: &Rc<RefCell<Self>>, res: &mut Vec<(Rc<RefCell<Self>>, T2)>, f: &F)
where
    F: Fn(&TreeNode<T>) -> T2,
{
    res.push((Rc::clone(me), f(&me.borrow())));
    for child in &me.borrow().children {
        TreeNode::dfs_impl(&child, res, f);
    }
}

fn dfs<F, T2>(me: &Rc<RefCell<Self>>, f: F) -> Vec<(Rc<RefCell<Self>>, T2)>
where
    F: Fn(&TreeNode<T>) -> T2,
{
    let mut res = Vec::new();
    TreeNode::dfs_impl(me, &mut res, &f);
    res
}

But this leads us to your last question.

Finally, are there other options to Rc<RefCell<...>>? It is working great but keeping track of the types are annoying.

The alternative is to use the types directly. Instead of Vec<Rc<RefCell<TreeNode>> just use Vec<TreeNode>. If you just do that with your current code you'll end up in a morass of lifetime errors. The particular problem that you'll find is cur_path in Solver::solve. Right now it is a vector of Rcs which is fine, but a Vec of mutable borrows to different parts of the same tree structure and that won't work very well since Rust's borrow checker can't ensure that you use that safely.

One approach would be reimplement the solution recursively. In particular, anytime you need to move down a layer in the tree, you should call a method on that tree node to take care of it. It should return when it leaves tree (i.e. sees a cd ..) so that the directory above can continue processing.

Another approach would be to use indexes instead of references. The idea here is that your tree structure has methods to manipulate the tree and takes indexes to determine which node is being manipulated instead of operating on references or Rcs. This avoids issues with the borrow checker and avoids the annoyance of dealing with Rc<RefCell. The drawback is that Rust borrow checker can't ensure that you use the indexes correctly.

Answer 2

To simplify the code, you can remove the unused ITER_ERR, PARSE_ERR, and SPLIT_ERR constants and the Debug trait bounds for the NodeValTrait trait. You can also remove the lines variable and directly iterate over the cmds split by "$ ". You can also use match statements instead of if statements for the cmd[0] comparisons.

Answer 3

Here is a temporary solution that I settled with (in particular for question 2 and 4), as well as some self-reviews after learning Rust for 3 more days!

Most of my reviews are already mentioned by Winston Ewert in his answer. However, some one other coding style review I have is:

Prefer if-else if-else for few branches with long codes

For example, the code

match prefix {
    "dir" => {
        cur.borrow_mut().insert_and_find(dir);
    }
    _ => {
        let prefix = prefix.parse().expect(PARSE_ERR);
        child.borrow_mut().val(prefix);
    }
}

Only has two branches, and should be replaced by

if prefix == "dir" {
    // ...
} else {
    // ...
}

This way, the intention of the code is clearer, and overall reduces indentation.

---

And now to answer my own questions:

Question 1: Answered by Winston Ewert - essentially, pass in &f instead so that dfs does not take ownership of f.

Question 2: One option is to use an arena, such as the typed-arena crate or just a Vec<TreeNode>. They function as a local memory allocators. Whenever a TreeNode allocation is required, the arena can do that and return a reference, either a direct &TreeNode reference or an uint index to an internal buffer. This also means the lifetime of nodes are tied to the arena instead of to their parents. Further discussion under Winston Ewert's answer.

Question 3: Answered above.

Question 4: Answered by Winston Ewert - essentially, use TreeNode<T> directly and rewrite the algorithms around the data structure, or use the arena method mentioned above.

You can see my implementation at this commit.