frame-alloc 0.1.0

extern crate std;

use crate::util::lock::Lock;
use crate::{InterruptControl, NoInterruptControl};
use core::cell::{Cell, RefCell, UnsafeCell};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread;
use std::vec::Vec;

// Per-thread interrupt-state bookkeeping -> tests isolated from one another.
std::thread_local! {
    static DEPTH: Cell<i64> = const { Cell::new(0) }; // current disable nesting
    static MAXD: Cell<i64> = const { Cell::new(0) };  // peak nesting observed
    static DIS: Cell<u64> = const { Cell::new(0) };   // disable() calls
    static RES: Cell<u64> = const { Cell::new(0) };   // restore() calls
}

fn depth() -> i64 {
    DEPTH.with(Cell::get)
}
fn reset() {
    DEPTH.set(0);
    MAXD.set(0);
    DIS.set(0);
    RES.set(0);
}

/// A fake [`InterruptControl`] that records nesting depth. `State` carries the
/// *prior* depth.
struct CountingIrq;
// SAFETY: `disable`/`restore` are faithful inverses on the depth counter.
unsafe impl InterruptControl for CountingIrq {
    type State = i64;
    const INIT: i64 = i64::MIN;
    fn disable() -> i64 {
        DIS.set(DIS.get() + 1);
        let prev = DEPTH.with(Cell::get);
        let now = prev + 1;
        DEPTH.set(now);
        MAXD.with(|c| {
            if now > c.get() {
                c.set(now)
            }
        });
        prev
    }
    unsafe fn restore(prev: i64) {
        RES.set(RES.get() + 1);
        assert!(
            prev >= 0,
            "restore saw the uninitialised INIT sentinel — release without acquire?"
        );
        DEPTH.set(prev);
    }
}

#[test]
fn single_acquire_release_balances() {
    let lock = Lock::<CountingIrq>::new();
    reset();
    lock.acquire();
    assert_eq!(depth(), 1, "acquire must disable once");
    lock.release();
    assert_eq!(depth(), 0, "release must restore");
    assert_eq!(DIS.with(Cell::get), 1);
    assert_eq!(RES.with(Cell::get), 1);
}

#[test]
fn nested_locks_restore_in_order() {
    let outer = Lock::<CountingIrq>::new();
    let inner = Lock::<CountingIrq>::new();
    reset();
    outer.acquire();
    inner.acquire();
    assert_eq!(depth(), 2, "two nested acquires disable twice");
    inner.release();
    assert_eq!(
        depth(),
        1,
        "inner release restores to still-disabled, not enabled"
    );
    outer.release();
    assert_eq!(depth(), 0, "only the outermost release re-enables");
    assert_eq!(MAXD.with(Cell::get), 2);
    assert_eq!(DIS.with(Cell::get), RES.with(Cell::get));
}

#[test]
fn no_interrupt_control_is_zero_sized() {
    assert_eq!(
        size_of::<Lock<NoInterruptControl>>(),
        size_of::<AtomicBool>()
    );
}

/// A lock plus the non-atomic state it guards. The `unsafe impl Sync` asserts
/// that `Lock` itself serialises all access to `counter`; if the lock fails to
/// provide mutual exclusion the counter races and updates are lost.
struct Shared {
    lock: Lock<NoInterruptControl>,
    counter: UnsafeCell<u64>,
    flag: AtomicBool,
}
// SAFETY: `counter` is only ever touched while holding `lock`; `flag` is atomic.
unsafe impl Sync for Shared {}

#[cfg_attr(audit, ignore = "audit builds are sequential-only")]
#[test]
fn mutual_exclusion_under_contention() {
    const THREADS: u64 = if cfg!(miri) { 4 } else { 8 };
    const ITERS: u64 = if cfg!(miri) { 100 } else { 50_000 };
    let shared = Arc::new(Shared {
        lock: Lock::new(),
        counter: UnsafeCell::new(0),
        flag: AtomicBool::new(false),
    });
    let mut handles = Vec::new();
    for _ in 0..THREADS {
        let s = shared.clone();
        handles.push(thread::spawn(move || {
            for _ in 0..ITERS {
                s.lock.acquire();
                // SAFETY: the lock serialises access to `counter`. A non-atomic
                // read-modify-write so that a broken lock loses updates.
                unsafe {
                    let c = s.counter.get();
                    *c = (*c).wrapping_add(1);
                }
                s.lock.release();
            }
        }));
    }
    for h in handles {
        h.join().unwrap();
    }
    // SAFETY: all threads have joined, so we have exclusive access.
    let total = unsafe { *shared.counter.get() };
    assert_eq!(
        total,
        THREADS * ITERS,
        "lost updates: lock failed to provide mutual exclusion"
    );
}

#[cfg_attr(audit, ignore = "audit builds are sequential-only")]
#[test]
fn acquire_blocks_while_held() {
    let shared = Arc::new(Shared {
        lock: Lock::new(),
        counter: UnsafeCell::new(0),
        flag: AtomicBool::new(false),
    });
    shared.lock.acquire();

    let s = shared.clone();
    let handle = thread::spawn(move || {
        s.lock.acquire();
        s.flag.store(true, Ordering::SeqCst);
        s.lock.release();
    });

    thread::sleep(std::time::Duration::from_millis(50));
    assert!(
        !shared.flag.load(Ordering::SeqCst),
        "second acquire returned while the lock was still held"
    );

    shared.lock.release();
    handle.join().unwrap();
    assert!(
        shared.flag.load(Ordering::SeqCst),
        "thread should proceed once the lock is released"
    );
}

// Per-thread stack of live tokens handed out by `disable`, checked by `restore`.
std::thread_local! {
    static TOKENS: RefCell<Vec<u64>> = const { RefCell::new(Vec::new()) };
    static NEXT_TOKEN: Cell<u64> = const { Cell::new(1) };
}

/// An [`InterruptControl`] that hands every `disable` a distinct token and, on
/// `restore`, asserts it received the exact token saved by its matching
/// `acquire`.
struct TokenIrq;
// SAFETY: `restore` checks it is the faithful inverse of its paired `disable`.
unsafe impl InterruptControl for TokenIrq {
    type State = u64;
    const INIT: u64 = 0;
    fn disable() -> u64 {
        let t = NEXT_TOKEN.with(|n| {
            let v = n.get();
            n.set(v + 1);
            v
        });
        TOKENS.with(|s| s.borrow_mut().push(t));
        t
    }
    unsafe fn restore(state: u64) {
        TOKENS.with(|s| {
            let expected = s
                .borrow_mut()
                .pop()
                .expect("restore without a matching disable");
            assert_eq!(
                state, expected,
                "release restored a different state than its acquire saved"
            );
        });
    }
}

#[test]
fn release_restores_each_acquires_own_state() {
    NEXT_TOKEN.set(1);
    TOKENS.with(|s| s.borrow_mut().clear());

    let a = Lock::<TokenIrq>::new();
    let b = Lock::<TokenIrq>::new();
    a.acquire();
    b.acquire();
    b.release();
    a.release();

    TOKENS.with(|s| {
        assert!(
            s.borrow().is_empty(),
            "every disabled token must be restored"
        )
    });
}