TL;DR; It’s possible to implement a binary tree using Rust’s trait solver.

I didn’t set out to implement a binary tree. It started with a requirement to store one or more implementors of a trait in a struct. Dynamic dispatch seemed like the obvious solution, but it quickly became clear that said trait wasn’t dyn-compatible/object safe¹. Next up was considering using an enum with all the possible trait implementors, but then I’d have to match the concrete type for every operation which very quickly gets too complex.

The use case involved composing all the implementors of the aforementioned trait to act as a single implementor, so that’s where I started—creating a struct that composes two implementors into a single one. Now I can use this structure to compose two or more implementors (and can handle the one case with an empty struct and a no-op implementation of the trait). Problem solved, right?

Well, there was just one issue. I needed to construct the tree before I actually had any instances of the types implementing the trait, and then set each one later on. I would however know upfront which types could be added to the tree. Simple, right? Just wrap each type in an Option, and ensure I can get a mutable reference to it when the time comes to add each instance. If only I could settle for storing each type in a well-known location in the tree, we could say problem solved.

But I couldn’t settle, and was hungry for an adventure, so down the rabbit hole we went!

The tree needed a setter method that would propagate down the tree until it found the right node. So, add a trait and use type state to determine which node we’re setting. Here’s my first attempt:

#[derive(Default)]
struct MaybeNode<T> {
    value: Option<T>,
}

impl<T> SetNode<NodeRight, T> for MaybeNode<T> {
    fn set_node(&mut self, value: T) {
        self.value = Some(value)
    }
}

impl<T> SetNode<NodeLeft, T> for MaybeNode<T> {
    fn set_node(&mut self, value: T) {
        self.value = Some(value)
    }
}

struct EitherNode<A, B> {
    right: A,
    left: B,
}

struct NodeRight;
struct NodeLeft;

pub trait SetNode<Node, T> {
    fn set_node(&mut self, value: T);
}

impl<A1, A2, B> SetNode<NodeRight, A2> for EitherNode<A1, B>
where
    A1: SetNode<NodeRight, A2>,
{
    fn set_node(&mut self, value: A2) {
        SetNode::<NodeRight, A2>::set_node(&mut self.right, value);
    }
}

impl<A, B1, B2> SetNode<NodeLeft, B2> for EitherNode<A, B1>
where
    B1: SetNode<NodeLeft, B2>,
{
    fn set_node(&mut self, value: B2) {
        SetNode::<NodeLeft, B2>::set_node(&mut self.left, value);
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut c = EitherNode {
        right: MaybeNode::<String>::default(),
        left: EitherNode {
            right: MaybeNode::<usize>::default(),
            left: EitherNode {
                right: MaybeNode::<u32>::default(),
                left: EitherNode {
                    right: MaybeNode::<u128>::default(),
                    left: MaybeNode::<u64>::default(),
                },
            },
        },
    };
    c.set_node("Hello string".to_owned());
    // c.set_slot(0usize); // Error
    c.set_node(0u64);
    // c.set_slot(0u32); // Error
    Ok(())
}

Notice that while we can successfully set nodes that follow a consistent path (e.g. right.right, or left.left.left), this simply won’t work if we try to set a path that has any deviation from that. Oops!

The reason this failed is because the trait bound is too restrictive, encoding which node should match, locking us in to either left or right the whole way down. Naively attempting to loosen that bound would result in coherence² issues, which is why we’re using type state in the first place. I found myself thinking that maybe Rust just isn’t capable of representing this sort of thing.

I’m a bit embarrassed to admit that at this point I pasted my code into Claude³ to see if any insight could be gained from that. While it’d be easy to think there’s no harm in that, I’ve become aware of many reasons one might want to consider—which may be the discussion of another blog post. One such reason is that it can erode ones ability to think.

It took me a bit of time to filter through all the confidently given and completely wrong answers it gave. While most of them were obviously wrong, I still got tripped up a bit on one that seemed like it might be true—that I’d need to define typestate for each possible path, similar to if we were working with a tuple of varying sizes⁴. But, several (more targeted) prompts and a few threads later, there was one piece of output that eventually caught my eye which—after doing a quick search— looks like it was pulled from Sgeo’s hlist crate⁵, and as it turned out, pointed me right at a solution (but perhaps a well placed nap could have done the same?—we’ll never know):

// ...
pub struct PathA<Next>(PhantomData<Next>);
pub struct PathB<Next>(PhantomData<Next>);
pub struct PathEnd;

impl<T> SetNode<PathEnd, T> for MaybeNode<T> {
  // ...
}

impl<A, B, Next, T> SetNode<PathA<Next>, T> for EitherNode<A, B>
where
    A: SetNode<Next, T>,
{
  // ...
}

impl<A, B, Next, T> SetNode<PathB<Next>, T> for EitherNode<A, B>
where
    B: SetNode<Next, T>,
{
  // ...
}

Yay, it works! But why? It turns out that Rust’s trait solver is in fact quite powerful and capable of recursion. Setting a breakpoint for each set_node fn makes it clear what’s happening. If we’re calling set_node for a type located at right.right.left, Rust’s trait solver can apparently iterate through each layer, finding a matching impl at each step (all at compile time of course!):

SetNode<PathA<
  SetNode<PathA<
      SetNode<PathB<
          SetNode<PathEnd, T>
      >>
  >>
>>

You can play around with a full implementation here. I even included a macro_rules! to abstract away placement of each type in the tree:

type MyTree = binary_tree_type![u32, u64, usize];

// vs

type MyTree = EitherNode<
  EitherNode<
    Node<u32>, Node<u64>,
  >,
  Node<usize>,
>;

I’m not sure if I’ll end up using this, but it was a fun learning experience!

p.s. An LLM didn’t write this, I just really like mdashes!