Global Arena Rust

global_arena.rs

use std::cell::RefCell;

// `generational_arena` is a special type of arena that allows elements to be recycled, and the
// handles allocated from the arena do not hold a borrow on the arena. The handles are later used to
// retrieve concrete borrows on the arena temporarily. The check for live-ness is essentially
// performed at runtime, rather than being checked by the borrow checker.
use generational_arena::{Arena, Index};

// Allocate `ARENA` at global/thread scope. The arena is wrapped in a RefCell so we can get
// runtime-checked mutable borrows of the arena.
thread_local! {
    static ARENA: RefCell<Arena<Node>> = RefCell::new(Arena::new());
}

/// The `Node` construct. Notice that its children are stored as `Index`es. This means you have to
/// handle these `Index` handles to the `ARENA` to get back a temporary borrow of a `Node` object.
#[derive(Debug, Default)]
struct Node {
    left: Option<Index>,
    right: Option<Index>,
}

/// Simple main function that builds a tree with a depth of 3
fn main() {
    let root = alloc_node();
    build_tree(root, 3);
}

/// Allocates a new Node. Panics if there are outstanding borrows of `ARENA`.
fn alloc_node() -> Index {
    // With `thread_local`, we can only access the global within the context of the `with` callback.
    ARENA.with(|a| {
        // `borrow_mut` performs a runtime check. It borrows the arena temporarily, inserts a node,
        // and returns the `Index`. `Index` does not store a borrow, and therefore does not retain
        // exclusive access to the Arena. However, this does mean that the borrow checking is,
        // essentially, deferred to runtime later - it's on the programmer to make sure they line
        // this all up correctly.
        a.borrow_mut().insert(Node::default())
    })
}

/// Builds a tree of nodes to `depth` starting with `node`.
fn build_tree(node: Index, depth: u32) {
    if depth == 0 {
        return;
    }

    let left = alloc_node();
    let right = alloc_node();

    build_tree(left, depth - 1);
    build_tree(right, depth - 1);

    // Here we update the child nodes. Again, the values are `Index`es - essentially handles to
    // retrieve the actual nodes at a later time.
    ARENA.with(|a| {
        let mut a = a.borrow_mut();
        let a = a.get_mut(node).unwrap();
        a.left = Some(left);
        a.right = Some(right);
    })
}