#![feature(allocator_api)]
#![cfg_attr(not(test), allow(dead_code, unused_features))]

//! Minimized standalone witness for allocator-backed `Arc::make_mut`
//! panic-unsafety.
//!
//! Stable regression-style witness:
//!
//! ```sh
//! cargo +nightly test --manifest-path alloc-arc-make-mut-ub-repro/Cargo.toml \
//!   arc_make_mut_preserves_live_arc_after_allocator_clone_panic -- --exact --nocapture
//! ```
//!
//! Miri UB witness:
//!
//! ```sh
//! cargo +nightly miri test --manifest-path alloc-arc-make-mut-ub-repro/Cargo.toml \
//!   miri_arc_make_mut_clone_drop_clone_then_read_original_after_allocator_clone_panic \
//!   -- --exact --nocapture
//! ```

use std::{
    alloc::{AllocError, Allocator, Global, Layout},
    mem::forget,
    panic::{AssertUnwindSafe, catch_unwind},
    ptr::NonNull,
    sync::{
        Arc, Weak,
        atomic::{AtomicUsize, Ordering::SeqCst},
    },
};

#[derive(Debug)]
struct PanicAlloc {
    state: Arc<PanicAllocState>,
}

#[derive(Debug)]
struct PanicAllocState {
    clone_count: AtomicUsize,
    panic_on: AtomicUsize,
}

impl PanicAlloc {
    fn new() -> Self {
        Self {
            state: Arc::new(PanicAllocState {
                clone_count: AtomicUsize::new(0),
                panic_on: AtomicUsize::new(usize::MAX),
            }),
        }
    }

    fn arm(&self, panic_on: usize) {
        self.state.clone_count.store(0, SeqCst);
        self.state.panic_on.store(panic_on, SeqCst);
    }

    fn disarm(&self) {
        self.state.panic_on.store(usize::MAX, SeqCst);
    }
}

impl Clone for PanicAlloc {
    fn clone(&self) -> Self {
        let clone_number = self.state.clone_count.fetch_add(1, SeqCst) + 1;
        if clone_number == self.state.panic_on.load(SeqCst) {
            panic!("allocator clone panic");
        }
        Self {
            state: Arc::clone(&self.state),
        }
    }
}

unsafe impl Allocator for PanicAlloc {
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        Global.allocate(layout)
    }

    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
        unsafe { Global.deallocate(ptr, layout) }
    }
}

fn panics_silently<T>(f: impl FnOnce() -> T) -> bool {
    let previous_hook = std::panic::take_hook();
    std::panic::set_hook(Box::new(|_| {}));
    let panicked = catch_unwind(AssertUnwindSafe(f)).is_err();
    std::panic::set_hook(previous_hook);
    panicked
}

fn corrupt_arc_make_mut_pair() -> (Arc<String, PanicAlloc>, Weak<String, PanicAlloc>) {
    let alloc = PanicAlloc::new();
    let mut arc = Arc::new_in(String::from("hello"), alloc.clone());
    let weak = Arc::downgrade(&arc);

    alloc.arm(1);
    let make_mut_panicked = panics_silently(|| {
        let _ = Arc::make_mut(&mut arc);
    });
    alloc.disarm();

    assert!(make_mut_panicked, "setup should panic inside Arc::make_mut");
    (arc, weak)
}

#[test]
fn arc_make_mut_preserves_live_arc_after_allocator_clone_panic() {
    let (arc, weak) = corrupt_arc_make_mut_pair();
    let strong_after_panic = Arc::strong_count(&arc);
    let weak_upgrades_after_panic = weak.upgrade().is_some();

    // Avoid exercising later drop behavior when validating the immediate
    // post-unwind state on a vulnerable toolchain.
    forget(weak);
    forget(arc);

    assert_eq!(
        strong_after_panic, 1,
        "the original Arc should still report one strong reference after the panic"
    );
    assert!(
        weak_upgrades_after_panic,
        "the existing Weak should still upgrade after the panic"
    );
}

#[cfg(miri)]
#[test]
fn miri_arc_make_mut_clone_drop_clone_then_read_original_after_allocator_clone_panic() {
    let (arc, weak) = corrupt_arc_make_mut_pair();

    // Safe code after the caught panic: resurrect the corrupted `Arc` with a
    // clone, drop that clone so it tears down the shared `String`, then read
    // through the original handle.
    let cloned = Arc::clone(&arc);
    drop(cloned);
    assert_eq!(arc.as_bytes()[0], b'h');
    drop(weak);
}