hegeltest 0.35.0

//! Stateful (model-based) testing support.
//!
//! State machines are defined using the [`state_machine`](crate::state_machine) attribute macro.
//! Methods annotated with `#[rule]` become rules (actions applied to the state machine) and
//! methods annotated with `#[invariant]` become invariants (checked on the machine's initial
//! and final state, and sampled in between — each invariant runs after any given rule with
//! probability `1 / stateful_step_count`, so its expected cost per test case stays constant
//! as the step count grows). Both take a [`TestCase`] parameter and borrow the state machine: rules
//! typically have signature `fn(&mut self, tc: TestCase)` and invariants
//! `fn(&self, tc: TestCase)`, but either kind of method may use `&self` or `&mut self`.
//!
//! To run a state machine, call [`run()`] inside a Hegel test.
//!
//! Example:
//! ```rust
//! use hegel::TestCase;
//! use hegel::generators as gs;
//!
//! struct IntegerStack {
//!     stack: Vec<i32>,
//! }
//!
//! #[hegel::state_machine]
//! impl IntegerStack {
//!     #[rule]
//!     fn push(&mut self, tc: TestCase) {
//!         let integers = gs::integers::<i32>;
//!         let element = tc.draw(integers());
//!         self.stack.push(element);
//!     }
//!
//!     #[rule]
//!     fn pop(&mut self, _: TestCase) {
//!         self.stack.pop();
//!     }
//!
//!     #[rule]
//!     fn pop_push(&mut self, tc: TestCase) {
//!         let integers = gs::integers::<i32>;
//!         let element = tc.draw(integers());
//!         let initial = self.stack.clone();
//!         self.stack.push(element);
//!         let popped = self.stack.pop().unwrap();
//!         assert_eq!(popped, element);
//!         assert_eq!(self.stack, initial);
//!     }
//!
//!     #[rule]
//!     fn push_pop(&mut self, tc: TestCase) {
//!         let initial = self.stack.clone();
//!         let element = self.stack.pop();
//!         tc.assume(element.is_some());
//!         let element = element.unwrap();
//!         self.stack.push(element);
//!         assert_eq!(self.stack, initial);
//!     }
//!
//!     #[invariant]
//!     fn len_agrees_with_is_empty(&self, _: TestCase) {
//!         assert_eq!(self.stack.is_empty(), self.stack.len() == 0);
//!     }
//! }
//!
//! #[hegel::test]
//! fn test_integer_stack(tc: TestCase) {
//!     let stack = IntegerStack { stack: Vec::new() };
//!     hegel::stateful::run(stack, tc);
//! }
//! ```
//!
//! # Assumptions in rules
//!
//! A violated assumption inside a rule is not transactional: anything the
//! rule did before the assumption failed has already happened and is not
//! rolled back. Always place assumptions at the start of the rule, before
//! mutating the state machine or the system under test. This applies
//! equally to draws from [`filter`](crate::generators::Generator::filter)ed
//! generators, which reject the rule like a failed assumption when they run
//! out of retries.
//!
//! # Concurrent state machines
//!
//! Concurrent state machines are defined using the
//! [`concurrent_state_machine`](crate::concurrent_state_machine) attribute macro. These work
//! similarly, but the rules are run concurrently from a number of worker threads. See
//! [`run_concurrent()`] for a detailed explanation of the execution model.
//!
//! Example:
//! ```rust
//! use std::sync::Mutex;
//! use std::sync::atomic::{AtomicUsize, Ordering};
//! use hegel::TestCase;
//! use hegel::generators as gs;
//!
//! struct SharedStack {
//!     stack: Mutex<Vec<i32>>,
//!     pushes: AtomicUsize,
//! }
//!
//! #[hegel::concurrent_state_machine]
//! impl SharedStack {
//!     #[rule(group = "ops")]
//!     fn push(&self, tc: TestCase) {
//!         let element = tc.draw(gs::integers::<i32>());
//!         self.stack.lock().unwrap_or_else(|e| e.into_inner()).push(element);
//!         self.pushes.fetch_add(1, Ordering::SeqCst);
//!     }
//!
//!     #[rule(group = "ops")]
//!     fn pop(&self, _: TestCase) {
//!         self.stack.lock().unwrap_or_else(|e| e.into_inner()).pop();
//!     }
//!
//!     #[rule(group = "audit")]
//!     fn audit(&self, tc: TestCase) {
//!         let stack = self.stack.lock().unwrap_or_else(|e| e.into_inner());
//!         tc.note(&format!("stack holds {} elements", stack.len()));
//!     }
//!
//!     #[invariant]
//!     fn never_grows_past_pushes(&self, _: TestCase) {
//!         let len = self.stack.lock().unwrap_or_else(|e| e.into_inner()).len();
//!         assert!(len <= self.pushes.load(Ordering::SeqCst));
//!     }
//! }
//!
//! #[hegel::test]
//! fn test_shared_stack(tc: TestCase) {
//!     let stack = SharedStack {
//!         stack: Mutex::new(Vec::new()),
//!         pushes: AtomicUsize::new(0),
//!     };
//!     hegel::stateful::run_concurrent(
//!         stack,
//!         tc,
//!         1,          // minimum concurrency
//!         3           // maximum concurrency
//!     );
//! }
//! ```

use crate::TestCase;
use crate::control::{
    AssumeFailed, InternalError, InvalidArgument, LoopDone, StopTest, hegel_internal_assert,
    raise_control, with_test_context,
};
use crate::ffi::{PoolHandle, StateMachineHandle};
use crate::generators::Generator;
use crate::run_lifecycle::{self, PanicInfo};
use crate::test_case::{labels, raise_for_rc};
use std::cell::{Cell, RefCell};
use std::collections::HashMap;
use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
use std::sync::{Mutex, mpsc};

/// The concurrency group a `#[rule]` without a `group = "..."` argument is
/// assigned to. All unannotated rules share this one group, so a machine
/// with no group annotations is maximally concurrent — and in a machine
/// that mixes annotated and unannotated rules, the unannotated rules never
/// overlap with any named group's rules (see
/// [`ConcurrentStateMachine`]).
pub const ANONYMOUS_GROUP: &str = "<anonymous>";

thread_local! {
    /// The worker-thread index [`run_concurrent`] assigns to each of its
    /// worker threads, used to tag that worker's draw/note output lines.
    /// `None` outside a concurrent stateful worker.
    static WORKER_INDEX: Cell<Option<usize>> = const { Cell::new(None) };
}

/// The calling thread's concurrent-worker index, if it is one of
/// [`run_concurrent`]'s worker threads. Read by the output machinery to tag
/// each worker line with its worker's index and time offset, and to regroup
/// the failure report's buffered lines worker by worker within each round.
pub(crate) fn current_worker_index() -> Option<usize> {
    WORKER_INDEX.with(|cell| cell.get())
}

/// A rule that can be applied to the state machine during testing.
pub struct Rule<M: ?Sized> {
    pub name: String,
    pub apply: fn(&mut M, TestCase),
}

impl<M> Rule<M> {
    /// Create a new rule with a name and an apply function.
    pub fn new(name: &str, apply: fn(&mut M, TestCase)) -> Self {
        Rule {
            name: name.to_string(),
            apply,
        }
    }
}

/// A pool of previously generated values.
///
/// Create one with [`pool()`] and populate it with [`add`](Pool::add). To draw
/// from the pool, use the generators it hands out rather than reading from it
/// directly:
///
/// - [`values_reusable`](Pool::values_reusable) returns a generator over `&T` —
///   drawing from it yields a reference to a value in the pool without removing
///   it.
/// - [`values_consumed`](Pool::values_consumed) returns a generator over `T` —
///   drawing from it removes a value from the pool and yields it by value.
///
/// Both generators are used through [`tc.draw`](TestCase::draw), so the chosen
/// value is recorded in the failing-test replay and the choice shrinks like any
/// other draw.
pub struct Pool<T> {
    pool: crate::ffi::PoolHandle,
    tc: TestCase,
    values: HashMap<i64, T>,
}

/// Ask the engine for a variable id from `pool`, consuming it if `consume`.
fn pool_generate(tc: &TestCase, pool: &crate::ffi::PoolHandle, consume: bool) -> i64 {
    match tc.with_ctc(|ctc| ctc.pool_generate(pool, consume)) {
        Ok(id) => id,
        Err(rc) => raise_for_rc(rc),
    }
}

impl<T> Pool<T> {
    /// Returns true if no values are in the pool.
    pub fn is_empty(&self) -> bool {
        self.values.is_empty()
    }

    /// Number of values currently in the pool.
    pub fn len(&self) -> usize {
        self.values.len()
    }

    /// Add a value to the pool.
    pub fn add(&mut self, v: T) {
        let variable_id: i64 = match self.tc.with_ctc(|ctc| ctc.pool_add(&self.pool)) {
            Ok(id) => id,
            Err(rc) => raise_for_rc(rc),
        };
        if self.values.contains_key(&variable_id) {
            panic!("unexpected variable id in map"); // nocov
        }
        self.values.insert(variable_id, v);
    }

    /// A generator over references to values in the pool.
    ///
    /// Drawing from it yields a `&T` borrowing a value in the pool, without
    /// removing it. Drawing rejects the current test case (as if by
    /// `assume(false)`) when the pool is empty.
    pub fn values_reusable(&self) -> ValuesReusable<'_, T> {
        ValuesReusable {
            pool: &self.pool,
            values: &self.values,
        }
    }

    /// A generator that consumes values from the pool.
    ///
    /// Drawing from it removes a value from the pool and yields it by value.
    /// Once consumed, that value is never drawn again. Drawing rejects the
    /// current test case (as if by `assume(false)`) when the pool is empty.
    pub fn values_consumed(&mut self) -> ValuesConsumed<'_, T> {
        ValuesConsumed {
            pool: &self.pool,
            values: RefCell::new(&mut self.values),
        }
    }
}

/// A generator over references to the values in a [`Pool`].
///
/// Returned by [`Pool::values_reusable`]. Borrows the pool, so the references it
/// produces stay valid for as long as the generator is alive.
pub struct ValuesReusable<'a, T> {
    pool: &'a crate::ffi::PoolHandle,
    values: &'a HashMap<i64, T>,
}

impl<'a, T> Generator<&'a T> for ValuesReusable<'a, T> {
    fn do_draw(&self, tc: &TestCase) -> &'a T {
        tc.assume(!self.values.is_empty());
        let variable_id = pool_generate(tc, self.pool, false);
        self.values.get(&variable_id).unwrap()
    }
}

impl<'a, T: crate::PrettyPrintable> crate::generators::PrintableGenerator<&'a T>
    for ValuesReusable<'a, T>
{
    fn do_draw_and_print(&self, tc: &TestCase, printer: &mut crate::PrettyPrinter) -> &'a T {
        crate::generators::draw_and_print_value(self, tc, printer)
    }
}

/// A generator that consumes values from a [`Pool`], removing each value it
/// yields.
///
/// Returned by [`Pool::values_consumed`]. Borrows the pool mutably; the inner
/// [`RefCell`] is what lets it remove a value during a draw, which only has
/// shared access to the generator.
pub struct ValuesConsumed<'a, T> {
    pool: &'a crate::ffi::PoolHandle,
    values: RefCell<&'a mut HashMap<i64, T>>,
}

impl<T> Generator<T> for ValuesConsumed<'_, T> {
    fn do_draw(&self, tc: &TestCase) -> T {
        tc.assume(!self.values.borrow().is_empty());
        let variable_id = pool_generate(tc, self.pool, true);
        self.values.borrow_mut().remove(&variable_id).unwrap()
    }
}

impl<T: crate::PrettyPrintable> crate::generators::PrintableGenerator<T> for ValuesConsumed<'_, T> {
    fn do_draw_and_print(&self, tc: &TestCase, printer: &mut crate::PrettyPrinter) -> T {
        crate::generators::draw_and_print_value(self, tc, printer)
    }
}

/// Create a new value pool for stateful tests.
pub fn pool<T>(tc: &TestCase) -> Pool<T> {
    let pool = match tc.with_ctc(|ctc| ctc.new_pool()) {
        Ok(handle) => handle,
        Err(rc) => raise_for_rc(rc), // nocov
    };
    Pool {
        pool,
        tc: tc.clone(),
        values: HashMap::new(),
    }
}

/// A pool of previously generated values that [`run_concurrent`]'s worker
/// threads may share.
///
/// The concurrent counterpart of [`Pool`], designed for `&self` access from
/// rules: create one with [`concurrent_pool()`], store it in the model, and
/// call [`add`](ConcurrentPool::add) / draw from
/// [`values_reusable`](ConcurrentPool::values_reusable) /
/// [`values_consumed`](ConcurrentPool::values_consumed) from any worker.
/// `ConcurrentPool<T>` is `Sync` whenever `T: Send`, so a model holding one
/// satisfies [`run_concurrent`]'s `Sync` bound.
///
/// The engine performs each pool draw's empty check, selection, and
/// consumption atomically, so concurrent workers cannot double-consume a
/// value or race the emptiness check — exactly one worker receives each
/// consumed value. The accepted trade-off is that a shared pool couples the
/// workers' streams: which value a draw resolves to (and whether it rejects
/// as empty) depends on what other workers have added or consumed in the
/// meantime. In a nondeterministic run — which is what any run with
/// concurrency > 1 becomes — nothing downstream depends on that
/// independence anyway.
pub struct ConcurrentPool<T> {
    handle: PoolHandle,
    values: Mutex<HashMap<i64, T>>,
}

impl<T> ConcurrentPool<T> {
    /// Acquire the pool's map, recovering from poisoning: a panic can
    /// unwind while the guard is held (canonically a `Clone` impl panicking
    /// inside a reusable draw, which must clone under the lock), and
    /// letting that one panic turn every later pool operation on every
    /// worker into a `PoisonError` panic would bury the real failure under
    /// fake ones. Recovery is sound because no panic point sits between map
    /// mutations: a panicking clone reads the map without modifying it, and
    /// `add` inserts with an engine-issued id in a single operation, so the
    /// guarded map is consistent whenever the guard is droppable.
    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<i64, T>> {
        self.values.lock().unwrap_or_else(|e| e.into_inner())
    }

    /// Returns true if no values are in the pool.
    pub fn is_empty(&self) -> bool {
        self.lock().is_empty()
    }

    /// Number of values currently in the pool.
    pub fn len(&self) -> usize {
        self.lock().len()
    }

    /// Add a value to the pool. `tc` is the calling worker's test case.
    ///
    /// The pool's lock is held across the engine registration and the map
    /// insert, keeping the frontend map in lockstep with the engine's pool
    /// state: without this, a concurrent consumer could be handed a
    /// variable id whose value hasn't been inserted yet.
    pub fn add(&self, tc: &TestCase, v: T) {
        let mut values = self.lock();
        match tc.with_ctc(|ctc| ctc.pool_add(&self.handle)) {
            Ok(variable_id) => {
                let previous = values.insert(variable_id, v);
                hegel_internal_assert!(previous.is_none(), "unexpected variable id in map");
            }
            Err(rc) => {
                drop(values);
                raise_for_rc(rc)
            }
        }
    }

    /// A generator over copies of the values in the pool.
    ///
    /// Drawing from it yields a *clone* of a value in the pool, without
    /// removing it — another worker may consume the referenced value at any
    /// moment, so references cannot safely escape the pool's lock; store
    /// `Arc<T>` in the pool if cloning is expensive. Drawing rejects the
    /// current test case (as if by `assume(false)`) when the pool is empty.
    pub fn values_reusable(&self) -> ConcurrentValuesReusable<'_, T> {
        ConcurrentValuesReusable { pool: self }
    }

    /// A generator that consumes values from the pool.
    ///
    /// Drawing from it removes a value from the pool and yields it by
    /// value; the engine consumes the id atomically, so exactly one worker
    /// receives each value. Drawing rejects the current test case (as if by
    /// `assume(false)`) when the pool is empty.
    pub fn values_consumed(&self) -> ConcurrentValuesConsumed<'_, T> {
        ConcurrentValuesConsumed { pool: self }
    }
}

/// A generator over cloned values in a [`ConcurrentPool`]. Returned by
/// [`ConcurrentPool::values_reusable`].
pub struct ConcurrentValuesReusable<'a, T> {
    pool: &'a ConcurrentPool<T>,
}

impl<T: Clone> Generator<T> for ConcurrentValuesReusable<'_, T> {
    fn do_draw(&self, tc: &TestCase) -> T {
        let values = self.pool.lock();
        match tc.with_ctc(|ctc| ctc.pool_generate(&self.pool.handle, false)) {
            Ok(variable_id) => values.get(&variable_id).unwrap().clone(),
            Err(rc) => {
                drop(values);
                raise_for_rc(rc)
            }
        }
    }
}

impl<T: Clone + crate::PrettyPrintable> crate::generators::PrintableGenerator<T>
    for ConcurrentValuesReusable<'_, T>
{
    fn do_draw_and_print(&self, tc: &TestCase, printer: &mut crate::PrettyPrinter) -> T {
        crate::generators::draw_and_print_value(self, tc, printer)
    }
}

/// A generator that consumes values from a [`ConcurrentPool`], removing
/// each value it yields. Returned by [`ConcurrentPool::values_consumed`].
pub struct ConcurrentValuesConsumed<'a, T> {
    pool: &'a ConcurrentPool<T>,
}

impl<T> Generator<T> for ConcurrentValuesConsumed<'_, T> {
    fn do_draw(&self, tc: &TestCase) -> T {
        let mut values = self.pool.lock();
        match tc.with_ctc(|ctc| ctc.pool_generate(&self.pool.handle, true)) {
            Ok(variable_id) => values.remove(&variable_id).unwrap(),
            Err(rc) => {
                drop(values);
                raise_for_rc(rc)
            }
        }
    }
}

impl<T: crate::PrettyPrintable> crate::generators::PrintableGenerator<T>
    for ConcurrentValuesConsumed<'_, T>
{
    fn do_draw_and_print(&self, tc: &TestCase, printer: &mut crate::PrettyPrinter) -> T {
        crate::generators::draw_and_print_value(self, tc, printer)
    }
}

/// Create a new value pool for concurrent stateful tests. See
/// [`ConcurrentPool`].
pub fn concurrent_pool<T>(tc: &TestCase) -> ConcurrentPool<T> {
    let handle = match tc.with_ctc(|ctc| ctc.new_pool()) {
        Ok(handle) => handle,
        Err(rc) => raise_for_rc(rc),
    };
    ConcurrentPool {
        handle,
        values: Mutex::new(HashMap::new()),
    }
}

/// Trait for defining a stateful test.
///
/// Implement this to define the rules (actions) and invariants (assertions)
/// of your state machine. Use `#[hegel::state_machine]` for a more
/// ergonomic way to define state machines.
pub trait StateMachine {
    /// The rules (actions) that can be applied to this state machine.
    fn rules(&self) -> Vec<Rule<Self>>;
    /// Invariants: checked on the machine's initial and final state, and
    /// sampled after rules in between with probability
    /// `1 / stateful_step_count` each.
    fn invariants(&self) -> Vec<Rule<Self>>;
}

/// Run invariants at a join point. With a machine, each invariant runs only
/// when the engine's sampling draw says to (see
/// [`machine_should_check_invariant`]); with `None` — the guaranteed checks
/// of the machine's initial and final state — every invariant runs.
fn check_invariants<M: StateMachine>(
    m: &mut M,
    invariants: &[Rule<M>],
    tc: &TestCase,
    machine: Option<&StateMachineHandle>,
) {
    for (index, invariant) in invariants.iter().enumerate() {
        if let Some(machine) = machine {
            if !machine_should_check_invariant(tc, machine, index as i64) {
                continue;
            }
        }
        let inv_tc = tc.child(2); // nocov
        (invariant.apply)(m, inv_tc); // nocov
    }
}

/// Ask the engine whether the machine should run another round, and which
/// concurrency group is current for it.
///
/// The start of the test case counts as a join point, so both runners call
/// this *before* running any rules: the machine has no current group until
/// the first call draws one (the engine rejects a rule request made before
/// it), and the first call is also what draws the test case's round cap.
/// Since the cap is at least 1, the first call of a fresh machine always
/// starts a round; `None` only ever follows earlier rounds.
fn machine_next_group(tc: &TestCase, machine: &StateMachineHandle) -> Option<usize> {
    match tc.with_ctc(|ctc| ctc.state_machine_next_group(machine)) {
        Ok(group) => group.map(|g| g as usize),
        Err(rc) => raise_for_rc(rc),
    }
}

/// Ask the engine for the next rule worker `worker_index` should run this round;
/// `None` once the thread's round budget is exhausted (the join point).
fn machine_next_rule(
    tc: &TestCase,
    machine: &StateMachineHandle,
    worker_index: i64,
) -> Option<i64> {
    match tc.with_ctc(|ctc| ctc.state_machine_next_rule(machine, worker_index)) {
        Ok(next) => next,
        Err(rc) => raise_for_rc(rc),
    }
}

/// Report that the rule most recently handed to worker `worker_index` was
/// rejected (a violated assumption), so the engine does not count it toward
/// the test case's budget.
fn machine_rule_rejected(tc: &TestCase, machine: &StateMachineHandle, worker_index: i64) {
    if let Err(rc) = tc.with_ctc(|ctc| ctc.state_machine_rule_rejected(machine, worker_index)) {
        raise_for_rc(rc);
    }
}

/// Ask the engine whether invariant `invariant_index` should run at the
/// current join point — a recorded draw that is true with probability
/// `1 / stateful_step_count`, making an invariant's expected sampled runs
/// per test case one regardless of step count.
fn machine_should_check_invariant(
    tc: &TestCase,
    machine: &StateMachineHandle,
    invariant_index: i64,
) -> bool {
    match tc.with_ctc(|ctc| ctc.state_machine_should_check_invariant(machine, invariant_index)) {
        Ok(should_check) => should_check,
        Err(rc) => raise_for_rc(rc),
    }
}

/// Execute a stateful test by repeatedly applying random rules and checking invariants.
///
/// A sequential machine is the special case of the engine's concurrent
/// state-machine protocol with a single group and concurrency 1: the engine
/// hands out exactly one rule per round, so the join points where sampled
/// invariant checks may run fall after each rule. Invariants run in full on
/// the machine's initial and final state; in between, each invariant runs at
/// a join point only when the engine's sampling draw (probability
/// `1 / stateful_step_count`) says to, keeping an invariant's expected cost
/// per test case constant as the step count grows. One consequence of the
/// join-point timing: a sampled check can land after a rule that stopped on
/// a violated assumption (rules are expected to reject before mutating the
/// model, and nothing restores model state on rejection anyway).
pub fn run<M: StateMachine>(mut m: M, tc: TestCase) {
    let rules = m.rules();
    let rule_names: Vec<&str> = rules.iter().map(|r| r.name.as_str()).collect();
    let rule_groups = vec![0i64; rules.len()];
    let invariants = m.invariants();
    let invariant_names: Vec<&str> = invariants.iter().map(|r| r.name.as_str()).collect();
    let machine = match tc
        .with_ctc(|ctc| ctc.new_state_machine(&rule_names, &rule_groups, &invariant_names, 1, 1))
    {
        Ok((handle, _)) => handle,
        Err(rc) => raise_for_rc(rc),
    };

    tc.note("Initial invariant check.");
    check_invariants(&mut m, &invariants, &tc, None);

    let mut steps_attempted: i64 = 0;

    loop {
        tc.start_span(labels::STATEFUL_RULE);
        if machine_next_group(&tc, &machine).is_none() {
            tc.stop_span(false);
            break;
        }
        // The engine hands out one rule per round at concurrency 1, but
        // that is engine policy, not protocol: pull rules until the join
        // point.
        let mut round_rejected = false;
        while let Some(rule_index) = machine_next_rule(&tc, &machine, 0) {
            hegel_internal_assert!(
                (0..rules.len() as i64).contains(&rule_index),
                "state_machine_next_rule returned out-of-range rule index {rule_index}"
            );
            let rule = &rules[rule_index as usize];
            tc.note(&format!("Step {}: {}", steps_attempted + 1, rule.name));

            let rule_tc = tc.child(2);
            let thunk = || (rule.apply)(&mut m, rule_tc);
            let result = catch_unwind(AssertUnwindSafe(thunk));

            steps_attempted += 1;
            match result {
                Ok(()) => {}
                Err(e) if e.downcast_ref::<AssumeFailed>().is_some() => {
                    machine_rule_rejected(&tc, &machine, 0);
                    round_rejected = true;
                    tc.note("Rule stopped early due to violated assumption.");
                }
                // Everything else — including StopTest, so an out-of-data
                // case is reported as an overrun instead of returning
                // normally with a half-applied rule — unwinds through the
                // caller.
                Err(e) => {
                    tc.stop_span(false);
                    resume_unwind(e)
                }
            };
        }
        tc.stop_span(round_rejected);

        check_invariants(&mut m, &invariants, &tc, Some(&machine));
    }

    tc.note("Final invariant check.");
    check_invariants(&mut m, &invariants, &tc, None);
}

/// A rule of a [`ConcurrentStateMachine`]: an action worker threads may
/// apply to the shared model during concurrent stateful testing.
///
/// Unlike a sequential [`Rule`], the apply function takes `&M`: the model is
/// shared by reference across the worker threads, so any mutable model
/// state needs interior mutability. `group` names the concurrency group the
/// rule belongs to; rules in the same group may run concurrently with each
/// other, rules in different groups never overlap.
pub struct ConcurrentRule<M: ?Sized> {
    pub name: String,
    pub group: String,
    pub apply: fn(&M, TestCase),
}

impl<M> ConcurrentRule<M> {
    /// Create a new rule with a name, a concurrency group, and an apply
    /// function. Pass [`ANONYMOUS_GROUP`] as the group for a rule without a
    /// group annotation.
    pub fn new(name: &str, group: &str, apply: fn(&M, TestCase)) -> Self {
        ConcurrentRule {
            name: name.to_string(),
            group: group.to_string(),
            apply,
        }
    }
}

/// An invariant of a [`ConcurrentStateMachine`], checked on the main thread
/// at every join point — between rounds of concurrent rule execution, while
/// all worker threads are parked.
pub struct ConcurrentInvariant<M: ?Sized> {
    pub name: String,
    pub apply: fn(&M, TestCase),
}

impl<M> ConcurrentInvariant<M> {
    /// Create a new invariant with a name and an apply function.
    pub fn new(name: &str, apply: fn(&M, TestCase)) -> Self {
        ConcurrentInvariant {
            name: name.to_string(),
            apply,
        }
    }
}

/// Trait for defining a concurrent stateful test.
///
/// Implement this to define the rules (actions), their concurrency-group
/// assignments, and the invariants of a model whose rules may run
/// *concurrently* against the system under test. Use
/// `#[hegel::concurrent_state_machine]` for a more ergonomic way to define
/// concurrent state machines, and [`run_concurrent`] to run one.
///
/// # Groups
///
/// At any moment exactly one group is *current*, and only rules belonging
/// to the current group are handed out — so rules in the same group may run
/// concurrently with each other, rules in different groups never overlap,
/// and the current group changes only at the join points between rounds.
/// Groups cannot express asymmetric overlap ("put may overlap get but not
/// delete"); that expressiveness limit is deliberate.
///
/// A rule without a group annotation is assigned to a single shared
/// anonymous group ([`ANONYMOUS_GROUP`]), so an unannotated machine is
/// maximally concurrent: any rule may overlap with any other, and naming
/// groups is how overlap gets *restricted*. **In a machine that mixes
/// annotated and unannotated rules, the unannotated rules form their own
/// group and therefore never overlap with any named group's rules** — do
/// not read "no group" as "unconstrained"; it is exactly backwards there.
pub trait ConcurrentStateMachine {
    /// The rules (actions) that worker threads may apply to this state
    /// machine, each with its concurrency-group assignment.
    fn rules(&self) -> Vec<ConcurrentRule<Self>>;
    /// Invariants, run on the main thread while the worker threads are
    /// parked: checked on the machine's initial and final state, and
    /// sampled at the join points in between with probability
    /// `1 / stateful_step_count` each.
    fn invariants(&self) -> Vec<ConcurrentInvariant<Self>>;
}

/// Concurrent counterpart of [`check_invariants`]: sampled per invariant
/// when given a machine, a guaranteed full sweep with `None`.
fn check_concurrent_invariants<M: ConcurrentStateMachine + ?Sized>(
    m: &M,
    invariants: &[ConcurrentInvariant<M>],
    tc: &TestCase,
    machine: Option<&StateMachineHandle>,
) {
    for (index, invariant) in invariants.iter().enumerate() {
        if let Some(machine) = machine {
            if !machine_should_check_invariant(tc, machine, index as i64) {
                continue;
            }
        }
        let inv_tc = tc.child(2);
        (invariant.apply)(m, inv_tc);
    }
}

/// What a worker reports back to the main thread at the end of its round.
enum WorkerEvent {
    /// The rule stream was exhausted normally.
    RoundDone,
    /// An `AssumeFailed` raised by a control draw (`next_rule`) — the
    /// engine concluded the family invalid mid-round (e.g. span nesting
    /// past its limit), so the whole case is invalid. A rule body's own
    /// `AssumeFailed` never reaches the main thread as this event:
    /// [`run_worker_round`] intercepts it (a rejected rule just ends early)
    /// and moves on to the next rule.
    Invalid,
    /// The family's choice budget is exhausted (`StopTest`): the whole case
    /// is an overrun.
    Overrun,
    /// A run-aborting control payload (`InternalError`, `InvalidArgument`)
    /// or a `LoopDone`, ferried verbatim for the main thread to re-raise.
    ControlPayload(Box<dyn std::any::Any + Send>),
    /// A real panic, with the worker-side capture to re-install on the main
    /// thread.
    Panicked {
        payload: Box<dyn std::any::Any + Send>,
        info: Option<PanicInfo>,
    },
    /// Synthesized by the main thread when a worker exited without
    /// reporting: never constructed by workers.
    Died,
}

/// Classify an unwind caught inside a worker as the event it ferries to
/// the main thread.
fn classify_worker_unwind(e: Box<dyn std::any::Any + Send>) -> WorkerEvent {
    if e.downcast_ref::<AssumeFailed>().is_some() {
        return WorkerEvent::Invalid;
    }
    if e.downcast_ref::<StopTest>().is_some() {
        return WorkerEvent::Overrun;
    }
    if e.downcast_ref::<InvalidArgument>().is_some()
        || e.downcast_ref::<InternalError>().is_some()
        || e.downcast_ref::<LoopDone>().is_some()
    {
        return WorkerEvent::ControlPayload(e);
    }
    WorkerEvent::Panicked {
        payload: e,
        info: run_lifecycle::take_panic_info(),
    }
}

/// One worker's round: pull rules for `worker` until the engine signals the
/// join point or something terminal happens. Every unwind source — the
/// `next_rule` control draws included — runs under `catch_unwind`, so no
/// unwind ever escapes the worker thread: a worker that died without
/// reporting would leave the main thread parked forever waiting for its
/// event.
fn run_worker_round<M: ConcurrentStateMachine + ?Sized>(
    worker: usize,
    tc: &TestCase,
    m: &M,
    rules: &[ConcurrentRule<M>],
    machine: &StateMachineHandle,
) -> WorkerEvent {
    loop {
        let next = catch_unwind(AssertUnwindSafe(|| {
            let next = machine_next_rule(tc, machine, worker as i64);
            if let Some(rule_index) = next {
                hegel_internal_assert!(
                    (0..rules.len() as i64).contains(&rule_index),
                    "state_machine_next_rule returned out-of-range rule index {rule_index}"
                );
            }
            next
        }));
        let rule_index = match next {
            Ok(Some(rule_index)) => rule_index,
            Ok(None) => return WorkerEvent::RoundDone,
            Err(e) => return classify_worker_unwind(e),
        };

        let rule = &rules[rule_index as usize];
        tc.note(&format!("Rule: {}", rule.name));
        let rule_tc = tc.child(2);
        let result = catch_unwind(AssertUnwindSafe(|| (rule.apply)(m, rule_tc)));
        match result {
            Ok(()) => {}
            Err(e) => match classify_worker_unwind(e) {
                WorkerEvent::Invalid => {
                    let rejected = catch_unwind(AssertUnwindSafe(|| {
                        machine_rule_rejected(tc, machine, worker as i64);
                    }));
                    if let Err(e) = rejected {
                        return classify_worker_unwind(e);
                    }
                    tc.note("Rule stopped early due to violated assumption.");
                }
                event => return event,
            },
        }
    }
}

/// A worker thread's whole-test-case loop: enter the test context (the
/// panic hook captures nothing on a thread outside it, and internal errors
/// must raise catchably), mirror the main thread's backtrace-capture
/// setting, then run one round per `rounds` message until the channel
/// closes.
///
/// Each `rounds` message carries a fresh [`TestCase`] clone made by the
/// main thread right after it noted the round's header, so the clone's
/// print region — where everything this worker draws or notes during the
/// round lands — is anchored under the header that labels the round.
/// Cloning once per worker up front would instead pool a worker's output
/// across all rounds at its single anchor, detaching the headers from the
/// rounds they label.
///
/// The closed channel is what terminates workers: the main thread holds
/// the senders as locals of its `thread::scope` body, so *any* exit from
/// that body — the normal end of the test case or an unwind from a join
/// point (a panicking invariant, an invariant's `AssumeFailed`, a draw
/// that exhausts the budget) — drops them, wakes the parked workers, and
/// lets the scope's implicit join complete instead of hanging.
fn worker_loop<M: ConcurrentStateMachine + ?Sized>(
    worker: usize,
    m: &M,
    rules: &[ConcurrentRule<M>],
    machine: &StateMachineHandle,
    capture_backtraces: bool,
    rounds: mpsc::Receiver<TestCase>,
    events: mpsc::Sender<WorkerEvent>,
) {
    WORKER_INDEX.with(|cell| cell.set(Some(worker)));
    run_lifecycle::set_backtrace_capture(capture_backtraces);
    with_test_context(|| {
        while let Ok(tc) = rounds.recv() {
            let event = run_worker_round(worker, &tc, m, rules, machine);
            if events.send(event).is_err() {
                break;
            }
        }
    });
}

/// Execute a concurrent stateful test. Execution proceeds in *rounds*. For
/// each round, the engine picks a random concurrency group; every worker
/// thread then runs a short (possibly empty) random sequence of rules from
/// that group — and only that group — concurrently with the other workers. Rules in the same
/// group may overlap each other, and rules in different groups never
/// overlap. Once every worker has finished its rules for the round, the
/// main thread runs the invariants the engine's sampling draws select
/// (probability `1 / stateful_step_count` each); every invariant runs in
/// full on the machine's initial and final state.
///
/// The number of worker threads is drawn per test case, when the state
/// machine is created, in `[min_concurrency, max_concurrency]` and weighted
/// toward `max_concurrency` (concurrency bugs need concurrency); pass
/// `min_concurrency == max_concurrency` for a fixed level.
///
/// # Nondeterminism
///
/// Concurrency bugs are nondeterministic — thread scheduling is outside
/// Hegel's control — so calling `run_concurrent` with `max_concurrency > 1`
/// makes the whole run nondeterministic. Failures are reported from the
/// discovering execution, with no replay, shrinking, database persistence,
/// or reproduce blob, with at most one failure per run.
///
/// # Abandoned rules and lock poisoning
///
/// A rule abandoned mid-execution by a rejected assumption or a failed draw
/// can poison any lock `std::sync::Mutex` it was holding. To avoid this,
/// perform all draws upfront within each rule.
///
/// # Example
///
/// ```no_run
/// use std::sync::Mutex;
/// use hegel::TestCase;
/// use hegel::generators as gs;
///
/// struct CounterTest {
///     counter: Mutex<i64>,
/// }
///
/// #[hegel::concurrent_state_machine]
/// impl CounterTest {
///     #[rule(group = "write")]
///     fn increment(&self, _: TestCase) {
///         *self.counter.lock().unwrap_or_else(|e| e.into_inner()) += 1;
///     }
///
///     #[rule(group = "read")]
///     fn read(&self, _: TestCase) {
///         let _ = *self.counter.lock().unwrap_or_else(|e| e.into_inner());
///     }
///
///     #[invariant]
///     fn non_negative(&self, _: TestCase) {
///         assert!(*self.counter.lock().unwrap_or_else(|e| e.into_inner()) >= 0);
///     }
/// }
///
/// #[hegel::test]
/// fn test_counter(tc: TestCase) {
///     let m = CounterTest { counter: Mutex::new(0) };
///     hegel::stateful::run_concurrent(m, tc, 1, 3);
/// }
/// ```
pub fn run_concurrent<M: ConcurrentStateMachine + Sync>(
    m: M,
    tc: TestCase,
    min_concurrency: i64,
    max_concurrency: i64,
) {
    let rules = m.rules();
    let invariants = m.invariants();
    let rule_names: Vec<&str> = rules.iter().map(|r| r.name.as_str()).collect();
    let invariant_names: Vec<&str> = invariants.iter().map(|r| r.name.as_str()).collect();
    let mut group_names: Vec<&str> = Vec::new();
    let mut rule_groups: Vec<i64> = Vec::with_capacity(rules.len());
    for rule in &rules {
        let index = group_names
            .iter()
            .position(|name| *name == rule.group)
            .unwrap_or_else(|| {
                group_names.push(rule.group.as_str());
                group_names.len() - 1
            });
        rule_groups.push(index as i64);
    }

    let (machine, concurrency) = match tc.with_ctc(|ctc| {
        ctc.new_state_machine(
            &rule_names,
            &rule_groups,
            &invariant_names,
            min_concurrency,
            max_concurrency,
        )
    }) {
        Ok(created) => created,
        Err(rc) => raise_for_rc(rc),
    };
    tc.note(&format!("Concurrency level: {concurrency}"));

    tc.note("Initial invariant check.");
    check_concurrent_invariants(&m, &invariants, &tc, None);

    let capture_backtraces = run_lifecycle::backtrace_capture_enabled();
    let concurrency = concurrency as usize;
    let m = &m;
    let rules = &rules;
    let machine = &machine;

    std::thread::scope(|scope| {
        let mut round_txs: Vec<mpsc::Sender<TestCase>> = Vec::with_capacity(concurrency);
        let mut event_rxs: Vec<mpsc::Receiver<WorkerEvent>> = Vec::with_capacity(concurrency);
        for worker in 0..concurrency {
            let (round_tx, round_rx) = mpsc::channel();
            let (event_tx, event_rx) = mpsc::channel();
            round_txs.push(round_tx);
            event_rxs.push(event_rx);
            scope.spawn(move || {
                worker_loop(
                    worker,
                    m,
                    rules,
                    machine,
                    capture_backtraces,
                    round_rx,
                    event_tx,
                );
            });
        }

        let mut round = 0u64;
        while let Some(group) = machine_next_group(&tc, machine) {
            hegel_internal_assert!(
                group < group_names.len(),
                "state_machine_next_group returned unknown group id {group}"
            );
            round += 1;
            tc.note(&format!(
                "---------------- Round {round}: group {:?} ----------------",
                group_names[group]
            ));

            for tx in &round_txs {
                let _ = tx.send(tc.clone());
            }
            let events: Vec<WorkerEvent> = event_rxs
                .iter()
                .map(|rx| rx.recv().unwrap_or(WorkerEvent::Died))
                .collect();

            resolve_round(events, &tc);

            check_concurrent_invariants(m, &invariants, &tc, Some(machine));
        }

        tc.note("Final invariant check.");
        check_concurrent_invariants(m, &invariants, &tc, None);
    });
}

/// Classify one round's worker events and, for a terminal round, re-raise
/// on the main thread. Precedence: **control payloads win over overrun, and
/// overrun wins over panic** — a control payload signals a framework or
/// usage bug that must not be masked, and a panic that co-occurs with an
/// engine-side family conclusion (overrun, invalid) is not trustworthy: the
/// concluding draw abandoned a rule mid-execution, and that abandonment's
/// side effects (canonically a poisoned model lock) can induce panics in
/// other workers that no real schedule of the user's rules could produce.
/// Misclassification costs are asymmetric, too: a dropped genuine panic
/// merely discards this case and resurfaces in a later one, while a
/// reported fake panic halts the run with a false bug. Dropped panics are
/// noted into the case's output buffer, so they stay visible whenever that
/// buffer is shown. Among several panics, the lowest worker index wins.
fn resolve_round(events: Vec<WorkerEvent>, tc: &TestCase) {
    struct WorkerPanic {
        worker: usize,
        payload: Box<dyn std::any::Any + Send>,
        info: Option<PanicInfo>,
    }

    let mut control: Option<Box<dyn std::any::Any + Send>> = None;
    let mut saw_overrun = false;
    let mut saw_invalid = false;
    let mut panics: Vec<WorkerPanic> = Vec::new();
    for (worker, event) in events.into_iter().enumerate() {
        match event {
            WorkerEvent::RoundDone => {}
            WorkerEvent::Invalid => saw_invalid = true,
            WorkerEvent::Overrun => saw_overrun = true,
            WorkerEvent::ControlPayload(payload) => {
                if control.is_none() {
                    control = Some(payload);
                }
            }
            WorkerEvent::Panicked { payload, info } => panics.push(WorkerPanic {
                worker,
                payload,
                info,
            }),
            WorkerEvent::Died => {
                if control.is_none() {
                    control = Some(Box::new(InternalError(format!(
                        "Internal error in hegel: concurrent stateful worker {worker} exited \
                         without reporting an outcome. This is a bug in hegel itself; please \
                         report it at https://github.com/hegeldev/hegel-rust/issues"
                    ))));
                }
            }
        }
    }

    let note_dropped = |dropped: &[WorkerPanic]| {
        for p in dropped {
            let location = p
                .info
                .as_ref()
                .map_or("<unknown>", |(_, _, location, _)| location.as_str());
            tc.note(&format!(
                "Dropped concurrent panic from worker {} at {}: {}",
                p.worker,
                location,
                run_lifecycle::panic_message(&p.payload)
            ));
        }
    };

    if let Some(payload) = control {
        resume_unwind(payload);
    }
    if saw_overrun || saw_invalid {
        note_dropped(&panics);
        if saw_overrun {
            raise_control(StopTest);
        }
        raise_control(AssumeFailed);
    }
    if !panics.is_empty() {
        note_dropped(&panics[1..]);
        let winner = panics.remove(0);
        if let Some(info) = winner.info {
            run_lifecycle::install_panic_info(info);
        }
        resume_unwind(winner.payload);
    }
}

#[cfg(test)]
#[path = "../tests/embedded/stateful_tests.rs"]
mod tests;