1102 Advanced

Red-Black Tree — Balanced Insert

Functional Programming

Tutorial Video

Text description (accessibility)

This video demonstrates the "Red-Black Tree — Balanced Insert" functional Rust example. Difficulty level: Advanced. Key concepts covered: Functional Programming. Implement a purely functional red-black tree supporting insert and membership test. Key difference from OCaml: 1. **Pattern depth:** OCaml matches three levels deep in one arm; Rust nests

Tutorial

The Problem

Implement a purely functional red-black tree supporting insert and membership test. Every insert returns a new tree; the old tree is unchanged. Balance is maintained after each insert using Okasaki's four-case rewrite rule.

🎯 Learning Outcomes

• How Rust enum with Box models recursive algebraic data types

• How Okasaki's elegant multi-case balance function translates from OCaml

pattern matching to nested Rust if let chains

• Why ref bindings are needed when pattern-matching on owned enum values

without consuming them

• How impl FromIterator<T> integrates a custom collection with Rust's

iterator ecosystem

🦀 The Rust Way

Rust cannot destructure deeply into Box<T> in a single match arm (no stable box patterns). Instead, the four cases are written as nested if let chains. ref bindings borrow the child nodes without moving them, keeping the owned left/right values available for the default fallthrough. Immutability is explicit — insert takes &self and returns a new RbTree<T>.

Code Example

#[derive(Debug, Clone, PartialEq)]
pub enum Color { Red, Black }

#[derive(Debug, Clone, PartialEq)]
pub enum RbTree<T> {
    E,
    T(Color, Box<RbTree<T>>, T, Box<RbTree<T>>),
}

impl<T: Ord + Clone> RbTree<T> {
    pub fn insert(&self, x: T) -> Self {
        match self.ins(&x) {
            RbTree::T(_, a, y, b) => RbTree::T(Color::Black, a, y, b),
            RbTree::E => RbTree::E,
        }
    }
}

impl<T: Ord + Clone> FromIterator<T> for RbTree<T> {
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        iter.into_iter().fold(RbTree::empty(), |t, x| t.insert(x))
    }
}

type color = Red | Black
type 'a rbtree = E | T of color * 'a rbtree * 'a * 'a rbtree

let balance = function
  | Black, T (Red, T (Red, a, x, b), y, c), z, d
  | Black, T (Red, a, x, T (Red, b, y, c)), z, d
  | Black, a, x, T (Red, T (Red, b, y, c), z, d)
  | Black, a, x, T (Red, b, y, T (Red, c, z, d)) ->
    T (Red, T (Black, a, x, b), y, T (Black, c, z, d))
  | color, a, x, b -> T (color, a, x, b)

let insert x t =
  let rec ins = function
    | E -> T (Red, E, x, E)
    | T (color, a, y, b) ->
      if x < y then balance (color, ins a, y, b)
      else if x > y then balance (color, a, y, ins b)
      else T (color, a, y, b)
  in
  match ins t with T (_, a, y, b) -> T (Black, a, y, b) | E -> E

Key Differences

Pattern depth: OCaml matches three levels deep in one arm; Rust nests

if let blocks because Box can't be destructured in a stable match.

Ownership: OCaml values are GC-managed; Rust needs Box for the

recursive tree nodes and Clone to copy subtrees that appear in both the old and new tree.

Root invariant: Both paint the root Black after insert — OCaml via a

second match ins t with, Rust via the same idiom on the ins result.

FromIterator: Rust's trait system lets RbTree participate in .collect()

idiomatically; OCaml uses List.fold_left ad hoc.

OCaml Approach

OCaml's balance is a single function on a 4-tuple (color, left, val, right). All four red-red violation patterns are listed in one match, sharing the same right-hand side. The type system ensures completeness with a catch-all arm. Immutability is free — all values are persistent by default.

Full Source

#![allow(dead_code)]
//! 1102: Red-Black Tree with Okasaki's Functional Balancing
//!
//! A purely functional red-black tree using Okasaki's elegant pattern-matching
//! balance function. Insert returns a new tree sharing structure with the old one.

#[derive(Debug, Clone, PartialEq)]
pub enum Color {
    Red,
    Black,
}

/// Purely functional red-black tree.
/// `E` = empty leaf; `T(color, left, value, right)` = internal node.
#[derive(Debug, Clone, PartialEq)]
pub enum RbTree<T> {
    E,
    T(Color, Box<RbTree<T>>, T, Box<RbTree<T>>),
}

impl<T: Ord + Clone> RbTree<T> {
    pub fn empty() -> Self {
        RbTree::E
    }

    /// Insert `x`, returning a new balanced tree. Root is always painted Black.
    pub fn insert(&self, x: T) -> Self {
        match self.ins(&x) {
            RbTree::T(_, a, y, b) => RbTree::T(Color::Black, a, y, b),
            RbTree::E => RbTree::E,
        }
    }

    /// Recursive insert that may return a Red root (fixed by `insert`).
    fn ins(&self, x: &T) -> Self {
        match self {
            RbTree::E => RbTree::T(
                Color::Red,
                Box::new(RbTree::E),
                x.clone(),
                Box::new(RbTree::E),
            ),
            // color, a, y, b are references (&Color, &Box<RbTree<T>>, &T, &Box<RbTree<T>>)
            RbTree::T(color, a, y, b) => {
                if x < y {
                    balance(color.clone(), a.ins(x), y.clone(), (**b).clone())
                } else if x > y {
                    balance(color.clone(), (**a).clone(), y.clone(), b.ins(x))
                } else {
                    self.clone() // duplicate: no change
                }
            }
        }
    }

    /// Test membership: O(log n).
    pub fn member(&self, x: &T) -> bool {
        match self {
            RbTree::E => false,
            RbTree::T(_, a, y, b) => {
                if x == y {
                    true
                } else if x < y {
                    a.member(x)
                } else {
                    b.member(x)
                }
            }
        }
    }

    /// In-order traversal — returns elements in sorted order.
    pub fn to_sorted_vec(&self) -> Vec<T> {
        match self {
            RbTree::E => vec![],
            RbTree::T(_, a, v, b) => {
                let mut out = a.to_sorted_vec();
                out.push(v.clone());
                out.extend(b.to_sorted_vec());
                out
            }
        }
    }

    /// Check that the root node is Black (RB invariant after insert).
    pub fn root_is_black(&self) -> bool {
        match self {
            RbTree::E => true,
            RbTree::T(color, _, _, _) => *color == Color::Black,
        }
    }
}

impl<T: Ord + Clone> FromIterator<T> for RbTree<T> {
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        iter.into_iter().fold(RbTree::empty(), |t, x| t.insert(x))
    }
}

/// Okasaki's balance: fix any red-red violation after insert into a Black node.
///
/// All four symmetric cases rewrite to:
///   `T(Red, T(Black, a, x, b), y, T(Black, c, z, d))`
///
/// Using `ref` bindings to borrow fields without consuming `left`/`right`,
/// so the owned values remain available for the default fallthrough.
fn balance<T: Clone>(color: Color, left: RbTree<T>, val: T, right: RbTree<T>) -> RbTree<T> {
    use Color::{Black, Red};

    if color == Black {
        // Case 1: T(Black, T(Red, T(Red,a,x,b), y,c), z, d)  — left-left
        if let RbTree::T(Red, ref ll, ref lv, ref lr) = left {
            if let RbTree::T(Red, ref a, ref x, ref b) = **ll {
                return RbTree::T(
                    Red,
                    Box::new(RbTree::T(Black, a.clone(), x.clone(), b.clone())),
                    lv.clone(),
                    Box::new(RbTree::T(Black, lr.clone(), val, Box::new(right))),
                );
            }
            // Case 2: T(Black, T(Red,a,x, T(Red,b,y,c)), z, d)  — left-right
            if let RbTree::T(Red, ref b, ref y, ref c) = **lr {
                return RbTree::T(
                    Red,
                    Box::new(RbTree::T(Black, ll.clone(), lv.clone(), b.clone())),
                    y.clone(),
                    Box::new(RbTree::T(Black, c.clone(), val, Box::new(right))),
                );
            }
        }
        // Case 3: T(Black, a, x, T(Red, T(Red,b,y,c), z,d))  — right-left
        if let RbTree::T(Red, ref rl, ref rv, ref rr) = right {
            if let RbTree::T(Red, ref b, ref y, ref c) = **rl {
                return RbTree::T(
                    Red,
                    Box::new(RbTree::T(Black, Box::new(left), val, b.clone())),
                    y.clone(),
                    Box::new(RbTree::T(Black, c.clone(), rv.clone(), rr.clone())),
                );
            }
            // Case 4: T(Black, a, x, T(Red, b, y, T(Red,c,z,d)))  — right-right
            if let RbTree::T(Red, ref c, ref z, ref d) = **rr {
                return RbTree::T(
                    Red,
                    Box::new(RbTree::T(Black, Box::new(left), val, rl.clone())),
                    rv.clone(),
                    Box::new(RbTree::T(Black, c.clone(), z.clone(), d.clone())),
                );
            }
        }
    }

    RbTree::T(color, Box::new(left), val, Box::new(right))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_empty_tree() {
        let t: RbTree<i32> = RbTree::empty();
        assert!(!t.member(&1));
        assert_eq!(t.to_sorted_vec(), Vec::<i32>::new());
        assert!(t.root_is_black());
    }

    #[test]
    fn test_single_element() {
        let t = RbTree::empty().insert(42);
        assert!(t.member(&42));
        assert!(!t.member(&0));
        assert_eq!(t.to_sorted_vec(), vec![42]);
        assert!(t.root_is_black());
    }

    #[test]
    fn test_sorted_order_multiple_elements() {
        let t = RbTree::from_iter([5, 3, 7, 1, 4, 6, 8, 2, 9]);
        assert_eq!(t.to_sorted_vec(), vec![1, 2, 3, 4, 5, 6, 7, 8, 9]);
    }

    #[test]
    fn test_duplicates_ignored() {
        let t = RbTree::from_iter([3, 1, 4, 1, 5, 9, 2, 6, 5]);
        assert_eq!(t.to_sorted_vec(), vec![1, 2, 3, 4, 5, 6, 9]);
    }

    #[test]
    fn test_root_always_black() {
        let t = RbTree::from_iter([5, 3, 7, 1, 4, 6, 8]);
        assert!(t.root_is_black());
    }

    #[test]
    fn test_member_all_inserted_values() {
        let values = [5, 3, 7, 1, 4, 6, 8, 2, 9];
        let t = RbTree::from_iter(values);
        for v in values {
            assert!(t.member(&v));
        }
        assert!(!t.member(&10));
        assert!(!t.member(&0));
    }

    #[test]
    fn test_ascending_insert_stays_balanced() {
        // Worst case for naive BST: ascending order. RB tree must stay balanced.
        let t = RbTree::from_iter(1..=16);
        assert_eq!(t.to_sorted_vec(), (1..=16).collect::<Vec<_>>());
        assert!(t.root_is_black());
    }

    #[test]
    fn test_string_tree() {
        let t = RbTree::from_iter(["banana", "apple", "cherry", "date"]);
        assert_eq!(t.to_sorted_vec(), vec!["apple", "banana", "cherry", "date"]);
        assert!(t.member(&"apple"));
        assert!(!t.member(&"elderberry"));
    }
}

type color = Red | Black
type 'a rbtree = E | T of color * 'a rbtree * 'a * 'a rbtree

(* Okasaki's 4-case balance — all cases produce T(Red, T(Black,a,x,b), y, T(Black,c,z,d)) *)
let balance = function
  | Black, T (Red, T (Red, a, x, b), y, c), z, d
  | Black, T (Red, a, x, T (Red, b, y, c)), z, d
  | Black, a, x, T (Red, T (Red, b, y, c), z, d)
  | Black, a, x, T (Red, b, y, T (Red, c, z, d)) ->
    T (Red, T (Black, a, x, b), y, T (Black, c, z, d))
  | color, a, x, b -> T (color, a, x, b)

(* Recursive insert: new leaf is Red, balance repairs any double-Red violation *)
let insert x t =
  let rec ins = function
    | E -> T (Red, E, x, E)
    | T (color, a, y, b) ->
      if x < y then balance (color, ins a, y, b)
      else if x > y then balance (color, a, y, ins b)
      else T (color, a, y, b)
  in
  (* Always blacken the root *)
  match ins t with T (_, a, y, b) -> T (Black, a, y, b) | E -> E

let rec mem x = function
  | E -> false
  | T (_, a, y, b) -> x = y || (if x < y then mem x a else mem x b)

let rec to_list = function
  | E -> []
  | T (_, a, v, b) -> to_list a @ [v] @ to_list b

let () =
  let t = List.fold_left (fun t x -> insert x t) E [5; 3; 7; 1; 4; 6; 8; 2; 9] in
  assert (to_list t = [1; 2; 3; 4; 5; 6; 7; 8; 9]);
  assert (mem 5 t);
  assert (mem 1 t);
  assert (mem 9 t);
  assert (not (mem 0 t));
  assert (not (mem 10 t));

  (* Root must always be Black *)
  (match t with T (Black, _, _, _) -> () | _ -> assert false);

  (* Duplicates are ignored: set semantics *)
  let t2 = List.fold_left (fun t x -> insert x t) E [3; 3; 1; 2; 1] in
  assert (to_list t2 = [1; 2; 3]);

  (* Sorted insertion exercises all balance cases *)
  let t3 = List.fold_left (fun t x -> insert x t) E [1; 2; 3; 4; 5] in
  assert (to_list t3 = [1; 2; 3; 4; 5]);

  print_endline "ok"

✓ Tests Rust test suite

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_empty_tree() {
        let t: RbTree<i32> = RbTree::empty();
        assert!(!t.member(&1));
        assert_eq!(t.to_sorted_vec(), Vec::<i32>::new());
        assert!(t.root_is_black());
    }

    #[test]
    fn test_single_element() {
        let t = RbTree::empty().insert(42);
        assert!(t.member(&42));
        assert!(!t.member(&0));
        assert_eq!(t.to_sorted_vec(), vec![42]);
        assert!(t.root_is_black());
    }

    #[test]
    fn test_sorted_order_multiple_elements() {
        let t = RbTree::from_iter([5, 3, 7, 1, 4, 6, 8, 2, 9]);
        assert_eq!(t.to_sorted_vec(), vec![1, 2, 3, 4, 5, 6, 7, 8, 9]);
    }

    #[test]
    fn test_duplicates_ignored() {
        let t = RbTree::from_iter([3, 1, 4, 1, 5, 9, 2, 6, 5]);
        assert_eq!(t.to_sorted_vec(), vec![1, 2, 3, 4, 5, 6, 9]);
    }

    #[test]
    fn test_root_always_black() {
        let t = RbTree::from_iter([5, 3, 7, 1, 4, 6, 8]);
        assert!(t.root_is_black());
    }

    #[test]
    fn test_member_all_inserted_values() {
        let values = [5, 3, 7, 1, 4, 6, 8, 2, 9];
        let t = RbTree::from_iter(values);
        for v in values {
            assert!(t.member(&v));
        }
        assert!(!t.member(&10));
        assert!(!t.member(&0));
    }

    #[test]
    fn test_ascending_insert_stays_balanced() {
        // Worst case for naive BST: ascending order. RB tree must stay balanced.
        let t = RbTree::from_iter(1..=16);
        assert_eq!(t.to_sorted_vec(), (1..=16).collect::<Vec<_>>());
        assert!(t.root_is_black());
    }

    #[test]
    fn test_string_tree() {
        let t = RbTree::from_iter(["banana", "apple", "cherry", "date"]);
        assert_eq!(t.to_sorted_vec(), vec!["apple", "banana", "cherry", "date"]);
        assert!(t.member(&"apple"));
        assert!(!t.member(&"elderberry"));
    }
}

Deep Comparison

OCaml vs Rust: Red-Black Tree with Okasaki's Functional Balancing

Side-by-Side Code

OCaml

type color = Red | Black
type 'a rbtree = E | T of color * 'a rbtree * 'a * 'a rbtree

let balance = function
  | Black, T (Red, T (Red, a, x, b), y, c), z, d
  | Black, T (Red, a, x, T (Red, b, y, c)), z, d
  | Black, a, x, T (Red, T (Red, b, y, c), z, d)
  | Black, a, x, T (Red, b, y, T (Red, c, z, d)) ->
    T (Red, T (Black, a, x, b), y, T (Black, c, z, d))
  | color, a, x, b -> T (color, a, x, b)

let insert x t =
  let rec ins = function
    | E -> T (Red, E, x, E)
    | T (color, a, y, b) ->
      if x < y then balance (color, ins a, y, b)
      else if x > y then balance (color, a, y, ins b)
      else T (color, a, y, b)
  in
  match ins t with T (_, a, y, b) -> T (Black, a, y, b) | E -> E

Rust (idiomatic — implements `FromIterator`, integrates with `.collect()`)

#[derive(Debug, Clone, PartialEq)]
pub enum Color { Red, Black }

#[derive(Debug, Clone, PartialEq)]
pub enum RbTree<T> {
    E,
    T(Color, Box<RbTree<T>>, T, Box<RbTree<T>>),
}

impl<T: Ord + Clone> RbTree<T> {
    pub fn insert(&self, x: T) -> Self {
        match self.ins(&x) {
            RbTree::T(_, a, y, b) => RbTree::T(Color::Black, a, y, b),
            RbTree::E => RbTree::E,
        }
    }
}

impl<T: Ord + Clone> FromIterator<T> for RbTree<T> {
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        iter.into_iter().fold(RbTree::empty(), |t, x| t.insert(x))
    }
}

Rust (Okasaki balance — four nested `if let` cases)

fn balance<T: Clone>(color: Color, left: RbTree<T>, val: T, right: RbTree<T>) -> RbTree<T> {
    use Color::{Black, Red};

    if color == Black {
        // Case 1: left-left red violation
        if let RbTree::T(Red, ref ll, ref lv, ref lr) = left {
            if let RbTree::T(Red, ref a, ref x, ref b) = **ll {
                return RbTree::T(Red,
                    Box::new(RbTree::T(Black, a.clone(), x.clone(), b.clone())),
                    lv.clone(),
                    Box::new(RbTree::T(Black, lr.clone(), val, Box::new(right))));
            }
            // Case 2: left-right red violation
            if let RbTree::T(Red, ref b, ref y, ref c) = **lr {
                return RbTree::T(Red,
                    Box::new(RbTree::T(Black, ll.clone(), lv.clone(), b.clone())),
                    y.clone(),
                    Box::new(RbTree::T(Black, c.clone(), val, Box::new(right))));
            }
        }
        // Case 3: right-left red violation
        if let RbTree::T(Red, ref rl, ref rv, ref rr) = right {
            if let RbTree::T(Red, ref b, ref y, ref c) = **rl {
                return RbTree::T(Red,
                    Box::new(RbTree::T(Black, Box::new(left), val, b.clone())),
                    y.clone(),
                    Box::new(RbTree::T(Black, c.clone(), rv.clone(), rr.clone())));
            }
            // Case 4: right-right red violation
            if let RbTree::T(Red, ref c, ref z, ref d) = **rr {
                return RbTree::T(Red,
                    Box::new(RbTree::T(Black, Box::new(left), val, rl.clone())),
                    rv.clone(),
                    Box::new(RbTree::T(Black, c.clone(), z.clone(), d.clone())));
            }
        }
    }
    RbTree::T(color, Box::new(left), val, Box::new(right))
}

Type Signatures

Concept	OCaml	Rust
Tree type	`'a rbtree`	`RbTree<T>`
Empty node	`E`	`RbTree::E`
Internal node	`T of color * 'a rbtree * 'a * 'a rbtree`	`T(Color, Box<RbTree<T>>, T, Box<RbTree<T>>)`
Insert signature	`val insert : 'a -> 'a rbtree -> 'a rbtree`	`fn insert(&self, x: T) -> Self`
Balance signature	`(color * 'a rbtree * 'a * 'a rbtree) -> 'a rbtree`	`fn balance<T: Clone>(Color, RbTree<T>, T, RbTree<T>) -> RbTree<T>`
Membership	`val mem : 'a -> 'a rbtree -> bool`	`fn member(&self, x: &T) -> bool`
Collection building	`List.fold_left (fun t x -> insert x t) E xs`	`iter.collect::<RbTree<_>>()`

Key Insights

Box for recursion: OCaml's GC manages recursive types transparently.

Rust requires Box<RbTree<T>> to give the recursive type a known size on the stack. Every child node is heap-allocated.

**ref bindings preserve ownership:** OCaml patterns always bind by value

(GC semantics). In Rust, matching on an owned enum would move its fields. ref ll borrows instead, keeping left intact for the default fallthrough. NLL (non-lexical lifetimes) ensures those borrows end before the owned left is moved in cases 3 and 4.

Box prevents deep match patterns: OCaml's single function arm matches

three levels deep in one expression. Rust stable cannot destructure through Box in a match pattern, so the four cases become nested if let chains — same logic, more syntactic scaffolding.

**Clone as the price of persistence:** Okasaki's trees share subtrees

structurally. In Rust, every node along the path from root to the insert point is rebuilt (cloned), while unchanged subtrees are shared via Box (reference counting would also work but isn't needed here).

**FromIterator vs fold_left:** OCaml uses List.fold_left ad hoc.

Rust's trait system lets RbTree implement FromIterator<T>, making it a first-class collection that works with .collect() and iterator adapters — a clean integration with the standard library.

When to Use Each Style

**Use idiomatic Rust (FromIterator + insert)** when building a tree from a data source — ranges, vecs, file lines. The .collect() call is zero-overhead compared to a manual fold and expresses intent clearly.

**Use the recursive ins + balance style** when you want to see the direct structural parallel to Okasaki's original formulation — useful for teaching, porting other purely functional algorithms, or verifying correctness against the textbook.

Exercises

Implement contains on the red-black tree and write a test that verifies every inserted element can be found and no non-inserted element is found.

Add a method black_height that returns the number of black nodes on any root-to-leaf path and assert it is the same for all paths.

Implement a from_iter constructor that builds a balanced red-black tree from an unsorted iterator and verify it produces a valid BST by checking that to_sorted_vec yields elements in order.

Open Source Repos

functional-rust

View the source for this example on GitHub — OCaml and Rust side by side in the repo.

Rust