#![feature(alloc_error_hook)]

use std::{
    alloc::{GlobalAlloc, System},
    panic::AssertUnwindSafe,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
};

pub fn main() {
    std::alloc::set_alloc_error_hook(|l| panic!("Failed to alloc: {l:?}"));
    let mut arc = Arc::new(vec![vec![1]]);

    let _weak = Arc::downgrade(&arc); // create a weak to hit the right code path in make_mut

    std::panic::catch_unwind(AssertUnwindSafe(|| {
        ERR_ON_ALLOC.store(true, Ordering::SeqCst); // make the next allocation error
        Arc::make_mut(&mut arc); // sets strong count to 0, then panics out of the method (reallocates because of _weak)
    }))
    .unwrap_err();

    eprintln!("Survived alloc error");

    assert_eq!(Arc::strong_count(&arc), 0);

    drop(arc.clone()); // strong count 0 -> 1 -> 0 leads to vec being dropped

    dbg!(arc); // SIGSEGV
}

// for control of when allocation errors.
static ERR_ON_ALLOC: AtomicBool = AtomicBool::new(false);
struct AllocErrOnDemand;
unsafe impl GlobalAlloc for AllocErrOnDemand {
    unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
        if ERR_ON_ALLOC.swap(false, Ordering::SeqCst) {
            std::ptr::null_mut()
        } else {
            unsafe { System.alloc(layout) }
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
        unsafe { System.dealloc(ptr, layout) }
    }
}
#[global_allocator]
static GLOBAL_ALLOC: AllocErrOnDemand = AllocErrOnDemand;