DEV Community

No More Tears, No More Knots: Arena-Allocated Trees in Rust

Ben Lovy on December 19, 2019

Enter The Arena When programming in Rust, it's not always straightforward to directly translate idioms you know. One such category is t...

Read full post

Ben Lovy • Dec 7 '22 • Edited

You're ahead of me - I've been slacking this year! You're right - T is a stand-in for whatever type you want as long as it implements the PartialEq trait. There is a usage example in the playground link at the bottom of the post, here's a snippet showing some usage where T is String:

let mut tree: ArenaTree<String> = ArenaTree::default();
let hello = tree.node("Hello".into());
let world = tree.node("World".into());
tree.arena[hello].children.push(world);
tree.arena[world].parent = Some(hello);

The values stored in the hello and world variables are just indexes into the arena vector. You can replace String with any type that implements PartialEq:

#[derive(PartialEq)]
struct CoolData {
    data: String
}
// ...
let mut tree: ArenaTree<CoolData> = ArenaTree::default();
let hello_struct = CoolData { data: "Hello".to_string() };
let hello = tree.node(hello_struct);

The rest will work the same. Does this help?

Ben Lovy • Dec 7 '22

Rust is definitely a difficult choice especially once you start hitting things like trees. Educational for sure, but you'll hit a lot more friction than other languages for some of the data structures you'll need. Good luck!

David Wickes • Dec 19 '19

This is interesting! The two thoughts I had were:

this is not how I solved this problem for AoC.
this really highlights that array indexes are just pointers.

I also had a helluva time trying to make a tree in Rust. This is ace!

Ben Lovy • Dec 19 '19

this really highlights that array indexes are just pointers.

Boring as it's been, being forced to drill data structures in C++ has admittedly been pretty helpful. Nothing is complicated.

Willing to share how you solved it? I saw this solution when I first read the problem and didn't spend much time thinking of other strategies.

David Wickes • Dec 19 '19

Willingly.

I used a reverse lookup hash table. Each key was the name of a node, and its value was the name of the parent node. This was taking advantage of the way the nodes only ever had one parent.

Distances between nodes was just the sum of the distance between them and their first common ancestor. First common ancestor was calculated by walking one of the nodes all the way up to root, collecting the node names, then walk the other one up until you find a name that's in the other list.

I'd like to say I was being clever, but I definitely stumbled across this one.

Ben Lovy • Dec 19 '19 • Edited

walk the other one up until you find a name that's in the other list

Oh, but that IS clever. At a glance though, you're stepping through the first list for each check as you walk up the second - is there a more efficient way to perform that check, or am I overestimating the computational hit? Not that I'd expect it to practically matter...

David Wickes • Dec 19 '19

is there a more efficient way to perform that check, or am I overestimating the computational hit?

Yeah, it's definitely not optimal. Repeatedly 'finding' in a list is annoying. If I was trying to optimize...

You could store the 'list' as yet another hash table, keying off the node name with values set to the distance, as there's no need to store that information as an ordered sequence if you save the distance. That way you could query it to find out if a node is present, step forward on the other path when its not, or add the distance returned from the hash to the distance travelled from the other path. That's a bit neater.

I'm beginning to wonder whether I just try to solve everything with hashes...

David Wickes • Dec 19 '19

is there a more efficient way to perform that check, or am I overestimating the computational hit?

Yeah, it's definitely not optimal. Repeatedly 'finding' in a list is annoying. If I was trying to optimize...

I'm beginning to wonder whether I just try to solve everything with hashes...

Ben Lovy • Dec 19 '19

It's hashes all the way down...

Andrey Orst • Dec 20 '19

Nice code, I had thought about doing trees this way myself, but I feel that this tree is not as flexible as pointer based tree. For instance, if I would like to add one tree as a children to another tree I would have to copy whole tree node by node updating each index. This is kinda slow process. In comparison, in pointer based tree we just say that this root node is now a child of that tree.

Another example is deleting subtrees. If we were to remove some nodes with deallocating the memory, then we're removing data from vector, which in turn will move all values in a vector for each node we delete. This is like super unoptimized for a tree.

I think this solution is great for grow-only trees, and not for trees that change frequently.

Ben Lovy • Dec 20 '19

Both great points. As always, you've gotta fit the data structure to the task, and this solves some problems but not all.

For point two, could you just "leak" those nodes, and then de-allocate in bulk as part of some other operation if needed? @mkm mentioned the node leaking issue, but it seems like an easier solve than the problem you describe, where each node is actually moving for each de-allocation.

Andrey Orst • Dec 20 '19

As always, you've gotta fit the data structure to the task, and this solves some problems but not all.

Exactly. There are always some tradeoffs.

For point two, could you just "leak" those nodes, and then de-allocate in bulk as part of some other operation if needed?

I think leak is an option, however only if trees are consistently small, and their lifetime determines that those will be deallocated relatively soon. It is a leak after all, and leaks are more painful over time as they tend to grow.

Though in my case (not related to AOC) it wasn't an option because I parse code to a huge tree structure and then do big amount of node manipulation on reduction, which also has tree attachment as a some part of reduction in some cases, which in turn will be removed in the end. So after all reductions the resulting arena would be bloated with so many trees that the memory usage would actually be quite big. The tree would die in the end, of course, but this is not an option for really huge structures that I can theoretically parse.

Ben Lovy • Dec 20 '19

Interesting. I think I'm going to keep playing with this code to see if it can be made more practical. It was so much more ergonomic than any alternative for solving the AoC problem, and took such a minimal amount of planning and code, that I feel all of these concerns could be at least mitigated if not totally resolved with a little more logic built in.

Andrey Orst • Dec 20 '19

I slightly doubt, but kudos if you will come up with something. Leaking nodes seems the only performant way for removal, because if we remove items from vector, we'll have to update IDs of all nodes that were allocated after removed ID in the arena, which can be huge amount of work if we decided to remove leftmost root's child. In this sense I think HashMap is more affordable, though we're going to have to sacrifice some performance to hashing and looking up.

Ben Lovy • Dec 20 '19

HashMap

Heh, that was exactly what I had in mind. You're probably right.

update IDs of all nodes that were allocated after removed ID in the arena

Could this be mitigated by not indexing arena directly, and instead searching for node.idx? All edit operations would just need to ensure they're managing the stored indices properly, but to me it looks like you wouldn't ever need to change an index of a node after it's allocated. It won't match the actual index in the arena, but I'm not sure that matters.

It's not something I know anything about, but I think there are crates designed around sparse data structures as well. Assuming trade-offs but for trees expected to have long lives over which they grow and shrink a lot, could be useful.

bretthancox • Dec 19 '19

I'm in the middle of day 6 of AoC right now and I'm hitting issues with the shortest route between me and Santa. This article is a great reminder that rust's type system gives me more flexibility than I'm used to. I've been trying to force HashMap and vectors to work, but your solution is so elegant. Thanks for the write up. A fresh perspective is what I needed.

Ben Lovy • Dec 19 '19

Rust and Clojure are a great pair for these challenges. They both afford a high degree of expressive freedom but from completely different directions.

Andrey Frolov • May 7 '22

Hi there, thanks for the post. What about if I want to have a different trees/subtrees builded from the common "pool" of nodes? Like a two separate heads of the list pointed to the same node.

Does this data structure fits well for such case? (of course if we wrap all nodes in Rc)

Ben Lovy • May 9 '22

Hi - because all nodes only refer to each other via a Copy index, instead of directly pointing to each other, you should be able to set up such structures without even needing to wrap in Rc. The nodes are only accessed via the arena as needed, the only reference to them is just a number that can be passed around without issue.

Andrey Frolov • May 9 '22

Thank you for the clarification

Ben Lovy • May 9 '22

The caveat would be simultaneous access - you would need to employ a lock of some sort if you expect multiple threads or tasks to mutably borrow the data in a given node.

Marko Mikulicic • Dec 19 '19

This is a nice way to stay memory safe with respect to the general memory, but with respect to your high-level "memory cells" (the node entries in the arena) you're on your own: the compiler won't protect you from accidentally "leaking" nodes in a long living tree (where a leak is just a node whose index is not pointed to by other nodes) nor preventing two nodes to point to the same node (breaking the tree invariant etc)

Ben Lovy • Dec 19 '19 • Edited

Great points, absolutely. There is some work to be done here for practical use beyond coding challenges.

I'm very much a novice - it sounds to me like point 2 could be resolved by just writing better methods that verify this property on insertion. Would that suffice?

Leaked nodes are a little trickier - you could conceivably write up a little "stop-the-world" style cleanup that sweeps through for unlinked nodes, but it's not immediately clear to me how to build that in to this structure without requiring the consumer to be aware.

Ryan Palo • Dec 19 '19

This is neat! I ended up using a HashMap of nodes and their children’s names, but that involves a lot of duplication. I was reading through the rust docs and saw that they recommended refcell> for stuff like this, but it was pretty intimidating.

Ben Lovy • Dec 19 '19

Totally - I've gone the Rc<RefCell<T>> route for other projects, and it does get you there but it always feels hacky and brittle. In C++, you really just use the Rc-equivalent shared_ptr when you actually want multiple ownership over some data, not to skirt around language semantics when its convenient for you - this strategy is kinda antithetical to what draws me to Rust in the first place! There are definitely legitimate uses of runtime interior mutability but I don't feel this is one of them.

u_dev • Jul 30 '22

This is awesome
One thing i would like to understand is, the ideal look up time for a tree is O(logn) but that doesnt seem to be followed here since you are iterating all the nodes in the vector.

for node in &self.arena {
            if node.val == val {
                return node.idx;
            }
        }

this seems to take O(n) time. Is this a known tradeoff?

jchimene • Jul 7 '20

Thanks!
Still an issue in
impl ArenaTree
where
T: PartialEq
{
fn node(&mut self, val: T) -> usize {
//first see if it exists
for node in &self.arena {
if node.val == val {
return node.idx;
}
}
// Otherwise, add new node
let idx = self.arena.len();
self.arena.push(Node::new(idx, name));
idx
}
}
I think you want to s/name/val/

Gerard Klijs • Dec 19 '19

Nice solution to a common rust problem. Thanks for sharing.

Kamil Dydo • Feb 7 '21 • Edited

When creating a node you're setting idx to be the length of the current Vec of children. If you then remove the one before last element and create another one at the end 2 last elements will have the same index.

Mike Lezhnin • Dec 19 '19

I think it is better to keep number of nodes you have in a separate variable, so that you can easily remove nodes from the tree while leaving them inside of the vector.

Ben Lovy • Dec 19 '19

Great point - this code conflates ArenaTree::size() with the length of the vector, which are separate concepts. Your implementation is absolutely more correct.

chichid • Nov 5 '22 • Edited

You sure have worked around the borrow checker issue and "sprinkling" 'a all over the place but here is what you didn't mention:

You now made it a lot easier to create memory leaks. And at best there will be nodes laying around for more time than they should be. Basically, no borrow checking, but also you need to manage that memory.
Say goodbye to performant thread safety. If you want to do some parallel processing for all the nodes in your tree your only bet now is to sprinkle Arcs all over the place. To give you an idea: Arc::new for simple struct may take 200 to 300% more per allocation. Not to mention the runtime cost of atomic resource counting.
The amount of copy that this will introduce is going to slow down your application if you have big tree (ex. a GUI).

I'd go with dealing with the Borrow checker when necessary instead of working around it. Even through it introduces ugly verbose lifecycles everywhere.

Felipe Gusmao • May 24 '23

Thanks for the article. Loved it.