hegeltest 0.32.4

use super::{Generator, TestCase};
use crate::control::LeafBudgetExceeded;
use crate::ffi::RecursionHandle;
use crate::test_case::{labels, raise_for_rc};
use std::marker::PhantomData;
use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
use std::sync::Arc;

const DEFAULT_MAX_DEPTH: usize = 32;
const DEFAULT_MAX_LEAVES: usize = 100;

/// The leaf generator and branch function of a [`recursive()`] generator,
/// type-erased so that [`SubtreeGenerator`] (which appears in the branch
/// function's own signature) does not need to name their types.
trait SubtreeDraw<T>: Send + Sync {
    fn draw_leaf(&self, tc: &TestCase) -> T;
    fn draw_branch(&self, tc: &TestCase, subtrees: SubtreeGenerator<T>) -> T;
}

struct RecursiveCore<G, F, R> {
    leaf: G,
    branch: F,
    _phantom: PhantomData<fn() -> R>,
}

impl<T, G, F, R> SubtreeDraw<T> for RecursiveCore<G, F, R>
where
    G: Generator<T> + Send + Sync,
    F: Fn(SubtreeGenerator<T>) -> R + Send + Sync,
    R: Generator<T>,
{
    fn draw_leaf(&self, tc: &TestCase) -> T {
        self.leaf.do_draw(tc)
    }

    fn draw_branch(&self, tc: &TestCase, subtrees: SubtreeGenerator<T>) -> T {
        (self.branch)(subtrees).do_draw(tc)
    }
}

/// The generator a [`recursive()`] branch function receives, producing the
/// recursive sub-values of the value under construction.
///
/// Each value it generates is itself either a leaf or a further branch. It
/// is `Clone`, so a branch function needing several sub-value generators
/// (e.g. for the fields of a [`tuples!`](crate::tuples)) can clone it.
/// Cloning is needed rather than borrowing (`tuples!(&subtrees, &subtrees)`)
/// because the generator the branch function returns would otherwise borrow
/// the function's own parameter.
pub struct SubtreeGenerator<T> {
    core: Arc<dyn SubtreeDraw<T>>,
    recursion: Arc<RecursionHandle>,
    depth: u64,
}

impl<T> Clone for SubtreeGenerator<T> {
    fn clone(&self) -> Self {
        SubtreeGenerator {
            core: Arc::clone(&self.core),
            recursion: Arc::clone(&self.recursion),
            depth: self.depth,
        }
    }
}

impl<T> SubtreeGenerator<T> {
    fn child(&self) -> Self {
        SubtreeGenerator {
            core: Arc::clone(&self.core),
            recursion: Arc::clone(&self.recursion),
            depth: self.depth + 1,
        }
    }
}

impl<T> Generator<T> for SubtreeGenerator<T> {
    fn do_draw(&self, tc: &TestCase) -> T {
        tc.start_span(labels::RECURSIVE);
        let branch = match tc.with_ctc(|ctc| ctc.recursion_branch(&self.recursion, self.depth)) {
            Ok(branch) => branch,
            Err(rc) => raise_for_rc(rc),
        };
        let result = if branch {
            self.core.draw_branch(tc, self.child())
        } else {
            if let Err(rc) = tc.with_ctc(|ctc| ctc.recursion_leaf(&self.recursion)) {
                raise_for_rc(rc);
            }
            self.core.draw_leaf(tc)
        };
        tc.stop_span(false);
        result
    }
}

/// Generator for recursively defined data. Created by [`recursive()`].
pub struct RecursiveGenerator<T> {
    core: Arc<dyn SubtreeDraw<T>>,
    max_depth: usize,
    max_leaves: usize,
}

impl<T> RecursiveGenerator<T> {
    /// Set the maximum nesting depth of branches (default 32).
    ///
    /// Sub-values at this depth are always leaves, so a `max_depth` of 0
    /// generates only leaves.
    pub fn max_depth(mut self, max_depth: usize) -> Self {
        self.max_depth = max_depth;
        self
    }

    /// Set the maximum number of leaf values in one generated value
    /// (default 100).
    ///
    /// A generation attempt that draws more than `max_leaves` leaves is
    /// discarded and retried with a lower branching probability; if several
    /// retries in a row fail to fit, the test case is rejected as if by
    /// [`assume`](crate::TestCase::assume).
    pub fn max_leaves(mut self, max_leaves: usize) -> Self {
        self.max_leaves = max_leaves;
        self
    }
}

impl<T> Generator<T> for RecursiveGenerator<T> {
    fn do_draw(&self, tc: &TestCase) -> T {
        let base_span_depth = tc.open_span_depth();
        let recursion = match tc
            .with_ctc(|ctc| ctc.new_recursion(self.max_depth as u64, self.max_leaves as u64))
        {
            Ok(recursion) => Arc::new(recursion),
            Err(rc) => raise_for_rc(rc),
        };
        loop {
            let root = SubtreeGenerator {
                core: Arc::clone(&self.core),
                recursion: Arc::clone(&recursion),
                depth: 0,
            };
            match catch_unwind(AssertUnwindSafe(|| root.do_draw(tc))) {
                Ok(value) => return value,
                Err(payload) if payload.downcast_ref::<LeafBudgetExceeded>().is_some() => {
                    match tc.with_ctc(|ctc| ctc.recursion_retry(&recursion)) {
                        Ok(()) => tc.reset_open_spans_to(base_span_depth),
                        Err(rc) => raise_for_rc(rc),
                    }
                }
                Err(payload) => resume_unwind(payload),
            }
        }
    }
}

/// Generate recursively defined data, such as trees or JSON documents.
///
/// `leaf` generates the non-recursive base cases. `branch` builds one level
/// of recursive structure: it receives a [`SubtreeGenerator`] producing
/// sub-values of the same type and returns a generator that combines some
/// number of them into a compound value, e.g. by collecting them with
/// [`vecs()`](super::vecs) and mapping the result into a node type. It is
/// called afresh for each branch node generated.
///
/// Generated values are leaves or branches of leaves, branches of those, and
/// so on. Sizes vary from a single leaf up to the limits set by
/// [`max_depth`](RecursiveGenerator::max_depth) (a hard depth cap) and
/// [`max_leaves`](RecursiveGenerator::max_leaves) (attempts that draw more
/// leaves than this are discarded and retried with a lower branching
/// probability).
///
/// # Example
///
/// ```no_run
/// use hegel::generators::{self as gs, Generator};
///
/// #[derive(Debug)]
/// enum Json {
///     Number(f64),
///     Array(Vec<Json>),
/// }
///
/// #[hegel::test]
/// fn my_test(tc: hegel::TestCase) {
///     let value = tc.draw(gs::recursive(
///         gs::floats::<f64>().map(Json::Number),
///         |json| gs::vecs(json).max_size(5).map(Json::Array),
///     ));
/// }
/// ```
pub fn recursive<T, G, F, R>(leaf: G, branch: F) -> RecursiveGenerator<T>
where
    T: 'static,
    G: Generator<T> + Send + Sync + 'static,
    F: Fn(SubtreeGenerator<T>) -> R + Send + Sync + 'static,
    R: Generator<T> + 'static,
{
    RecursiveGenerator {
        core: Arc::new(RecursiveCore {
            leaf,
            branch,
            _phantom: PhantomData,
        }),
        max_depth: DEFAULT_MAX_DEPTH,
        max_leaves: DEFAULT_MAX_LEAVES,
    }
}