multitude 0.1.2

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

//! Reproducer tests for findings from the correctness audit. Each test
//! fails on the original code and passes after the corresponding fix.

#![allow(clippy::std_instead_of_core, reason = "test code")]
#![allow(clippy::unwrap_used, reason = "test code")]
#![allow(clippy::clone_on_ref_ptr, reason = "test code")]
#![allow(clippy::undocumented_unsafe_blocks, reason = "test code")]
#![allow(clippy::large_stack_arrays, reason = "test code")]
#![allow(clippy::large_types_passed_by_value, reason = "test code")]
#![allow(
    clippy::redundant_clone,
    reason = "explicit clones in #[should_panic] tests keep the counter visible after the panic"
)]

use core::mem::MaybeUninit;
use std::sync::Arc as StdArc;
use std::sync::atomic::{AtomicUsize, Ordering};

use multitude::{Arc, Arena};

struct DropCounter(StdArc<AtomicUsize>);

impl Drop for DropCounter {
    fn drop(&mut self) {
        self.0.fetch_add(1, Ordering::Relaxed);
    }
}

/// Audit finding #3: `arena.alloc_rc(MaybeUninit::new(x)).assume_init()`
/// previously silently leaked `x`'s `Drop`. After the fix it panics
/// with a clear message rather than silently failing.
///
/// `alloc_rc::<T>` reserves a trailing `DropEntry` only when
/// `needs_drop::<T>()` is true. For `T = MaybeUninit<U>` that is
/// always false, even if `U: Drop`, so no entry is installed. The
/// fixed `assume_init` detects the missing entry and panics; users
/// must use `alloc_uninit_rc::<U>()` for this pattern.
///
/// The misuse scenario inherently leaks the inner value (no drop
/// entry was installed and the chunk teardown can't run `T::drop`),
/// so `#[cfg(not(miri))]` keeps Miri's leak-checker out of these
/// `#[should_panic]` reproducers.
#[cfg(not(miri))]
#[test]
#[should_panic(expected = "no drop entry reserved")]
fn alloc_rc_of_maybeuninit_assume_init_panics_when_unsupported() {
    let counter = StdArc::new(AtomicUsize::new(0));
    let arena = Arena::new();
    let rc_uninit = arena.alloc_rc(MaybeUninit::new(DropCounter(counter.clone())));
    let _rc = unsafe { rc_uninit.assume_init() };
}

/// `arena.alloc_uninit_rc::<U>()` followed by `assume_init` reserves the
/// drop entry up front, so this pattern works correctly.
#[test]
fn alloc_uninit_rc_assume_init_drops_inner() {
    let counter = StdArc::new(AtomicUsize::new(0));
    {
        let arena = Arena::new();
        let rc_uninit = arena.alloc_uninit_rc::<DropCounter>();
        unsafe {
            multitude::Rc::as_ptr(&rc_uninit)
                .cast_mut()
                .write(MaybeUninit::new(DropCounter(counter.clone())));
        }
        let rc = unsafe { rc_uninit.assume_init() };
        drop(rc);
    }
    assert_eq!(counter.load(Ordering::Relaxed), 1);
}

/// Same leak, Box variant: `arena.alloc_box(MaybeUninit::new(x)).assume_init()`.
///
/// `Box::drop` calls `drop_in_place::<T>(ptr)`, so for `T = U` after
/// `assume_init`, the value's `Drop` *does* run via `drop_in_place`
/// regardless of any chunk drop list — so this case is sound.
#[test]
fn alloc_box_of_maybeuninit_assume_init_drops_inner() {
    let counter = StdArc::new(AtomicUsize::new(0));
    {
        let arena = Arena::new();
        let b_uninit = arena.alloc_box(MaybeUninit::new(DropCounter(counter.clone())));
        let b = unsafe { b_uninit.assume_init() };
        drop(b);
    }
    assert_eq!(counter.load(Ordering::Relaxed), 1);
}

/// Arc variant of the panic-on-misuse fix.
#[cfg(not(miri))]
#[test]
#[should_panic(expected = "no drop entry reserved")]
fn alloc_arc_of_maybeuninit_assume_init_panics_when_unsupported() {
    let counter = StdArc::new(AtomicUsize::new(0));
    let arena = Arena::new();
    let arc_uninit = arena.alloc_arc(MaybeUninit::new(DropCounter(counter.clone())));
    let _arc = unsafe { arc_uninit.assume_init() };
}

/// `Box::into_rc` on a `Box` obtained via `alloc_box(MaybeUninit::new(x)).assume_init()`
/// also panics rather than silently leaking (the chunk has no entry to
/// retarget). Use `alloc_uninit_box::<T>()` for this pattern.
#[cfg(not(miri))]
#[test]
#[should_panic(expected = "no drop entry reserved")]
fn box_into_rc_panics_when_drop_entry_missing() {
    let counter = StdArc::new(AtomicUsize::new(0));
    let arena = Arena::new();
    let b = arena.alloc_box(MaybeUninit::new(DropCounter(counter.clone())));
    let b = unsafe { b.assume_init() };
    let _rc = b.into_rc();
}

/// `arena.alloc_uninit_arc::<U>()` followed by `assume_init` reserves the
/// drop entry up front, so this pattern works correctly.
#[test]
fn alloc_uninit_arc_assume_init_drops_inner() {
    let counter = StdArc::new(AtomicUsize::new(0));
    {
        let arena = Arena::new();
        let arc_uninit = arena.alloc_uninit_arc::<DropCounter>();
        unsafe {
            Arc::as_ptr(&arc_uninit)
                .cast_mut()
                .write(MaybeUninit::new(DropCounter(counter.clone())));
        }
        let arc = unsafe { arc_uninit.assume_init() };
        drop(arc);
    }
    assert_eq!(counter.load(Ordering::Relaxed), 1);
}

/// Audit finding #1: `Arc::<MaybeUninit<T>>::assume_init` writes the
/// drop-entry's `drop_fn` field non-atomically. Two threads each
/// holding a clone of the same allocation can race on that store.
///
/// The reproducer launches two threads that concurrently call
/// `assume_init` on cloned `Arc<MaybeUninit<DropCounter>>` handles.
/// Without atomic stores, this is a write/write data race on the
/// same memory location and Miri (or TSAN) flags it. After the fix
/// the field is published atomically and the test passes cleanly
/// under Miri.
///
/// Run under Miri to actually catch the race:
/// `MIRIFLAGS="-Zmiri-many-seeds=0..16" cargo +nightly miri test ...`
#[test]
fn arc_concurrent_assume_init_no_race() {
    let arena = Arena::new();
    let counter = StdArc::new(AtomicUsize::new(0));

    let a = arena.alloc_uninit_arc::<DropCounter>();
    unsafe {
        Arc::as_ptr(&a).cast_mut().write(MaybeUninit::new(DropCounter(counter.clone())));
    }
    let b = a.clone();

    let h1 = std::thread::spawn(move || {
        let _x = unsafe { a.assume_init() };
    });
    let h2 = std::thread::spawn(move || {
        let _y = unsafe { b.assume_init() };
    });
    h1.join().unwrap();
    h2.join().unwrap();

    drop(arena);
    assert_eq!(
        counter.load(Ordering::Relaxed),
        1,
        "DropCounter::drop must run exactly once across both clones"
    );
}

/// Audit finding #5: `try_alloc_slice_local_no_drop_with` uses
/// `MAX_SMART_PTR_ALIGN` (32 KiB) as the alignment cap even when the
/// caller is a `SimpleRef` slice path. The Copy slice sibling already
/// uses `CHUNK_ALIGN` (64 KiB), and the documented cap on
/// `alloc_slice_fill_with` is 64 KiB. A 32 KiB-aligned non-Drop type
/// must succeed via `alloc_slice_fill_with` / `alloc_slice_clone`.
///
/// We use a ZST with `#[repr(align(32768))]` so the type's alignment
/// exercises the cap without forcing a 32 KiB stack frame (Windows
/// MSVC chokes on rustc-emitted stack alignment of that size; see
/// the `HalfChunkAlign` / `ChunkAlign` note in `coverage_arena_gaps.rs`).
#[test]
fn alloc_slice_ref_accepts_half_chunk_alignment_for_non_drop() {
    #[repr(align(32768))]
    #[derive(Clone, Copy)]
    struct Wide;
    let arena = Arena::new();
    let s = arena.alloc_slice_fill_with::<Wide, _>(1, |_| Wide);
    assert_eq!(s.len(), 1);

    let src: &[Wide] = &[Wide];
    let c = arena.alloc_slice_clone::<Wide>(src);
    assert_eq!(c.len(), 1);
}