197: Trampoline — Heap-Based Tail Call Optimization

// Trampoline — Heap-Based Tail Call Optimization
//
// Deep mutual recursion causes stack overflow in Rust (no TCO guarantee).
// A trampoline converts it to a heap loop: instead of recursing, return
// a "thunk" (Bounce) that the loop will call next.
//
// type Step<A> = Done(A) | More(Box<dyn FnOnce() -> Step<A>>)

// ── The Trampoline type ───────────────────────────────────────────────────

enum Step<A> {
    Done(A),
    More(Box<dyn FnOnce() -> Step<A>>),
}

impl<A> Step<A> {
    fn done(x: A) -> Self {
        Step::Done(x)
    }

    fn more<F: FnOnce() -> Step<A> + 'static>(f: F) -> Self {
        Step::More(Box::new(f))
    }
}

/// Run the trampoline loop — O(1) stack regardless of depth.
fn run<A>(mut step: Step<A>) -> A {
    loop {
        match step {
            Step::Done(x) => return x,
            Step::More(f) => step = f(),
        }
    }
}

// ── Mutual recursion: even/odd — would stack-overflow at ~100k without trampoline

fn is_even_t(n: u64) -> Step<bool> {
    if n == 0 {
        Step::done(true)
    } else {
        Step::more(move || is_odd_t(n - 1))
    }
}

fn is_odd_t(n: u64) -> Step<bool> {
    if n == 0 {
        Step::done(false)
    } else {
        Step::more(move || is_even_t(n - 1))
    }
}

// ── Factorial via trampoline ───────────────────────────────────────────────

fn factorial_t(n: u64) -> u64 {
    fn go(n: u64, acc: u64) -> Step<u64> {
        if n <= 1 {
            Step::done(acc)
        } else {
            Step::more(move || go(n - 1, n * acc))
        }
    }
    run(go(n, 1))
}

// ── Fibonacci via trampoline (CPS to avoid two recursive calls) ───────────

fn fibonacci_t(n: u64) -> u64 {
    // Use CPS to reduce to one recursive call at a time.
    fn go(n: u64, k: Box<dyn FnOnce(u64) -> Step<u64>>) -> Step<u64> {
        if n <= 1 {
            k(n)
        } else {
            Step::more(move || {
                go(n - 1, Box::new(move |a| {
                    go(n - 2, Box::new(move |b| k(a + b)))
                }))
            })
        }
    }
    run(go(n, Box::new(Step::done)))
}

// ── Ackermann via trampoline ──────────────────────────────────────────────

fn ackermann_t(m: u64, n: u64) -> u64 {
    fn go(m: u64, n: u64) -> Step<u64> {
        if m == 0 {
            Step::done(n + 1)
        } else if n == 0 {
            Step::more(move || go(m - 1, 1))
        } else {
            // ackermann(m, n) = ackermann(m-1, ackermann(m, n-1))
            // We need the inner result first — use a CPS/two-step approach.
            // Since the trampoline is single-step, we chain via a closure.
            Step::more(move || {
                // Compute inner = ackermann(m, n-1), then use it
                // We run the inner trampoline eagerly (small values only):
                let inner = run(go(m, n - 1));
                go(m - 1, inner)
            })
        }
    }
    run(go(m, n))
}

// ── Flatten a deeply-nested list via trampoline ───────────────────────────

#[derive(Debug, Clone)]
enum Tree {
    Leaf(i32),
    Node(Vec<Tree>),
}

fn flatten_t(tree: Tree) -> Vec<i32> {
    fn go(trees: Vec<Tree>, acc: Vec<i32>) -> Step<Vec<i32>> {
        if trees.is_empty() {
            Step::done(acc)
        } else {
            let mut rest = trees;
            let head = rest.remove(0);
            match head {
                Tree::Leaf(v) => {
                    let mut new_acc = acc;
                    new_acc.push(v);
                    Step::more(move || go(rest, new_acc))
                }
                Tree::Node(children) => {
                    let mut new_trees = children;
                    new_trees.extend(rest);
                    Step::more(move || go(new_trees, acc))
                }
            }
        }
    }
    run(go(vec![tree], vec![]))
}

fn main() {
    // ── Even/odd — large numbers ───────────────────────────────────────
    println!("is_even(1_000_000) = {}", run(is_even_t(1_000_000)));
    println!("is_odd(999_999)    = {}", run(is_odd_t(999_999)));

    println!();

    // ── Factorial ─────────────────────────────────────────────────────
    println!("10! = {}", factorial_t(10));
    println!("20! = {}", factorial_t(20));

    println!();

    // ── Fibonacci ─────────────────────────────────────────────────────
    for i in [0, 1, 5, 10, 15] {
        println!("fib({}) = {}", i, fibonacci_t(i));
    }

    println!();

    // ── Ackermann ─────────────────────────────────────────────────────
    println!("ack(3,4) = {}", ackermann_t(3, 4)); // 125

    println!();

    // ── Flatten tree ──────────────────────────────────────────────────
    let tree = Tree::Node(vec![
        Tree::Leaf(1),
        Tree::Node(vec![Tree::Leaf(2), Tree::Leaf(3)]),
        Tree::Node(vec![
            Tree::Node(vec![Tree::Leaf(4), Tree::Leaf(5)]),
            Tree::Leaf(6),
        ]),
    ]);
    println!("flatten = {:?}", flatten_t(tree));

    println!();
    println!("No stack overflow on deep recursion ✓");
}

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

    #[test]
    fn test_is_even_large() {
        assert!(run(is_even_t(1_000_000)));
        assert!(!run(is_even_t(999_999)));
    }

    #[test]
    fn test_is_odd_large() {
        assert!(run(is_odd_t(999_999)));
        assert!(!run(is_odd_t(1_000_000)));
    }

    #[test]
    fn test_factorial() {
        assert_eq!(factorial_t(0), 1);
        assert_eq!(factorial_t(1), 1);
        assert_eq!(factorial_t(5), 120);
        assert_eq!(factorial_t(10), 3_628_800);
    }

    #[test]
    fn test_fibonacci() {
        assert_eq!(fibonacci_t(0), 0);
        assert_eq!(fibonacci_t(1), 1);
        assert_eq!(fibonacci_t(10), 55);
        assert_eq!(fibonacci_t(15), 610);
    }

    #[test]
    fn test_ackermann() {
        assert_eq!(ackermann_t(0, 0), 1);
        assert_eq!(ackermann_t(1, 1), 3);
        assert_eq!(ackermann_t(2, 3), 9);
        assert_eq!(ackermann_t(3, 3), 61);
    }

    #[test]
    fn test_flatten_tree() {
        let tree = Tree::Node(vec![
            Tree::Leaf(1),
            Tree::Node(vec![Tree::Leaf(2), Tree::Leaf(3)]),
        ]);
        assert_eq!(flatten_t(tree), vec![1, 2, 3]);
    }
}