implement try_unwrap to allow a custom Drop to drop iteratively

[HuStmpHrrr]

introduce try_unwrap to implement downstream Drop to drop an object iteratively.

[HuStmpHrrr]

This program demonstrates that deeply nested objects cause a stack overflow. There are different ways to fix it but it would be helpful if this crate exposes try_unwrap.

//! Deallocating deeply nested hashconsed elements without overflowing the stack.
//!
//! Dropping a `HConsed<T>` whose `T` holds further `HConsed<T>`s drops the element in place, which
//! drops its children, and so on: the native stack grows with the *depth* of the element. Deep
//! enough elements therefore abort the process with a stack overflow, which is not a catchable
//! panic.
//!
//! ```sh
//! cargo run --example deep_drop                  # aborts while dropping the term
//! cargo run --example deep_drop consign          # aborts while dropping the consign
//! cargo run --example deep_drop pinned           # term drop is safe at any depth
//! cargo run --example deep_drop iterative        # both orders, no recursion
//! cargo run --example deep_drop iterative 1000000
//! ```

use hashconsing::{HConsign, HashConsign};

/// Deliberately small, so the overflow does not depend on the platform's default thread stack.
const STACK: usize = 1024 * 1024;
const DEPTH: usize = 100_000;

/// Nesting without a `Drop` implementation: the element is dropped in place, recursively.
mod recursive {
    use hashconsing::HConsed;

    pub type Term = HConsed<Node>;

    #[derive(Hash, Clone, PartialEq, Eq)]
    pub struct Node(pub Payload);

    #[derive(Hash, Clone, PartialEq, Eq)]
    pub enum Payload {
        Var(usize),
        Lam(Term),
    }
}

/// The same nesting, with a `Drop` implementation that frees the children iteratively.
mod iterative {
    use hashconsing::HConsed;

    pub type Term = HConsed<Node>;

    /// The `Drop` implementation lives on this wrapper rather than on [`Payload`] itself, because
    /// a child cannot be moved out of a type that implements `Drop` (E0509) — and moving the
    /// children out is precisely what has to happen before the parent is dropped.
    #[derive(Hash, Clone, PartialEq, Eq)]
    pub struct Node(pub Payload);

    #[derive(Hash, Clone, PartialEq, Eq)]
    pub enum Payload {
        Var(usize),
        Lam(Term),
    }

    /// Move the child element, if any, out of `payload` and onto `sink`.
    fn take_child(payload: &mut Payload, sink: &mut Vec<Term>) {
        match std::mem::replace(payload, Payload::Var(0)) {
            Payload::Lam(child) => sink.push(child),
            Payload::Var(_) => {}
        }
    }

    impl Drop for Node {
        fn drop(&mut self) {
            // `Vec::new` does not allocate, so a childless node costs nothing here.
            let mut worklist = Vec::new();
            take_child(&mut self.0, &mut worklist);
            while let Some(handle) = worklist.pop() {
                // Sole owner: take the element apart before it is dropped. Otherwise dropping
                // `handle` is just a reference count decrement, with nothing to free.
                if let Ok(mut node) = HConsed::try_unwrap(handle) {
                    take_child(&mut node.0, &mut worklist);
                    // `node` is dropped here, but it is childless by now: its own `Drop` finds an
                    // empty worklist and returns immediately. The stack never grows.
                }
            }
        }
    }
}

/// Build `Lam(Lam(... Var(0) ...))`, `depth` deep.
macro_rules! chain {
    ($m:ident, $depth:expr) => {{
        let mut factory: HConsign<$m::Node> = HConsign::empty();
        let mut term = factory.mk($m::Node($m::Payload::Var(0)));
        for _ in 0..$depth {
            term = factory.mk($m::Node($m::Payload::Lam(term)));
        }
        println!(
            "built a chain of depth {} ({} nodes)",
            $depth,
            factory.len()
        );
        (factory, term)
    }};
}

/// Drop the consign first, then the term: the overflow lands on the term.
fn run_term(depth: usize) {
    let (factory, term) = chain!(recursive, depth);

    // Frees nothing: every node is still reachable from `term`.
    drop(factory);
    println!("dropping the term (this is expected to overflow the stack) ...");
    drop(term);
    println!("... survived, so the chain was not deep enough; try a larger depth");
}

/// Drop the term first, then the consign: the overflow lands on the consign.
fn run_consign(depth: usize) {
    let (factory, term) = chain!(recursive, depth);

    // Frees the root node only: the consign's key for it still owns the rest of the chain.
    drop(term);
    println!("dropping the consign (this is expected to overflow the stack) ...");
    drop(factory);
    println!("... survived, so the chain was not deep enough; try a larger depth");
}

/// Drop the term while the consign is alive: nothing cascades, however deep the term is.
fn run_pinned(depth: usize) {
    let (mut factory, term) = chain!(recursive, depth);

    let child_count = match &term.get().0 {
        recursive::Payload::Lam(child) => child.arc_count(),
        recursive::Payload::Var(_) => 0,
    };
    println!(
        "root has {} strong reference (yours), its child has {child_count} \
         (the root's element, plus the consign's key for the root)",
        term.arc_count(),
    );

    println!("dropping the term ...");
    drop(term);
    println!(
        "... survived: only the root was freed, and the consign still has {} entries",
        factory.len()
    );

    // Freeing the rest needs the keys gone. `collect` does that without ever cascading, since
    // each removed key frees exactly one node, whose own key still pins its children — so it is
    // stack-safe however deep the chain is, at the cost of one pass per level.
    println!("collecting the consign ...");
    let start = std::time::Instant::now();
    factory.collect();
    println!(
        "... collected in {:?}, {} entries left: stack-safe, but quadratic in the depth",
        start.elapsed(),
        factory.len()
    );
}

/// The same two orders, with the iterative `Drop` implementation.
fn run_iterative(depth: usize) {
    let (factory, term) = chain!(iterative, depth);
    drop(factory);
    println!("dropping the term ...");
    drop(term);
    println!("... done, no recursion");

    let (factory, term) = chain!(iterative, depth);
    drop(term);
    println!("dropping the consign ...");
    drop(factory);
    println!("... done, no recursion");
}

fn main() {
    let mode = std::env::args().nth(1).unwrap_or_else(|| "term".into());
    let depth = std::env::args()
        .nth(2)
        .map(|d| d.parse().expect("depth must be a number"))
        .unwrap_or(DEPTH);

    let run: fn(usize) = match mode.as_str() {
        "term" => run_term,
        "consign" => run_consign,
        "pinned" => run_pinned,
        "iterative" => run_iterative,
        other => {
            eprintln!(
                "unknown mode `{other}`, expected `term`, `consign`, `pinned` or `iterative`"
            );
            std::process::exit(2)
        }
    };

    println!("mode: {mode}, stack: {} KiB", STACK / 1024);
    std::thread::Builder::new()
        .name(mode)
        .stack_size(STACK)
        .spawn(move || run(depth))
        .expect("failed to spawn thread")
        .join()
        .expect("thread panicked")
}