Complicating Rust to Make a Point
Last week, Monday Morning Haskell published a post about binary trees, which includes a code example in Rust. The article concludes that the problem, inverting the binary tree, is made difficult by Rust's ownership model, and that this applies to recursive data structures in general. So, as the resident Rust fanboy, I take it upon myself to defend and restore Rust's tarnished honour from such vile accusations! Along the way, we will have a look at different ways to represent the problem in Rust and the different trade-offs each entails. Let's dive right in!
Our baseline
The article uses a LeetCode Problem, which represents child nodes as Rc<RefCell<TreeNode>.
While this is a valid presentation, it is often needlessly complicated and makes solving the problem harder than it needs to be.
To me, that is the core of the problem, which the post falsely attributes to a limitation in Rust.
But let's have a look at this baseline code, as assembled from the post:
use std::cell::RefCell;
use std::rc::Rc;
#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
pub val: i32,
pub left: Option<Rc<RefCell<TreeNode>>>,
pub right: Option<Rc<RefCell<TreeNode>>>,
}
impl TreeNode {
#[inline]
pub fn new(val: i32) -> Self {
TreeNode {
val,
left: None,
right: None,
}
}
}
pub fn invert_tree(root: Option<Rc<RefCell<TreeNode>>>) -> Option<Rc<RefCell<TreeNode>>> {
if let Some(node) = root.clone() {
let mut node_ref = node.borrow_mut();
// Recursively invert left and right subtrees
let left = invert_tree(node_ref.left.clone());
let right = invert_tree(node_ref.right.clone());
// Swap them
node_ref.left = right;
node_ref.right = left;
}
return root;
}
Improving the RefCell-based implementation
To start, let's apply some minor fixes to the existing implementation. The result still isn't pretty, but it works.
use std::cell::RefCell;
use std::rc::Rc;
type Link = Option<Rc<RefCell<TreeNode>>>;
#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
pub val: i32,
pub left: Link,
pub right: Link,
}
impl TreeNode {
#[inline]
pub fn inner(val: i32, left: Link, right: Link) -> Self {
TreeNode { val, left, right }
}
#[inline]
pub fn leaf(val: i32) -> Self {
Self::inner(val, None, None)
}
}
pub fn invert_tree(root: Link) -> Option<Rc<RefCell<TreeNode>>> {
root.map(|root| {
{
let mut root_ref = root.borrow_mut();
// Swap them
(root_ref.left, root_ref.right) = (
invert_tree(root_ref.right.take()),
invert_tree(root_ref.left.take()),
);
}
root
})
}
At this point, I also wrote a conversion function and a property test, to ensure that both implementations produced equivalent trees.
While implementing the test, I encountered a minor hurdle. I originally implemented my test something like this:
let baseline_inverted = crate::baseline::invert_tree(tree.clone());
let improved_inverted = crate::improved::invert_tree(crate::improved::convert(tree.clone()));
However, with this order, the test kept failing. It took me a moment to notice the issue: Since our baseline tree uses interior mutability, inverting messes up the initial state for the other tests. I was able to fix this by changing the call order of different inversion functions:
let improved_inverted = crate::improved::invert_tree(crate::improved::convert(tree.clone()));
// Baseline needs to be inverted last, otherwise we mess up `tree`
let baseline_inverted = crate::baseline::invert_tree(tree);
Haskell-like implementation
Then, I added an implementation inspired by the Haskell implementation from the original post:
#[derive(Debug, Eq, PartialEq)]
pub enum TreeNode {
Nil,
Node(i32, Box<Self>, Box<Self>),
}
pub fn invert_tree(node: TreeNode) -> TreeNode {
match node {
TreeNode::Nil => TreeNode::Nil,
TreeNode::Node(val, left, right) => TreeNode::Node(
val,
Box::new(invert_tree(*right)),
Box::new(invert_tree(*left)),
),
}
}
We still need some indirection (Box<Self>) to avoid an unbounded size of the type.
However, using Box instead of Rc allows us to unwrap and rewrap the type with relative ease.
At this point, it should be clear that Rust can implement as easily as Haskell. Only the original formulation of the problem on LeetCode makes the resulting implementation awkward.
Idiomatic rust
The Haskell implementation is nice, but let's go with something a bit more idiomatic:
#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
pub val: i32,
pub left: Option<Box<Self>>,
pub right: Option<Box<Self>>,
}
pub fn invert_tree(root: Option<Box<TreeNode>>) -> Option<Box<TreeNode>> {
root.map(|root| {
let TreeNode { val, left, right } = *root;
let (left, right) = (invert_tree(right), invert_tree(left));
Box::new(TreeNode { val, left, right })
})
}
In contrast to the other implementation, we lean fully into the functional transformation from one tree to the other.
Instead of using an Rc, we use a Box to ensure unique ownership.
Mutable recursive implementation
Following from the previous implementation, we can create an implementation that uses mutability instead of the more functional transformation:
use core::mem;
#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
pub val: i32,
pub left: Option<Box<Self>>,
pub right: Option<Box<Self>>,
}
pub fn invert_tree(root: &mut Option<Box<TreeNode>>) {
let Some(root) = root.as_deref_mut() else {
return;
};
mem::swap(&mut root.left, &mut root.right);
invert_tree(&mut root.left);
invert_tree(&mut root.right);
}
We still use recursion, but instead of taking ownership, we only pass an exclusive reference which we use to modify the node.
Mutable iterative implementation
We can also adapt our mutable recursive implementation to use an iterative approach:
use core::mem;
#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
pub val: i32,
pub left: Option<Box<Self>>,
pub right: Option<Box<Self>>,
}
pub fn invert_tree(root: &mut Option<Box<TreeNode>>) {
let mut queue = VecDeque::from([root]);
while let Some(node) = queue.pop_front() {
let Some(node) = node.as_deref_mut() else {
continue;
};
mem::swap(&mut node.left, &mut node.right);
queue.push_back(&mut node.left);
queue.push_back(&mut node.right);
}
}
Such an implementation is made possible by using a mutable implementation. With the transformation approach, an iterative implementation would be impossible.
Better implementation using interior mutability
The core problem with the initial implementation is that it mixes interior mutability with functional transformation.
This is also what my initial testing attempt ran afoul of.
If we fully commit to functional transformation, we can use Rc::make_mut to ensure that we do not modify the tree underneath other owners.
Rc::make_mut returns a unique reference to the inner value, either directly when there is only one owner, or by first cloning the inner value when necessary.
use std::cell::RefCell;
use std::rc::Rc;
pub type Link = Option<Rc<RefCell<TreeNode>>>;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TreeNode {
pub val: i32,
pub left: Link,
pub right: Link,
}
pub fn invert_tree(root: Link) -> Option<Rc<RefCell<TreeNode>>> {
root.map(|mut root| {
let root_ref = Rc::make_mut(&mut root).get_mut();
// Swap them
(root_ref.left, root_ref.right) = (
invert_tree(root_ref.right.take()),
invert_tree(root_ref.left.take()),
);
root
})
}
This implementation is only marginally more complex than the Haskell implementation.
Conclusion
I hope that after this journey, you have seen a few different ways to implement a binary tree in Rust. There is still a few more (raw pointers, anyone?), but the selection already shows what can be achieved in entirely safe Rust.
It should also have shown that the original problem formulation for Rust was quite awkward and can be cleaned up significantly if we allow for minor changes in the function signature.