async_safe_defer/

sync.rs

use core::{
    fmt,
    marker::PhantomData,
    mem,
    ops::{Deref, DerefMut},
};

/// Decides whether a deferred action runs when its guard is dropped.
pub trait Strategy {
    /// Returns `true` when the action should run in the current drop context.
    fn should_run() -> bool;
}

/// Selects the action whenever an armed guard is dropped.
#[derive(Debug)]
pub enum Always {}

impl Strategy for Always {
    #[inline(always)]
    fn should_run() -> bool {
        true
    }
}

/// Selects the action when an armed guard is dropped outside panic unwinding.
///
/// Returning `Result::Err` is a successful exit for this policy because it does
/// not inspect return values.
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
#[derive(Debug)]
pub enum OnSuccess {}

#[cfg(feature = "std")]
impl Strategy for OnSuccess {
    #[inline]
    fn should_run() -> bool {
        !std::thread::panicking()
    }
}

/// Selects the action when an armed guard is dropped during panic unwinding.
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
#[derive(Debug)]
pub enum OnUnwind {}

#[cfg(feature = "std")]
impl Strategy for OnUnwind {
    #[inline]
    fn should_run() -> bool {
        std::thread::panicking()
    }
}

/// Holds a synchronous action for at-most-once deferred execution.
///
/// The strategy decides whether the action runs when the guard is dropped.
/// Calling [`DeferGuard::disarm`] returns the action without executing it, while
/// [`DeferGuard::run_now`] executes it immediately regardless of the strategy.
/// Panics from the strategy or a drop-time action propagate. A second panic
/// during unwinding may abort the process.
#[must_use = "store the guard so its deferred action runs at the intended scope exit"]
pub struct DeferGuard<F: FnOnce(), S: Strategy = Always> {
    action: Option<F>,
    strategy: PhantomData<fn() -> S>,
}

impl<F: FnOnce()> DeferGuard<F> {
    /// Arms `action` with the [`Always`] strategy.
    #[inline]
    pub const fn new(action: F) -> Self {
        Self::with_strategy(action)
    }
}

impl<F: FnOnce(), S: Strategy> DeferGuard<F, S> {
    /// Arms `action` with strategy `S`.
    #[inline]
    pub const fn with_strategy(action: F) -> Self {
        Self {
            action: Some(action),
            strategy: PhantomData,
        }
    }

    /// Disarms the guard and returns its action without executing it.
    #[inline]
    pub fn disarm(mut self) -> F {
        self.action
            .take()
            .expect("an armed defer guard always contains its action")
    }

    /// Executes the action immediately, regardless of `S`, and consumes the guard.
    #[inline]
    pub fn run_now(mut self) {
        if let Some(action) = self.action.take() {
            action();
        }
    }
}

impl<F: FnOnce(), S: Strategy> Drop for DeferGuard<F, S> {
    #[inline]
    fn drop(&mut self) {
        let Some(action) = self.action.take() else {
            return;
        };

        if S::should_run() {
            action();
        }
    }
}

impl<F: FnOnce(), S: Strategy> fmt::Debug for DeferGuard<F, S> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("DeferGuard")
            .field("armed", &self.action.is_some())
            .finish()
    }
}

enum ScopeState<T, F> {
    Armed { value: T, action: F },
    Disarmed,
}

/// Owns a value and conditionally passes it to an action when dropped.
///
/// The guarded value is accessible through [`Deref`] and [`DerefMut`]. The
/// action receives ownership of the final value and runs at most once.
/// Panics from the strategy or a drop-time action propagate. A second panic
/// during unwinding may abort the process.
#[must_use = "store the guard so its deferred action runs at the intended scope exit"]
pub struct ScopeGuard<T, F, S = Always>
where
    F: FnOnce(T),
    S: Strategy,
{
    state: ScopeState<T, F>,
    strategy: PhantomData<fn() -> S>,
}

impl<T, F, S> ScopeGuard<T, F, S>
where
    F: FnOnce(T),
    S: Strategy,
{
    /// Creates an armed guard using strategy `S`.
    #[inline]
    pub const fn with_strategy(value: T, action: F) -> Self {
        Self {
            state: ScopeState::Armed { value, action },
            strategy: PhantomData,
        }
    }

    /// Disarms the guard and returns its value without running the action.
    #[inline]
    pub fn into_inner(mut self) -> T {
        let (value, action) = match mem::replace(&mut self.state, ScopeState::Disarmed) {
            ScopeState::Armed { value, action } => (value, action),
            ScopeState::Disarmed => unreachable!("a live scope guard is armed"),
        };
        drop(action);
        value
    }

    /// Disarms the guard and returns both its value and action.
    #[inline]
    pub fn into_parts(mut self) -> (T, F) {
        match mem::replace(&mut self.state, ScopeState::Disarmed) {
            ScopeState::Armed { value, action } => (value, action),
            ScopeState::Disarmed => unreachable!("a live scope guard is armed"),
        }
    }

    /// Runs the action immediately, regardless of `S`, and consumes the guard.
    #[inline]
    pub fn run_now(mut self) {
        if let ScopeState::Armed { value, action } =
            mem::replace(&mut self.state, ScopeState::Disarmed)
        {
            action(value);
        }
    }

    fn value(&self) -> &T {
        match &self.state {
            ScopeState::Armed { value, .. } => value,
            ScopeState::Disarmed => unreachable!("a live scope guard is armed"),
        }
    }

    fn value_mut(&mut self) -> &mut T {
        match &mut self.state {
            ScopeState::Armed { value, .. } => value,
            ScopeState::Disarmed => unreachable!("a live scope guard is armed"),
        }
    }
}

impl<T, F, S> Deref for ScopeGuard<T, F, S>
where
    F: FnOnce(T),
    S: Strategy,
{
    type Target = T;

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.value()
    }
}

impl<T, F, S> DerefMut for ScopeGuard<T, F, S>
where
    F: FnOnce(T),
    S: Strategy,
{
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.value_mut()
    }
}

impl<T, F, S> Drop for ScopeGuard<T, F, S>
where
    F: FnOnce(T),
    S: Strategy,
{
    #[inline]
    fn drop(&mut self) {
        let ScopeState::Armed { value, action } =
            mem::replace(&mut self.state, ScopeState::Disarmed)
        else {
            return;
        };

        if S::should_run() {
            action(value);
        }
    }
}

impl<T, F, S> fmt::Debug for ScopeGuard<T, F, S>
where
    T: fmt::Debug,
    F: FnOnce(T),
    S: Strategy,
{
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("ScopeGuard")
            .field("value", self.value())
            .finish()
    }
}

/// Arms a new [`DeferGuard`] with the [`Always`] strategy.
#[inline]
pub const fn defer<F: FnOnce()>(action: F) -> DeferGuard<F> {
    DeferGuard::new(action)
}

/// Arms a new [`DeferGuard`] that runs on a non-panicking drop.
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
#[inline]
pub const fn defer_on_success<F: FnOnce()>(action: F) -> DeferGuard<F, OnSuccess> {
    DeferGuard::with_strategy(action)
}

/// Arms a new [`DeferGuard`] that runs while unwinding from a panic.
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
#[inline]
pub const fn defer_on_unwind<F: FnOnce()>(action: F) -> DeferGuard<F, OnUnwind> {
    DeferGuard::with_strategy(action)
}

/// Owns `value` and passes it to `action` when the guard is dropped.
#[inline]
pub const fn guard<T, F>(value: T, action: F) -> ScopeGuard<T, F>
where
    F: FnOnce(T),
{
    ScopeGuard::with_strategy(value, action)
}

/// Owns `value` and passes it to `action` on a non-panicking drop.
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
#[inline]
pub const fn guard_on_success<T, F>(value: T, action: F) -> ScopeGuard<T, F, OnSuccess>
where
    F: FnOnce(T),
{
    ScopeGuard::with_strategy(value, action)
}

/// Owns `value` and passes it to `action` while unwinding from a panic.
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
#[inline]
pub const fn guard_on_unwind<T, F>(value: T, action: F) -> ScopeGuard<T, F, OnUnwind>
where
    F: FnOnce(T),
{
    ScopeGuard::with_strategy(value, action)
}

/// Binds a deferred action that runs on a non-panicking scope exit.
///
/// The action must evaluate to `()`.
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
#[macro_export]
macro_rules! defer_on_success {
    (move $($body:tt)*) => {
        let _defer_on_success_guard = $crate::defer_on_success(move || { $($body)* });
    };
    ($($body:tt)*) => {
        let _defer_on_success_guard = $crate::defer_on_success(|| { $($body)* });
    };
}

/// Binds a deferred action that runs while unwinding from a panic.
///
/// The action must evaluate to `()`. If it panics, the second panic may abort
/// the process.
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
#[macro_export]
macro_rules! defer_on_unwind {
    (move $($body:tt)*) => {
        let _defer_on_unwind_guard = $crate::defer_on_unwind(move || { $($body)* });
    };
    ($($body:tt)*) => {
        let _defer_on_unwind_guard = $crate::defer_on_unwind(|| { $($body)* });
    };
}