alloc_chaos/

lib.rs

#![deny(unsafe_op_in_unsafe_fn)]
#![warn(missing_docs, rust_2018_idioms, unreachable_pub)]

//! Deterministic allocation-failure testing for Rust.
//!
//! `alloc-chaos` wraps the process global allocator and can run a test closure
//! repeatedly while failing allocation attempt `0`, then `1`, then `2`, and so
//! on. This is useful for libraries that intentionally use fallible allocation
//! APIs and want tests for their OOM paths.
//!
//! # Example
//!
//! ```no_run
//! #[global_allocator]
//! static GLOBAL: alloc_chaos::ChaosAllocator = alloc_chaos::ChaosAllocator::system();
//!
//! #[derive(Debug, Clone, Copy, PartialEq, Eq)]
//! struct OutOfMemory;
//!
//! fn build_buffer(size: usize) -> Result<Vec<u8>, OutOfMemory> {
//!     let mut bytes = Vec::new();
//!     bytes.try_reserve_exact(size).map_err(|_| OutOfMemory)?;
//!     bytes.resize(size, 0);
//!     Ok(bytes)
//! }
//!
//! #[test]
//! fn buffer_builder_handles_oom() {
//!     alloc_chaos::check(|| {
//!         match build_buffer(1024) {
//!             Ok(bytes) => assert_eq!(bytes.len(), 1024),
//!             Err(OutOfMemory) => {}
//!         }
//!     })
//!     .assert_success();
//! }
//! ```
//!
//! # Limits
//!
//! This crate verifies a concrete execution path. It does not prove that all
//! possible inputs or all possible control-flow branches are OOM-safe.
//!
//! Only one check can be active in a process, but allocation counting and
//! failure injection are limited to the thread executing the checked closure.
//! Allocations performed by other threads are ignored.
//!
//! The closure is executed many times. Keep all state needed by the code under
//! test inside the closure, or otherwise make the closure deterministic across
//! repeated runs.
//!
//! The in-process checker cannot contain process aborts. Code that uses
//! infallible allocation APIs may abort when an allocation fails instead of
//! returning a recoverable error.

use std::alloc::{GlobalAlloc, Layout, System};
use std::any::Any;
use std::cell::Cell;
use std::error::Error;
use std::fmt;
use std::ops::Range;
use std::panic::{self, AssertUnwindSafe};
use std::ptr;
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering};

const MODE_DISABLED: u8 = 0;
const MODE_COUNTING: u8 = 1;
const MODE_FAILING: u8 = 2;

const NO_TARGET: usize = usize::MAX;
const NO_INJECTED_INDEX: usize = usize::MAX;
const NO_REALLOC_NEW_SIZE: usize = usize::MAX;

const ALLOC_OP_NONE: u8 = 0;
const ALLOC_OP_ALLOC: u8 = 1;
const ALLOC_OP_ALLOC_ZEROED: u8 = 2;
const ALLOC_OP_REALLOC: u8 = 3;

static MODE: AtomicU8 = AtomicU8::new(MODE_DISABLED);
static TARGET: AtomicUsize = AtomicUsize::new(NO_TARGET);
static SEEN: AtomicUsize = AtomicUsize::new(0);
static INJECTED: AtomicBool = AtomicBool::new(false);
static INJECTED_INDEX: AtomicUsize = AtomicUsize::new(NO_INJECTED_INDEX);
static INJECTED_OP: AtomicU8 = AtomicU8::new(ALLOC_OP_NONE);
static INJECTED_SIZE: AtomicUsize = AtomicUsize::new(0);
static INJECTED_ALIGN: AtomicUsize = AtomicUsize::new(0);
static INJECTED_NEW_SIZE: AtomicUsize = AtomicUsize::new(NO_REALLOC_NEW_SIZE);
static CHECK_ACTIVE: AtomicBool = AtomicBool::new(false);
static ALLOCATOR_WAS_USED: AtomicBool = AtomicBool::new(false);

std::thread_local! {
    static TRACK_ALLOCATIONS_ON_THREAD: Cell<bool> = const { Cell::new(false) };
}

/// A global allocator wrapper that can deterministically fail selected
/// allocation attempts during an [`alloc_chaos`](crate) check.
///
/// Install it once in a test binary:
///
/// ```no_run
/// #[global_allocator]
/// static GLOBAL: alloc_chaos::ChaosAllocator = alloc_chaos::ChaosAllocator::system();
/// ```
///
/// The wrapper delegates to its inner allocator unless a check is active and
/// the current allocation attempt is the selected failure target.
pub struct ChaosAllocator<A = System> {
    inner: A,
}

impl ChaosAllocator<System> {
    /// Creates a wrapper around [`std::alloc::System`].
    #[must_use]
    pub const fn system() -> Self {
        Self { inner: System }
    }
}

impl<A> ChaosAllocator<A> {
    /// Creates a wrapper around a custom global allocator.
    #[must_use]
    pub const fn new(inner: A) -> Self {
        Self { inner }
    }
}

impl<A> fmt::Debug for ChaosAllocator<A> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ChaosAllocator").finish_non_exhaustive()
    }
}

// SAFETY: The wrapper preserves the `GlobalAlloc` contract by delegating all
// successful operations to the inner allocator. For injected failures it returns
// null from `alloc`, `alloc_zeroed`, or `realloc`. Returning null is the
// allocator-level failure signal. For `realloc`, returning null leaves the
// original allocation untouched, which is the required behavior.
unsafe impl<A> GlobalAlloc for ChaosAllocator<A>
where
    A: GlobalAlloc,
{
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        if should_inject_failure(AllocOp::Alloc, layout, None) {
            ptr::null_mut()
        } else {
            // SAFETY: Delegates the exact layout received from the caller to the
            // wrapped allocator.
            unsafe { self.inner.alloc(layout) }
        }
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        if should_inject_failure(AllocOp::AllocZeroed, layout, None) {
            ptr::null_mut()
        } else {
            // SAFETY: Delegates the exact layout received from the caller to the
            // wrapped allocator.
            unsafe { self.inner.alloc_zeroed(layout) }
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        // SAFETY: Delegates the pointer and layout received from the caller to
        // the wrapped allocator. Deallocation is never failed or counted.
        unsafe { self.inner.dealloc(ptr, layout) }
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        if should_inject_failure(AllocOp::Realloc, layout, Some(new_size)) {
            ptr::null_mut()
        } else {
            // SAFETY: Delegates the exact pointer, old layout, and requested new
            // size received from the caller to the wrapped allocator.
            unsafe { self.inner.realloc(ptr, layout, new_size) }
        }
    }
}

/// Runs `f` once to count allocations and then once per observed allocation,
/// failing one allocation attempt per run.
///
/// This is equivalent to `Check::new().run(f)`.
///
/// # Panics
///
/// Panics if another check is already active in the process. Use
/// [`try_check`] to handle that case explicitly.
pub fn check<F>(f: F) -> Report
where
    F: Fn(),
{
    Check::new().run(f)
}

/// Fallible variant of [`check`].
pub fn try_check<F>(f: F) -> Result<Report, AlreadyRunning>
where
    F: Fn(),
{
    Check::new().try_run(f)
}

/// Configuration for an allocation-failure check.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[must_use]
pub struct Check {
    max_failures: Option<usize>,
    stop_on_failure: bool,
    target_start: usize,
    target_end: Option<usize>,
    stability_runs: usize,
}

impl Check {
    /// Creates a check with no failure limit, no early stop, and one baseline
    /// run.
    pub const fn new() -> Self {
        Self {
            max_failures: None,
            stop_on_failure: false,
            target_start: 0,
            target_end: None,
            stability_runs: 1,
        }
    }

    /// Limits the number of allocation attempts that will be failed.
    ///
    /// This is useful when a baseline run performs many allocations and a full
    /// exhaustive check would be too expensive. A limited report is considered
    /// truncated, so [`Report::assert_success`] will fail unless the limit still
    /// covers every observed baseline allocation.
    pub const fn max_failures(mut self, max_failures: usize) -> Self {
        self.max_failures = Some(max_failures);
        self
    }

    /// Removes a previously configured failure limit.
    pub const fn unlimited_failures(mut self) -> Self {
        self.max_failures = None;
        self
    }

    /// Tests exactly one zero-based allocation attempt.
    ///
    /// This is primarily a reproduction and debugging aid after a full check
    /// has identified an interesting allocation number. A single-target report
    /// is considered truncated when the baseline contains any other allocation
    /// attempts, so [`Report::assert_success`] remains exhaustive by default.
    ///
    /// # Panics
    ///
    /// Panics if `target` is [`usize::MAX`].
    pub fn only_failure(mut self, target: usize) -> Self {
        self.target_start = target;
        self.target_end = Some(
            target
                .checked_add(1)
                .expect("alloc-chaos failure target must be less than usize::MAX"),
        );
        self
    }

    /// Tests a zero-based half-open range of allocation attempts.
    ///
    /// For example, `failure_range(30..40)` tests allocation attempts `30`
    /// through `39`. Ranges that do not cover every observed baseline
    /// allocation produce truncated reports.
    ///
    /// # Panics
    ///
    /// Panics if `range.start > range.end`.
    pub fn failure_range(mut self, range: Range<usize>) -> Self {
        assert!(
            range.start <= range.end,
            "alloc-chaos failure range start must be less than or equal to range end"
        );

        self.target_start = range.start;
        self.target_end = Some(range.end);
        self
    }

    /// Restores the default target selection, which tests every observed
    /// baseline allocation attempt.
    pub const fn all_failures(mut self) -> Self {
        self.target_start = 0;
        self.target_end = None;
        self
    }

    /// Configures how many baseline counting runs are used to check allocation
    /// sequence stability before injection begins.
    ///
    /// The default is `1`, which performs no extra stability comparison. Values
    /// greater than `1` rerun the closure in counting mode and require every
    /// baseline run to complete with the same allocation count. Passing `0` is
    /// treated as `1`.
    pub const fn stability_runs(mut self, runs: usize) -> Self {
        self.stability_runs = if runs == 0 { 1 } else { runs };
        self
    }

    /// Stops after the first failed iteration.
    ///
    /// A failed iteration is one that panics or does not reach the selected
    /// allocation attempt. Early-stop reports are considered truncated unless
    /// the stopped attempt was the final observed allocation.
    pub const fn stop_on_failure(mut self, enabled: bool) -> Self {
        self.stop_on_failure = enabled;
        self
    }

    /// Runs this check.
    ///
    /// # Panics
    ///
    /// Panics if another check is already active in the process. Use
    /// [`Check::try_run`] to handle that case explicitly.
    pub fn run<F>(self, f: F) -> Report
    where
        F: Fn(),
    {
        self.try_run(f)
            .expect("alloc-chaos check is already active in this process")
    }

    /// Fallible variant of [`Check::run`].
    pub fn try_run<F>(self, f: F) -> Result<Report, AlreadyRunning>
    where
        F: Fn(),
    {
        let _active = ActiveCheck::enter()?;

        let allocator_installed = probe_allocator_installed();
        let baseline = run_counting(&f);
        let baseline_allocations = baseline.observed_allocations;

        let mut stability_baselines = Vec::new();
        let mut baseline_stable = true;

        if baseline.outcome.is_completed() {
            stability_baselines.reserve_exact(self.stability_runs.saturating_sub(1));

            for _ in 1..self.stability_runs {
                let candidate = run_counting(&f);
                if !baseline_matches(&baseline, &candidate) {
                    baseline_stable = false;
                }
                stability_baselines.push(candidate);
            }
        }

        let mut attempts = Vec::new();
        let mut truncated = baseline_allocations > 0;

        if baseline.outcome.is_completed() && baseline_stable {
            let plan = FailurePlan::for_check(self, baseline_allocations);
            truncated = !plan.is_exhaustive_for(baseline_allocations);
            attempts.reserve_exact(plan.len());

            for target in plan.targets() {
                let attempt = run_failing(target, &f);
                let stop = self.stop_on_failure && !attempt.is_success();
                attempts.push(attempt);

                if stop {
                    truncated = !plan.is_exhaustive_after_stop(target, baseline_allocations);
                    break;
                }
            }
        }

        Ok(Report {
            baseline,
            stability_baselines,
            baseline_stable,
            attempts,
            truncated,
            allocator_installed,
        })
    }
}

impl Default for Check {
    fn default() -> Self {
        Self::new()
    }
}

/// Error returned when a check is requested while another check is already
/// active in this process.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AlreadyRunning;

impl fmt::Display for AlreadyRunning {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("an alloc-chaos check is already active in this process")
    }
}

impl Error for AlreadyRunning {}

/// Allocator operation that was selected for injected failure.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[must_use]
pub enum AllocOp {
    /// A call to [`GlobalAlloc::alloc`].
    Alloc,

    /// A call to [`GlobalAlloc::alloc_zeroed`].
    AllocZeroed,

    /// A call to [`GlobalAlloc::realloc`].
    Realloc,
}

impl AllocOp {
    const fn as_u8(self) -> u8 {
        match self {
            Self::Alloc => ALLOC_OP_ALLOC,
            Self::AllocZeroed => ALLOC_OP_ALLOC_ZEROED,
            Self::Realloc => ALLOC_OP_REALLOC,
        }
    }

    const fn from_u8(value: u8) -> Option<Self> {
        match value {
            ALLOC_OP_ALLOC => Some(Self::Alloc),
            ALLOC_OP_ALLOC_ZEROED => Some(Self::AllocZeroed),
            ALLOC_OP_REALLOC => Some(Self::Realloc),
            _ => None,
        }
    }
}

impl fmt::Display for AllocOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Alloc => f.write_str("alloc"),
            Self::AllocZeroed => f.write_str("alloc_zeroed"),
            Self::Realloc => f.write_str("realloc"),
        }
    }
}

/// Metadata for the allocation attempt that was selected for injected failure.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[must_use]
pub struct Allocation {
    index: usize,
    operation: AllocOp,
    size: usize,
    align: usize,
    new_size: Option<usize>,
}

impl Allocation {
    /// Zero-based allocation attempt index.
    #[must_use]
    pub fn index(&self) -> usize {
        self.index
    }

    /// Allocator operation used by this attempt.
    pub fn operation(&self) -> AllocOp {
        self.operation
    }

    /// Requested layout size in bytes.
    #[must_use]
    pub fn size(&self) -> usize {
        self.size
    }

    /// Requested layout alignment in bytes.
    #[must_use]
    pub fn align(&self) -> usize {
        self.align
    }

    /// Requested new size for [`AllocOp::Realloc`] attempts.
    #[must_use]
    pub fn new_size(&self) -> Option<usize> {
        self.new_size
    }
}

impl fmt::Display for Allocation {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} size={} align={}", self.operation, self.size, self.align)?;

        if let Some(new_size) = self.new_size {
            write!(f, " new_size={new_size}")?;
        }

        Ok(())
    }
}

/// Result of a full allocation-failure check.
#[derive(Debug, Clone, PartialEq, Eq)]
#[must_use]
pub struct Report {
    baseline: Baseline,
    stability_baselines: Vec<Baseline>,
    baseline_stable: bool,
    attempts: Vec<Attempt>,
    truncated: bool,
    allocator_installed: bool,
}

impl Report {
    /// Returns the primary baseline run.
    pub fn baseline(&self) -> &Baseline {
        &self.baseline
    }

    /// Returns extra baseline runs used to validate allocation-count stability.
    pub fn stability_baselines(&self) -> &[Baseline] {
        &self.stability_baselines
    }

    /// Returns whether every baseline run completed with the same allocation
    /// count as the primary baseline run.
    #[must_use]
    pub fn baseline_is_stable(&self) -> bool {
        self.baseline_stable
    }

    /// Returns the number of allocation attempts observed during the primary
    /// baseline run.
    #[must_use]
    pub fn baseline_allocations(&self) -> usize {
        self.baseline.observed_allocations
    }

    /// Returns all injected-failure runs.
    pub fn attempts(&self) -> &[Attempt] {
        &self.attempts
    }

    /// Returns all injected-failure runs that did not complete successfully.
    pub fn failed_attempts(&self) -> impl Iterator<Item = &Attempt> {
        self.attempts.iter().filter(|attempt| !attempt.is_success())
    }

    /// Returns the first injected-failure run that did not complete
    /// successfully.
    #[must_use]
    pub fn first_failure(&self) -> Option<&Attempt> {
        self.attempts.iter().find(|attempt| !attempt.is_success())
    }

    /// Returns the number of selected failure targets that were executed.
    #[must_use]
    pub fn tested_failures(&self) -> usize {
        self.attempts.len()
    }

    /// Returns the number of executed runs that actually reached and failed the
    /// selected allocation attempt.
    #[must_use]
    pub fn injected_failures(&self) -> usize {
        self.attempts
            .iter()
            .filter(|attempt| attempt.injected())
            .count()
    }

    /// Returns the number of observed baseline allocation attempts that did not
    /// have a corresponding injected failure.
    #[must_use]
    pub fn untested_failures(&self) -> usize {
        self.baseline_allocations()
            .saturating_sub(self.injected_failures())
    }

    /// Returns `true` if not every observed baseline allocation was tested.
    ///
    /// This can happen when the baseline run panics, when the baseline is
    /// unstable, when [`Check::max_failures`] limits the run count, when
    /// [`Check::only_failure`] or [`Check::failure_range`] selects a subset, or
    /// when [`Check::stop_on_failure`] stops the check early.
    #[must_use]
    pub fn is_truncated(&self) -> bool {
        self.truncated
    }

    /// Returns whether the [`ChaosAllocator`] was observed during an explicit
    /// allocator-installation probe at the start of the check.
    ///
    /// `false` usually means the global allocator wrapper was not installed.
    #[must_use]
    pub fn allocator_installed(&self) -> bool {
        self.allocator_installed
    }

    /// Returns `true` only for an exhaustive, valid, fully successful check.
    ///
    /// Success requires all of the following:
    ///
    /// - the global allocator wrapper was installed;
    /// - the primary baseline completed;
    /// - all configured stability baseline runs matched the primary baseline;
    /// - the report is not truncated;
    /// - every observed baseline allocation was tested; and
    /// - every injected run reached its target allocation and completed without
    ///   panic.
    #[must_use]
    pub fn is_success(&self) -> bool {
        self.allocator_installed
            && self.baseline.outcome.is_completed()
            && self.baseline_stable
            && !self.truncated
            && self.tested_failures() == self.baseline_allocations()
            && self.attempts.iter().all(Attempt::is_success)
    }

    /// Panics with a human-readable report if [`Report::is_success`] is `false`.
    pub fn assert_success(&self) {
        assert!(self.is_success(), "{self}");
    }
}

impl fmt::Display for Report {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "alloc-chaos report")?;
        writeln!(
            f,
            "  baseline: {}, {} allocation attempt(s)",
            self.baseline.outcome, self.baseline.observed_allocations
        )?;

        if !self.stability_baselines.is_empty() {
            writeln!(
                f,
                "  stability: {} across {} baseline run(s)",
                if self.baseline_stable {
                    "stable"
                } else {
                    "unstable"
                },
                self.stability_baselines.len() + 1
            )?;

            if !self.baseline_stable {
                for (index, baseline) in self.stability_baselines.iter().enumerate() {
                    writeln!(
                        f,
                        "    baseline #{}: {}, {} allocation attempt(s)",
                        index + 2,
                        baseline.outcome,
                        baseline.observed_allocations
                    )?;
                }
            }
        }

        writeln!(
            f,
            "  tested: {}/{} allocation failure target(s){}",
            self.tested_failures(),
            self.baseline.observed_allocations,
            if self.truncated { " (truncated)" } else { "" }
        )?;

        let injected_failures = self.injected_failures();
        if injected_failures != self.tested_failures() {
            writeln!(
                f,
                "  injected: {injected_failures}/{} selected target(s)",
                self.tested_failures()
            )?;
        }

        if self.is_truncated() || injected_failures != self.baseline_allocations() {
            writeln!(f, "  untested: {} allocation failure(s)", self.untested_failures())?;
        }
        writeln!(f, "  allocator installed: {}", self.allocator_installed)?;

        if self.is_success() {
            writeln!(f, "  status: exhaustive success")?;
        } else if !self.allocator_installed {
            writeln!(f, "  status: invalid check: allocator wrapper not observed")?;
        } else if !self.baseline.outcome.is_completed() {
            writeln!(f, "  status: failure")?;
            writeln!(f, "  baseline failed: {}", self.baseline.outcome)?;
        } else if !self.baseline_stable {
            writeln!(f, "  status: invalid check: unstable baseline")?;
        } else if self.first_failure().is_some() {
            writeln!(f, "  status: failure")?;

            for attempt in self.failed_attempts() {
                write!(
                    f,
                    "  allocation #{}: {}, injected={}, observed={} allocation attempt(s)",
                    attempt.target_allocation,
                    attempt.outcome,
                    attempt.injected,
                    attempt.observed_allocations
                )?;

                if let Some(allocation) = attempt.injected_allocation {
                    write!(f, ", {allocation}")?;
                }

                writeln!(f)?;
            }
        } else if self.truncated {
            writeln!(f, "  status: partial success")?;
        } else {
            writeln!(f, "  status: failure")?;
        }

        Ok(())
    }
}

/// Baseline run used to count allocation attempts before injecting failures.
#[derive(Debug, Clone, PartialEq, Eq)]
#[must_use]
pub struct Baseline {
    observed_allocations: usize,
    outcome: Outcome,
}

impl Baseline {
    /// Number of allocation attempts observed in the baseline run.
    #[must_use]
    pub fn observed_allocations(&self) -> usize {
        self.observed_allocations
    }

    /// Outcome of the baseline run.
    pub fn outcome(&self) -> &Outcome {
        &self.outcome
    }
}

/// One run with one selected allocation attempt failed.
#[derive(Debug, Clone, PartialEq, Eq)]
#[must_use]
pub struct Attempt {
    target_allocation: usize,
    observed_allocations: usize,
    injected: bool,
    injected_allocation: Option<Allocation>,
    outcome: Outcome,
}

impl Attempt {
    /// Zero-based allocation attempt selected for failure.
    #[must_use]
    pub fn target_allocation(&self) -> usize {
        self.target_allocation
    }

    /// Number of allocation attempts observed during this run.
    #[must_use]
    pub fn observed_allocations(&self) -> usize {
        self.observed_allocations
    }

    /// Returns whether the selected allocation attempt was actually reached and
    /// failed.
    #[must_use]
    pub fn injected(&self) -> bool {
        self.injected
    }

    /// Returns metadata for the allocation attempt that was failed, if the
    /// target allocation was reached.
    #[must_use]
    pub fn injected_allocation(&self) -> Option<Allocation> {
        self.injected_allocation
    }

    /// Outcome of this run.
    pub fn outcome(&self) -> &Outcome {
        &self.outcome
    }

    /// Returns `true` if the selected allocation was injected and the closure
    /// completed without panic.
    #[must_use]
    pub fn is_success(&self) -> bool {
        self.injected && self.outcome.is_completed()
    }
}

/// Outcome of one checked run.
#[derive(Debug, Clone, PartialEq, Eq)]
#[must_use]
pub enum Outcome {
    /// The closure returned normally.
    Completed,

    /// The closure panicked.
    Panicked(Panic),
}

impl Outcome {
    /// Returns `true` if the run returned normally.
    #[must_use]
    pub fn is_completed(&self) -> bool {
        matches!(self, Self::Completed)
    }

    /// Returns `true` if the run panicked.
    #[must_use]
    pub fn is_panicked(&self) -> bool {
        matches!(self, Self::Panicked(_))
    }
}

impl fmt::Display for Outcome {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Completed => f.write_str("completed"),
            Self::Panicked(panic) => write!(f, "panicked: {panic}"),
        }
    }
}

/// Captured panic payload summary.
#[derive(Debug, Clone, PartialEq, Eq)]
#[must_use]
pub struct Panic {
    message: Option<String>,
}

impl Panic {
    /// Returns the panic message when the payload was a `String` or
    /// `&'static str`.
    #[must_use]
    pub fn message(&self) -> Option<&str> {
        self.message.as_deref()
    }
}

impl fmt::Display for Panic {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.message() {
            Some(message) => f.write_str(message),
            None => f.write_str("non-string panic payload"),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct FailurePlan {
    start: usize,
    end: usize,
}

impl FailurePlan {
    fn for_check(check: Check, baseline_allocations: usize) -> Self {
        let selection_start = check.target_start.min(baseline_allocations);
        let selection_end = check
            .target_end
            .unwrap_or(baseline_allocations)
            .min(baseline_allocations);
        let selected_failures = selection_end.saturating_sub(selection_start);
        let failure_limit = check
            .max_failures
            .map_or(selected_failures, |max| max.min(selected_failures));

        Self {
            start: selection_start,
            end: selection_start + failure_limit,
        }
    }

    fn len(self) -> usize {
        self.end - self.start
    }

    fn targets(self) -> Range<usize> {
        self.start..self.end
    }

    fn is_exhaustive_for(self, baseline_allocations: usize) -> bool {
        self.start == 0 && self.end == baseline_allocations
    }

    fn is_exhaustive_after_stop(self, stopped_target: usize, baseline_allocations: usize) -> bool {
        self.start == 0 && stopped_target.saturating_add(1) >= baseline_allocations
    }
}

#[derive(Debug, Clone, Copy)]
enum RunMode {
    Count,
    Fail { target: usize },
}

#[derive(Debug, Clone, Copy)]
struct RunStats {
    observed_allocations: usize,
    injected: bool,
    injected_allocation: Option<Allocation>,
}

#[derive(Debug)]
struct ActiveCheck;

impl ActiveCheck {
    fn enter() -> Result<Self, AlreadyRunning> {
        CHECK_ACTIVE
            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
            .map(|_| Self)
            .map_err(|_| AlreadyRunning)
    }
}

impl Drop for ActiveCheck {
    fn drop(&mut self) {
        MODE.store(MODE_DISABLED, Ordering::SeqCst);
        CHECK_ACTIVE.store(false, Ordering::SeqCst);
    }
}

#[derive(Debug)]
struct ModeGuard {
    finished: bool,
    previous_thread_tracking: bool,
}

impl ModeGuard {
    fn enter(mode: RunMode) -> Self {
        let previous_thread_tracking = current_thread_tracking_enabled();
        set_current_thread_tracking(true);

        SEEN.store(0, Ordering::SeqCst);
        INJECTED.store(false, Ordering::SeqCst);
        INJECTED_INDEX.store(NO_INJECTED_INDEX, Ordering::SeqCst);
        INJECTED_OP.store(ALLOC_OP_NONE, Ordering::SeqCst);
        INJECTED_SIZE.store(0, Ordering::SeqCst);
        INJECTED_ALIGN.store(0, Ordering::SeqCst);
        INJECTED_NEW_SIZE.store(NO_REALLOC_NEW_SIZE, Ordering::SeqCst);

        match mode {
            RunMode::Count => {
                TARGET.store(NO_TARGET, Ordering::SeqCst);
                MODE.store(MODE_COUNTING, Ordering::SeqCst);
            }
            RunMode::Fail { target } => {
                TARGET.store(target, Ordering::SeqCst);
                MODE.store(MODE_FAILING, Ordering::SeqCst);
            }
        }

        Self {
            finished: false,
            previous_thread_tracking,
        }
    }

    fn finish(mut self) -> RunStats {
        MODE.store(MODE_DISABLED, Ordering::SeqCst);

        let injected = INJECTED.load(Ordering::SeqCst);
        let injected_allocation = if injected {
            AllocOp::from_u8(INJECTED_OP.load(Ordering::SeqCst)).map(|operation| Allocation {
                index: INJECTED_INDEX.load(Ordering::SeqCst),
                operation,
                size: INJECTED_SIZE.load(Ordering::SeqCst),
                align: INJECTED_ALIGN.load(Ordering::SeqCst),
                new_size: non_sentinel_usize(INJECTED_NEW_SIZE.load(Ordering::SeqCst)),
            })
        } else {
            None
        };

        let stats = RunStats {
            observed_allocations: SEEN.load(Ordering::SeqCst),
            injected,
            injected_allocation,
        };

        set_current_thread_tracking(self.previous_thread_tracking);
        self.finished = true;
        stats
    }
}

impl Drop for ModeGuard {
    fn drop(&mut self) {
        if !self.finished {
            MODE.store(MODE_DISABLED, Ordering::SeqCst);
            set_current_thread_tracking(self.previous_thread_tracking);
        }
    }
}

fn run_counting<F>(f: &F) -> Baseline
where
    F: Fn(),
{
    let (stats, outcome) = run_in_mode(RunMode::Count, f);
    Baseline {
        observed_allocations: stats.observed_allocations,
        outcome,
    }
}

fn run_failing<F>(target_allocation: usize, f: &F) -> Attempt
where
    F: Fn(),
{
    let (stats, outcome) = run_in_mode(
        RunMode::Fail {
            target: target_allocation,
        },
        f,
    );

    Attempt {
        target_allocation,
        observed_allocations: stats.observed_allocations,
        injected: stats.injected,
        injected_allocation: stats.injected_allocation,
        outcome,
    }
}

fn run_in_mode<F>(mode: RunMode, f: &F) -> (RunStats, Outcome)
where
    F: Fn(),
{
    let guard = ModeGuard::enter(mode);
    let result = panic::catch_unwind(AssertUnwindSafe(f));
    let stats = guard.finish();

    let outcome = match result {
        Ok(()) => Outcome::Completed,
        Err(payload) => Outcome::Panicked(panic_from_payload(payload)),
    };

    (stats, outcome)
}

fn panic_from_payload(payload: Box<dyn Any + Send>) -> Panic {
    let message = if let Some(message) = payload.downcast_ref::<&'static str>() {
        Some((*message).to_owned())
    } else {
        payload.downcast_ref::<String>().cloned()
    };

    Panic { message }
}

fn baseline_matches(reference: &Baseline, candidate: &Baseline) -> bool {
    reference.outcome.is_completed()
        && candidate.outcome.is_completed()
        && reference.observed_allocations == candidate.observed_allocations
}

fn set_current_thread_tracking(enabled: bool) {
    TRACK_ALLOCATIONS_ON_THREAD.with(|tracking| tracking.set(enabled));
}

fn current_thread_tracking_enabled() -> bool {
    TRACK_ALLOCATIONS_ON_THREAD
        .try_with(Cell::get)
        .unwrap_or(false)
}

fn should_inject_failure(operation: AllocOp, layout: Layout, new_size: Option<usize>) -> bool {
    ALLOCATOR_WAS_USED.store(true, Ordering::SeqCst);

    // Do not inject into panic machinery. If the code under test panics, the
    // checker should report that panic rather than potentially failing an
    // allocation while the panic runtime is building or formatting payloads.
    if std::thread::panicking() {
        return false;
    }

    let mode = MODE.load(Ordering::SeqCst);

    if mode == MODE_DISABLED || !current_thread_tracking_enabled() {
        return false;
    }

    let index = SEEN.fetch_add(1, Ordering::SeqCst);

    if mode == MODE_FAILING && index == TARGET.load(Ordering::SeqCst) {
        INJECTED_INDEX.store(index, Ordering::SeqCst);
        INJECTED_OP.store(operation.as_u8(), Ordering::SeqCst);
        INJECTED_SIZE.store(layout.size(), Ordering::SeqCst);
        INJECTED_ALIGN.store(layout.align(), Ordering::SeqCst);
        INJECTED_NEW_SIZE.store(new_size.unwrap_or(NO_REALLOC_NEW_SIZE), Ordering::SeqCst);
        INJECTED.store(true, Ordering::SeqCst);
        true
    } else {
        false
    }
}

fn probe_allocator_installed() -> bool {
    ALLOCATOR_WAS_USED.store(false, Ordering::SeqCst);

    let layout = Layout::from_size_align(1, 1).expect("valid allocator probe layout");

    // SAFETY: The layout is non-zero-sized and has a valid power-of-two
    // alignment. The returned pointer is deallocated with the same layout if
    // allocation succeeds.
    let ptr = unsafe { std::alloc::alloc(layout) };

    if !ptr.is_null() {
        // SAFETY: `ptr` was allocated by `std::alloc::alloc` with `layout`
        // above and has not been deallocated yet.
        unsafe { std::alloc::dealloc(ptr, layout) };
    }

    ALLOCATOR_WAS_USED.load(Ordering::SeqCst)
}

fn non_sentinel_usize(value: usize) -> Option<usize> {
    if value == NO_REALLOC_NEW_SIZE {
        None
    } else {
        Some(value)
    }
}

#[cfg(test)]
#[global_allocator]
static TEST_ALLOCATOR: ChaosAllocator = ChaosAllocator::system();

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::{Mutex, MutexGuard};

    static TEST_LOCK: Mutex<()> = Mutex::new(());

    fn test_lock() -> MutexGuard<'static, ()> {
        TEST_LOCK
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    fn two_fallible_allocations() {
        let mut first = Vec::<u8>::new();
        let mut second = Vec::<u8>::new();

        let _ = first.try_reserve_exact(64);
        let _ = second.try_reserve_exact(64);
    }

    fn with_quiet_expected_panics<R>(f: impl FnOnce() -> R) -> R {
        let previous_hook = std::panic::take_hook();
        std::panic::set_hook(Box::new(|_| {}));

        let result = std::panic::catch_unwind(AssertUnwindSafe(f));

        std::panic::set_hook(previous_hook);

        match result {
            Ok(value) => value,
            Err(payload) => std::panic::resume_unwind(payload),
        }
    }

    #[test]
    fn reports_zero_allocations_for_empty_closure() {
        let _lock = test_lock();

        let report = check(|| {});

        assert_eq!(report.baseline_allocations(), 0);
        assert!(report.attempts().is_empty());
        assert!(report.allocator_installed());
        assert!(report.is_success(), "{report}");
    }

    #[test]
    fn injects_try_reserve_failure() {
        let _lock = test_lock();

        let report = check(|| {
            let mut values = Vec::<u8>::new();
            let _ = values.try_reserve_exact(1024);
        });

        assert_eq!(report.baseline_allocations(), 1, "{report}");
        assert_eq!(report.attempts().len(), 1);
        assert!(report.attempts()[0].injected(), "{report}");
        assert!(report.attempts()[0].injected_allocation().is_some());
        assert!(report.is_success(), "{report}");
    }

    #[test]
    fn records_injected_allocation_metadata() {
        let _lock = test_lock();

        let report = check(|| {
            let mut values = Vec::<u8>::new();
            let _ = values.try_reserve_exact(1024);
        });

        let allocation = report.attempts()[0]
            .injected_allocation()
            .expect("allocation metadata should be recorded");

        assert_eq!(allocation.index(), 0);
        assert_eq!(allocation.operation(), AllocOp::Alloc);
        assert_eq!(allocation.size(), 1024);
        assert_eq!(allocation.align(), 1);
        assert_eq!(allocation.new_size(), None);
    }

    #[test]
    fn reports_baseline_panic() {
        let _lock = test_lock();

        let report = with_quiet_expected_panics(|| check(|| panic!("baseline failed")));

        assert!(!report.is_success());
        assert!(report.attempts().is_empty());
        assert!(matches!(report.baseline().outcome(), Outcome::Panicked(_)));
    }

    #[test]
    fn assert_success_panics_on_failure_report() {
        let _lock = test_lock();

        let report = with_quiet_expected_panics(|| check(|| panic!("expected test panic")));
        let panic =
            with_quiet_expected_panics(|| std::panic::catch_unwind(|| report.assert_success()));

        assert!(panic.is_err());
    }

    #[test]
    fn rejects_nested_checks() {
        let _lock = test_lock();

        let report = check(|| {
            assert_eq!(try_check(|| {}).unwrap_err(), AlreadyRunning);
        });

        assert!(report.is_success(), "{report}");
    }

    #[test]
    fn max_failures_limits_attempts_and_marks_report_truncated() {
        let _lock = test_lock();

        let report = Check::new().max_failures(1).run(two_fallible_allocations);

        assert!(report.baseline_allocations() >= 2, "{report}");
        assert_eq!(report.tested_failures(), 1);
        assert_eq!(report.attempts().len(), 1);
        assert!(report.is_truncated());
        assert_eq!(
            report.untested_failures(),
            report.baseline_allocations() - report.tested_failures()
        );
        assert!(!report.is_success(), "{report}");
    }

    #[test]
    fn unlimited_failures_removes_limit() {
        let _lock = test_lock();

        let report = Check::new()
            .max_failures(1)
            .unlimited_failures()
            .run(two_fallible_allocations);

        assert!(report.baseline_allocations() >= 2, "{report}");
        assert_eq!(report.tested_failures(), report.baseline_allocations());
        assert!(!report.is_truncated(), "{report}");
        assert!(report.is_success(), "{report}");
    }

    #[test]
    fn only_failure_reproduces_one_target_and_marks_report_truncated() {
        let _lock = test_lock();

        let report = Check::new().only_failure(1).run(two_fallible_allocations);

        assert!(report.baseline_allocations() >= 2, "{report}");
        assert_eq!(report.tested_failures(), 1);
        assert_eq!(report.attempts()[0].target_allocation(), 1);
        assert!(report.attempts()[0].is_success(), "{report}");
        assert!(report.is_truncated(), "{report}");
        assert!(!report.is_success(), "{report}");
    }

    #[test]
    fn failure_range_selects_requested_targets() {
        let _lock = test_lock();

        let report = Check::new()
            .failure_range(1..2)
            .run(two_fallible_allocations);

        assert!(report.baseline_allocations() >= 2, "{report}");
        assert_eq!(report.tested_failures(), 1);
        assert_eq!(report.attempts()[0].target_allocation(), 1);
        assert!(report.is_truncated(), "{report}");
    }

    #[test]
    fn failure_range_rejects_reversed_ranges() {
        let _lock = test_lock();

        let panic = std::panic::catch_unwind(|| {
            #[allow(clippy::reversed_empty_ranges)]
            let _ = Check::new().failure_range(2..1);
        });

        assert!(panic.is_err());
    }

    #[test]
    fn stop_on_failure_preserves_range_truncation() {
        let _lock = test_lock();

        let report = Check::new()
            .failure_range(1..2)
            .stop_on_failure(true)
            .run(|| {
                let mut first = Vec::<u8>::new();
                let _ = first.try_reserve_exact(64);

                let mut second = Vec::<u8>::new();
                if second.try_reserve_exact(64).is_err() {
                    panic!("second allocation failed");
                }
            });

        assert!(report.baseline_allocations() >= 2, "{report}");
        assert_eq!(report.tested_failures(), 1, "{report}");
        assert_eq!(report.attempts()[0].target_allocation(), 1);
        assert!(report.is_truncated(), "{report}");
        assert!(!report.is_success(), "{report}");
    }

    #[test]
    fn reports_unstable_baseline() {
        let _lock = test_lock();
        let runs = AtomicUsize::new(0);

        let report = Check::new().stability_runs(2).run(|| {
            if runs.fetch_add(1, Ordering::SeqCst) == 0 {
                let mut values = Vec::<u8>::new();
                let _ = values.try_reserve_exact(64);
            }
        });

        assert!(!report.baseline_is_stable(), "{report}");
        assert_eq!(report.stability_baselines().len(), 1);
        assert!(report.attempts().is_empty(), "{report}");
        assert!(report.is_truncated(), "{report}");
        assert!(!report.is_success(), "{report}");
    }

    #[test]
    fn reports_panic_in_injected_failure_run() {
        let _lock = test_lock();

        let report = with_quiet_expected_panics(|| {
            Check::new().max_failures(1).run(|| {
                let mut values = Vec::<u8>::new();
                if values.try_reserve_exact(1024).is_err() {
                    panic!("allocation failure was not handled");
                }
            })
        });

        assert!(report.baseline().outcome().is_completed(), "{report}");
        assert_eq!(report.attempts().len(), 1);
        assert!(report.attempts()[0].injected(), "{report}");
        assert!(matches!(
            report.attempts()[0].outcome(),
            Outcome::Panicked(panic)
                if panic.message() == Some("allocation failure was not handled")
        ));
        assert!(!report.is_success());
    }

    #[test]
    fn stop_on_failure_stops_after_first_failed_attempt() {
        let _lock = test_lock();

        let report = with_quiet_expected_panics(|| {
            Check::new().stop_on_failure(true).run(|| {
                let mut first = Vec::<u8>::new();
                if first.try_reserve_exact(64).is_err() {
                    panic!("first allocation failed");
                }

                let mut second = Vec::<u8>::new();
                let _ = second.try_reserve_exact(64);
            })
        });

        assert!(report.baseline_allocations() >= 2, "{report}");
        assert_eq!(report.tested_failures(), 1, "{report}");
        assert!(report.is_truncated(), "{report}");
        assert!(!report.is_success());
    }

    #[test]
    fn exposes_failed_attempts() {
        let _lock = test_lock();

        let report = with_quiet_expected_panics(|| {
            Check::new().max_failures(1).run(|| {
                let mut values = Vec::<u8>::new();
                if values.try_reserve_exact(1024).is_err() {
                    panic!("not handled");
                }
            })
        });

        assert!(report.first_failure().is_some());
        assert_eq!(report.failed_attempts().count(), 1);
    }

    #[test]
    fn reports_targets_that_are_not_reached_as_untested() {
        let _lock = test_lock();
        let run = AtomicUsize::new(0);

        let report = Check::new().run(|| {
            let current_run = run.fetch_add(1, Ordering::SeqCst);

            let mut first = Vec::<u8>::new();
            let _ = first.try_reserve_exact(64);

            if current_run != 2 {
                let mut second = Vec::<u8>::new();
                let _ = second.try_reserve_exact(64);
            }
        });

        assert_eq!(report.baseline_allocations(), 2, "{report}");
        assert_eq!(report.tested_failures(), 2, "{report}");
        assert_eq!(report.injected_failures(), 1, "{report}");
        assert_eq!(report.untested_failures(), 1, "{report}");
        assert!(!report.is_success(), "{report}");
    }

    #[test]
    fn display_mentions_untested_failures_when_truncated() {
        let _lock = test_lock();

        let report = Check::new().max_failures(1).run(two_fallible_allocations);
        let rendered = report.to_string();

        assert!(rendered.contains("truncated"), "{rendered}");
        assert!(rendered.contains("untested:"), "{rendered}");
        assert!(rendered.contains("partial success"), "{rendered}");
    }

    #[test]
    fn display_mentions_unstable_baseline() {
        let _lock = test_lock();
        let runs = AtomicUsize::new(0);

        let report = Check::new().stability_runs(2).run(|| {
            if runs.fetch_add(1, Ordering::SeqCst) == 0 {
                let mut values = Vec::<u8>::new();
                let _ = values.try_reserve_exact(64);
            }
        });
        let rendered = report.to_string();

        assert!(rendered.contains("unstable baseline"), "{rendered}");
    }
}