/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! Intern objects in memory.
//!
//! This is similar to [`internment` crate](https://github.com/droundy/internment)
//! but with changes for performance and flexibility.

#![cfg_attr(feature = "gazebo_lint", feature(plugin))]
#![cfg_attr(feature = "gazebo_lint", allow(deprecated))] // :(
#![cfg_attr(feature = "gazebo_lint", plugin(gazebo_lint))]

use std::borrow::Borrow;
use std::cmp::Ordering;
use std::collections::hash_map::DefaultHasher;
use std::fmt;
use std::fmt::Display;
use std::fmt::Formatter;
use std::hash::Hash;
use std::hash::Hasher;
use std::marker;
use std::mem;
use std::ops::Deref;
use std::ptr;

use allocative::Allocative;
use allocative::Visitor;
use dupe::Dupe;
use lock_free_hashtable::sharded::ShardedLockFreeRawTable;

pub struct StaticInterner<T: 'static, H = DefaultHasher> {
    table: ShardedLockFreeRawTable<Box<InternedData<T>>, 64>,
    _marker: marker::PhantomData<H>,
}

/// This structure is similar to `Hashed<T>`, but it is not parameterized by hash function.
#[derive(Debug)]
struct InternedData<T: 'static> {
    data: T,
    hash: u64,
}

/// An interned pointer.
///
/// Equality of this type is a pointer comparison.
/// But note, this works correctly only if `Intern` pointers created
/// from the same instance of `StaticInterner`.
#[derive(Debug)]
pub struct Intern<T: 'static> {
    pointer: &'static InternedData<T>,
}

// TODO(nga): derive.
impl<T: Allocative> Allocative for Intern<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        if mem::size_of::<T>() > 0 {
            let visitor = visitor.enter_shared(
                allocative::Key::new("pointer"),
                mem::size_of::<*const T>(),
                &**self as &T as *const T as *const (),
            );
            if let Some(mut visitor) = visitor {
                (**self).visit(&mut visitor);
                visitor.exit();
            }
        }
    }
}

impl<T: 'static> Copy for Intern<T> {}

impl<T: 'static> Clone for Intern<T> {
    #[inline]
    fn clone(&self) -> Self {
        *self
    }
}

impl<T: 'static> Dupe for Intern<T> {
    #[inline]
    fn dupe(&self) -> Self {
        *self
    }
}

impl<T: 'static> Deref for Intern<T> {
    type Target = T;

    #[inline]
    fn deref(&self) -> &T {
        &self.pointer.data
    }
}

impl<T: 'static> Intern<T> {
    #[inline]
    pub fn deref_static(&self) -> &'static T {
        &self.pointer.data
    }
}

impl<T> Hash for Intern<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        // We could hash only the pointer, since we only compare the pointers,
        // but users may expect hashing to be stable between runs.
        self.pointer.hash.hash(state);
    }
}

impl<T> PartialEq for Intern<T> {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        ptr::eq(self.pointer, other.pointer)
    }
}

impl<T> Eq for Intern<T> {}

impl<T: PartialOrd> PartialOrd for Intern<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.pointer.data.partial_cmp(&other.pointer.data)
    }
}

impl<T: Ord> Ord for Intern<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.pointer.data.cmp(&other.pointer.data)
    }
}

impl<T: Display> Display for Intern<T> {
    #[inline]
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        Display::fmt(&self.pointer.data, f)
    }
}

/// Hash the value before acquiring the lock.
struct Hashed<T, H> {
    hash: u64,
    value: T,
    _marker: marker::PhantomData<H>,
}

impl<T: Hash, H: Hasher + Default> Hashed<T, H> {
    /// Compute the hash.
    fn hash(value: &T) -> u64 {
        let mut hasher = H::default();
        value.hash(&mut hasher);
        hasher.finish()
    }

    fn new(value: T) -> Self {
        let hash = Self::hash(&value);
        Hashed {
            hash,
            value,
            _marker: marker::PhantomData,
        }
    }
}

/// `Self` is "equivalent" to `K` which means `Self` and `K` are logically "equal"
/// and hashes of `Self` and `Q` are equal.
pub trait Equiv<K: ?Sized> {
    /// Compare `Self` and `K`.
    fn equivalent(&self, key: &K) -> bool;
}

impl<Q: ?Sized, K: ?Sized> Equiv<K> for &Q
where
    Q: Eq,
    K: Borrow<Q>,
{
    #[inline]
    fn equivalent(&self, key: &K) -> bool {
        *self == key.borrow()
    }
}

impl<T: 'static, H> StaticInterner<T, H> {
    /// Create a new interner for given type.
    pub const fn new() -> StaticInterner<T, H> {
        StaticInterner {
            table: ShardedLockFreeRawTable::new(),
            _marker: marker::PhantomData,
        }
    }
}

impl<T: 'static, H: Hasher + Default> StaticInterner<T, H> {
    /// Allocate a value, or return previously allocated one.
    pub fn intern<Q>(&'static self, value: Q) -> Intern<T>
    where
        Q: Hash + Equiv<T> + Into<T>,
        T: Eq + Hash,
    {
        let hashed = Hashed::<_, H>::new(value);
        if let Some(pointer) = self
            .table
            .lookup(hashed.hash, |t| hashed.value.equivalent(&t.data))
        {
            return Intern { pointer };
        }

        self.intern_slow(hashed)
    }

    #[cold]
    fn intern_slow<Q>(&'static self, hashed_value: Hashed<Q, H>) -> Intern<T>
    where
        Q: Hash + Equiv<T> + Into<T>,
        T: Eq + Hash,
    {
        let pointer = Box::new(InternedData {
            data: hashed_value.value.into(),
            hash: hashed_value.hash,
        });
        let pointer = self
            .table
            .insert(
                hashed_value.hash,
                pointer,
                |a, b| a.hash == b.hash && a.data == b.data,
                |t| t.hash,
            )
            .0;
        Intern { pointer }
    }

    /// Get a value if it has been interned.
    pub fn get<Q>(&'static self, key: Q) -> Option<Intern<T>>
    where
        Q: Hash + Equiv<T>,
        T: Eq + Hash,
    {
        let hashed = Hashed::<_, H>::new(key);
        self.table
            .lookup(hashed.hash, |t| hashed.value.equivalent(&t.data))
            .map(|pointer| Intern { pointer })
    }

    /// Iterate over the interned values.
    #[inline]
    pub fn iter(&'static self) -> Iter<T, H> {
        Iter {
            iter: self.table.iter(),
            _marker: marker::PhantomData,
        }
    }
}

pub struct Iter<T: 'static, H: 'static> {
    iter: lock_free_hashtable::sharded::Iter<'static, Box<InternedData<T>>, 64>,
    _marker: marker::PhantomData<H>,
}

impl<T: 'static, H> Iterator for Iter<T, H> {
    type Item = Intern<T>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().map(|pointer| Intern { pointer })
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeSet;

    use crate::StaticInterner;

    static STRING_INTERNER: StaticInterner<String> = StaticInterner::new();

    #[test]
    fn test_intern() {
        assert_eq!(
            STRING_INTERNER.intern(&"hello".to_owned()),
            STRING_INTERNER.intern(&"hello".to_owned())
        );
        assert_eq!(
            STRING_INTERNER.intern(&"hello".to_owned()),
            STRING_INTERNER.intern("hello"),
        );
        assert_ne!(
            STRING_INTERNER.intern(&"hello".to_owned()),
            STRING_INTERNER.intern(&"world".to_owned())
        );
    }

    // Make sure things work with reallocation.
    #[test]
    fn test_resize() {
        let mut interned_strings = Vec::new();
        for i in 0..100000 {
            let s = i.to_string();
            let interned = STRING_INTERNER.intern(&s);
            assert_eq!(&s, &*interned);
            interned_strings.push(interned);
        }

        for s in &interned_strings {
            let interned = STRING_INTERNER.intern(&String::clone(s));
            assert_eq!(*s, interned);
        }
    }

    static TEST_GET_INTERNER: StaticInterner<String> = StaticInterner::new();
    #[test]
    fn test_get() {
        let interner = &TEST_GET_INTERNER;
        assert_eq!(interner.get("hello"), None);
        assert_eq!(interner.get(&"hello".to_owned()), None);

        let interned = interner.intern(&"hello".to_owned());
        assert_eq!(interner.get("hello"), Some(interned));
        assert_eq!(interner.get(&"hello".to_owned()), Some(interned));
        assert_eq!(interner.get("world"), None);
    }

    static TEST_ITER_INTERNER: StaticInterner<&'static str> = StaticInterner::new();
    #[test]
    fn test_iter() {
        let interner = &TEST_ITER_INTERNER;
        assert_eq!(
            interner
                .iter()
                .map(|v| *v)
                .collect::<BTreeSet<&'static str>>(),
            BTreeSet::from([])
        );
        interner.intern("hello");
        interner.intern("cat");
        interner.intern("world");

        assert_eq!(
            interner
                .iter()
                .map(|v| *v)
                .collect::<BTreeSet<&'static str>>(),
            BTreeSet::from(["hello", "cat", "world"])
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! Stress test for `LockFreeRawTable`.
//! It runs forever.

use std::collections::hash_map::DefaultHasher;
use std::hash::Hash;
use std::hash::Hasher;
use std::sync::atomic;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use std::sync::Condvar;
use std::sync::Mutex;
use std::thread;
use std::time::Duration;

use lock_free_hashtable::raw::LockFreeRawTable;

struct Barrier {
    rem: AtomicUsize,
}

impl Barrier {
    fn new(n: usize) -> Barrier {
        Barrier {
            rem: AtomicUsize::new(n),
        }
    }

    /// Wait for all threads to reach this point.
    fn wait(&self) {
        // We are not using any OS synchronization primitives here,
        // so we start all the work with as little delay as possible.
        let mut rem = self.rem.fetch_sub(1, atomic::Ordering::Relaxed) - 1;
        while rem != 0 {
            rem = self.rem.load(atomic::Ordering::Relaxed);
        }
    }
}

fn hash(key: u32) -> u64 {
    let mut hasher = DefaultHasher::new();
    key.hash(&mut hasher);
    hasher.finish()
}

fn hash_fn(key: &u32) -> u64 {
    hash(*key)
}

struct JobArg<'a> {
    /// Iteration number.
    i: u32,
    /// Thread number, 0..num_cpus.
    thread_i: u32,
    /// Number of threads.
    num_threads: u32,
    table: &'a LockFreeRawTable<Box<u32>>,
    /// Barrier can be used to synchronize threads in the middle of the job.
    barrier_0: &'a Barrier,
}

impl<'a> JobArg<'a> {
    /// Generate values unique to this job iteration among all threads.
    fn unique_values(&self, count: u32) -> impl Iterator<Item = u32> + '_ {
        (0..count).map(move |i| self.i + self.thread_i * count + i)
    }
}

fn run_stress_test(name: &str, job_count: u32, job_fn: impl Fn(JobArg) + Clone + Send + 'static) {
    println!("{name}:");

    let num_threads = num_cpus::get() as u32;

    struct JobData {
        /// Synchronize threads at the beginning of the job.
        start_barrier: Barrier,
        /// The table to be used by the job.
        table: LockFreeRawTable<Box<u32>>,
        /// Synchronize threads in the middle of the job.
        barrier_0: Barrier,
    }

    let jobs = (0..job_count)
        .map(|_| {
            let table = LockFreeRawTable::new();
            let start_barrier = Barrier::new(num_threads as usize);
            let barrier_0 = Barrier::new(num_threads as usize);
            Arc::new(JobData {
                table,
                start_barrier,
                barrier_0,
            })
        })
        .collect::<Vec<_>>();

    struct SharedState {
        /// Iteration.
        i: AtomicUsize,
        /// Threads still running.
        running_threads: Mutex<u32>,
        /// Wait for all threads to finish.
        condvar: Condvar,
    }

    let shared_state = Arc::new(SharedState {
        i: AtomicUsize::new(0),
        running_threads: Mutex::new(num_threads),
        condvar: Condvar::new(),
    });

    let mut threads = Vec::new();
    for thread_i in 0..num_threads {
        let jobs: Vec<Arc<JobData>> = jobs.clone();
        let shared_state = shared_state.clone();
        let job_fn = job_fn.clone();
        threads.push(thread::spawn(move || {
            for (i, job) in jobs.into_iter().enumerate() {
                if thread_i == 0 {
                    shared_state.i.store(i, atomic::Ordering::Relaxed);
                }
                job.start_barrier.wait();
                job_fn(JobArg {
                    thread_i,
                    i: i as u32,
                    num_threads,
                    table: &job.table,
                    barrier_0: &job.barrier_0,
                });
            }
            let mut rem_threads = shared_state.running_threads.lock().unwrap();
            *rem_threads = rem_threads.checked_sub(1).unwrap();
            if *rem_threads == 0 {
                shared_state.condvar.notify_all();
            }
        }));
    }

    // Release memory.
    drop(jobs);

    loop {
        let guard = shared_state.running_threads.lock().unwrap();
        if *guard == 0 {
            break;
        }
        let (_guard, _timeout) = shared_state
            .condvar
            .wait_timeout(guard, Duration::from_secs(1))
            .unwrap();
        println!("{}", shared_state.i.load(atomic::Ordering::Relaxed));
    }

    for thread in threads {
        thread.join().unwrap();
    }
}

/// Concurrently insert unique values into a table.
fn insert_unique(jobs: u32) {
    run_stress_test("insert_unique", jobs, |job| {
        for v in job.unique_values(1000) {
            job.table
                .insert(hash(v), Box::new(v), |a, b| a == b, hash_fn);
        }
        job.barrier_0.wait();
        for v in job.unique_values(1000) {
            assert!(job.table.lookup(hash(v), |a| *a == v).is_some());
        }
        assert_eq!(1000 * job.num_threads as usize, job.table.iter().count());
    });
}

/// Concurrently insert the same values into a table.
fn insert_same(jobs: u32) {
    run_stress_test("insert_same", jobs, |job| {
        for v in 0..10000 {
            let v = v + job.i;
            job.table
                .insert(hash(v), Box::new(v), |a, b| a == b, hash_fn);
        }
        job.barrier_0.wait();
        for v in 0..10000 {
            let v = v + job.i;
            assert!(job.table.lookup(hash(v), |a| *a == v).is_some());
        }
        assert_eq!(10000, job.table.iter().count());
    });
}

/// Concurrently insert unique values into a table and lookup them not long after.
fn insert_lookup(jobs: u32) {
    run_stress_test("insert_lookup", jobs, |job| {
        let mut prev = None;
        for v in job.unique_values(1000) {
            job.table
                .insert(hash(v), Box::new(v), |a, b| a == b, hash_fn);
            if let Some(prev) = prev {
                assert!(job.table.lookup(hash(prev), |a| *a == prev).is_some());
            }
            prev = Some(v);
        }

        job.barrier_0.wait();
        assert_eq!(1000 * job.num_threads as usize, job.table.iter().count());
    })
}

fn main() {
    println!("num_cpus: {}", num_cpus::get());

    let mut jobs = 100;
    loop {
        insert_unique(jobs);
        insert_same(jobs);
        insert_lookup(jobs);
        if jobs < 100000 {
            jobs *= 10;
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! Describes the types we can store in `LockFreeRawTable`.

use std::num::NonZeroU32;
use std::num::NonZeroU64;
use std::ptr;
use std::ptr::NonNull;
use std::sync::Arc;

/// Generalized non-null pointer.
///
/// Anything which can be stored in `LockFreeRawTable`.
pub trait AtomicValue {
    /// The value stored in the table.
    ///
    /// If the value is larger than max support atomic width on the platform,
    /// atomic operations will be spinlocked.
    type Raw: Copy;
    /// Dereferenced value.
    type Ref<'a>: Copy
    where
        Self: 'a;

    /// Null value. This is stored in `LockFreeRawTable` when the entry is missing.
    fn null() -> Self::Raw;
    /// Check if the value is null.
    fn is_null(this: Self::Raw) -> bool;
    /// Obtain the raw pointer. It must be non-null.
    fn into_raw(this: Self) -> Self::Raw;
    /// Construct the value from the raw pointer.
    unsafe fn from_raw(raw: Self::Raw) -> Self;
    /// Dereference the raw pointer.
    unsafe fn deref<'a>(raw: Self::Raw) -> Self::Ref<'a>;
}

impl<T> AtomicValue for Box<T> {
    type Raw = *mut T;
    type Ref<'a> = &'a T where Self: 'a;

    #[inline]
    fn null() -> Self::Raw {
        ptr::null_mut()
    }

    #[inline]
    fn is_null(this: Self::Raw) -> bool {
        this.is_null()
    }

    #[inline]
    fn into_raw(this: Self) -> Self::Raw {
        Box::into_raw(this)
    }

    #[inline]
    unsafe fn from_raw(raw: Self::Raw) -> Self {
        unsafe { Box::from_raw(raw) }
    }

    #[inline]
    unsafe fn deref<'a>(raw: Self::Raw) -> Self::Ref<'a> {
        unsafe { &*raw }
    }
}

impl<T> AtomicValue for Arc<T> {
    type Raw = *const T;
    type Ref<'a> = &'a T where Self: 'a;

    #[inline]
    fn null() -> Self::Raw {
        ptr::null()
    }

    #[inline]
    fn is_null(this: Self::Raw) -> bool {
        this.is_null()
    }

    #[inline]
    fn into_raw(this: Self) -> Self::Raw {
        Arc::into_raw(this)
    }

    #[inline]
    unsafe fn from_raw(raw: Self::Raw) -> Self {
        Arc::from_raw(raw)
    }

    #[inline]
    unsafe fn deref<'a>(raw: Self::Raw) -> Self::Ref<'a> {
        unsafe { &*raw }
    }
}

impl AtomicValue for NonZeroU64 {
    type Raw = u64;
    type Ref<'a> = NonZeroU64;

    #[inline]
    fn null() -> u64 {
        0
    }

    #[inline]
    fn is_null(this: u64) -> bool {
        this == 0
    }

    #[inline]
    fn into_raw(this: NonZeroU64) -> u64 {
        this.get()
    }

    #[inline]
    unsafe fn from_raw(raw: u64) -> NonZeroU64 {
        NonZeroU64::new_unchecked(raw)
    }

    #[inline]
    unsafe fn deref<'a>(raw: u64) -> Self::Ref<'a> {
        NonZeroU64::new_unchecked(raw)
    }
}

impl AtomicValue for NonZeroU32 {
    type Raw = u32;
    type Ref<'a> = NonZeroU32;

    #[inline]
    fn null() -> u32 {
        0
    }

    #[inline]
    fn is_null(this: u32) -> bool {
        this == 0
    }

    #[inline]
    fn into_raw(this: NonZeroU32) -> u32 {
        this.get()
    }

    #[inline]
    unsafe fn from_raw(raw: u32) -> NonZeroU32 {
        NonZeroU32::new_unchecked(raw)
    }

    #[inline]
    unsafe fn deref<'a>(raw: u32) -> Self::Ref<'a> {
        NonZeroU32::new_unchecked(raw)
    }
}

/// Raw pointers stored in the table.
#[derive(Copy, Clone)]
pub struct RawPtr<T>(pub NonNull<T>);

impl<T> AtomicValue for RawPtr<T> {
    type Raw = *mut T;
    type Ref<'a> = NonNull<T> where Self: 'a;

    #[inline]
    fn null() -> *mut T {
        ptr::null_mut()
    }

    #[inline]
    fn is_null(this: *mut T) -> bool {
        this.is_null()
    }

    #[inline]
    fn into_raw(this: RawPtr<T>) -> *mut T {
        this.0.as_ptr()
    }

    #[inline]
    unsafe fn from_raw(raw: *mut T) -> RawPtr<T> {
        RawPtr(NonNull::new_unchecked(raw))
    }

    #[inline]
    unsafe fn deref<'a>(raw: *mut T) -> Self::Ref<'a>
    where
        Self: 'a,
    {
        NonNull::new_unchecked(raw)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::mem;
use std::mem::ManuallyDrop;
use std::slice;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

use allocative::Allocative;
use allocative::Visitor;
use atomic::Atomic;

use crate::atomic_value::AtomicValue;

/// Fixed capacity hashtable.
/// When dropped, it will not drop the entries.
pub(crate) struct FixedCapTable<T: AtomicValue> {
    /// Current list of entries. Null if the entry is missing, otherwise a leaked `Box<T>`.
    entries: Box<[Atomic<T::Raw>]>,
    /// Number of entries in the table. We always use relaxed operations for this, meaning
    /// that any particular return value can always be wrong in either direction.
    size: AtomicUsize,
}

pub(crate) struct IterPtrs<'a, T: AtomicValue> {
    iter: slice::Iter<'a, Atomic<T::Raw>>,
}

impl<'a, T: AtomicValue> Iterator for IterPtrs<'a, T> {
    type Item = T::Raw;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.iter.next() {
                Some(entry) => {
                    let ptr = entry.load(Ordering::Acquire);
                    if !T::is_null(ptr) {
                        return Some(ptr);
                    }
                }
                None => return None,
            }
        }
    }
}

impl<T: AtomicValue + Allocative> FixedCapTable<T> {
    pub(crate) fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>, current: bool) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        {
            let mut visitor = visitor.enter_unique(
                allocative::Key::new("entries"),
                mem::size_of_val(&self.entries),
            );
            {
                let mut visitor = visitor.enter(
                    allocative::Key::new("capacity"),
                    mem::size_of_val::<[Atomic<T::Raw>]>(&*self.entries),
                );
                // We only visit pointers from current table.
                if current {
                    for ptr in self.iter_ptrs() {
                        let mut visitor =
                            visitor.enter(allocative::Key::new("entry"), mem::size_of::<T::Raw>());
                        unsafe {
                            // SAFETY: We "forget" the value after the iteration,
                            //   so it is safe to convert it to owned value.
                            let value = T::from_raw(ptr);
                            let value = ManuallyDrop::new(value);
                            if mem::size_of_val(&value) > mem::size_of::<T::Raw>() {
                                // Size will be computed incorrectly.
                            }

                            T::visit(&value, &mut visitor);
                        }
                        visitor.exit();
                    }
                }
                visitor.exit();
            }
            visitor.exit();
        }
        visitor.exit();
    }
}

/// Iterator over the entries.
pub struct Iter<'a, T: AtomicValue> {
    iter: IterPtrs<'a, T>,
}

impl<'a, T: AtomicValue> Iter<'a, T> {
    #[inline]
    pub(crate) fn empty() -> Iter<'a, T> {
        Iter {
            iter: IterPtrs { iter: [].iter() },
        }
    }
}

impl<'a, T: AtomicValue + 'a> Iterator for Iter<'a, T> {
    type Item = T::Ref<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().map(|ptr| unsafe { T::deref(ptr) })
    }
}

impl<T: AtomicValue> FixedCapTable<T> {
    pub(crate) fn with_capacity(cap: usize) -> FixedCapTable<T> {
        assert!(cap.is_power_of_two());
        let mut entries = Vec::new();
        entries.resize_with(cap, || Atomic::new(T::null()));
        FixedCapTable {
            entries: entries.into_boxed_slice(),
            size: AtomicUsize::new(0),
        }
    }

    #[inline]
    pub(crate) fn capacity(&self) -> usize {
        self.entries.len()
    }

    #[inline]
    pub(crate) fn need_resize(&self) -> bool {
        self.size.load(Ordering::Relaxed) >= self.entries.len() / 2
    }

    #[inline]
    pub(crate) fn lookup<'a>(
        &'a self,
        hash: u64,
        eq: impl Fn(T::Ref<'_>) -> bool,
    ) -> Option<T::Ref<'a>> {
        let mut index = hash as usize & (self.entries.len() - 1);
        for _ in 0..self.entries.len() {
            let entry = self.entries[index].load(Ordering::Acquire);
            if T::is_null(entry) {
                return None;
            }
            let entry = unsafe { T::deref(entry) };
            if eq(entry) {
                return Some(entry);
            }
            index = (index + 1) & (self.entries.len() - 1);
        }
        None
    }

    pub(crate) fn insert_unique_unchecked(&mut self, hash: u64, value: T::Raw) {
        let mut index = hash as usize & (self.entries.len() - 1);
        loop {
            let entry = self.entries[index].get_mut();
            if T::is_null(*entry) {
                *self.entries[index].get_mut() = value;
                *self.size.get_mut() += 1;
                return;
            }
            index = (index + 1) & (self.entries.len() - 1);
        }
    }

    /// Insert or lookup value.
    ///
    /// If the value is already in the table, it will be returned.
    /// If the capacity is full, the value will be returned as an error.
    /// Otherwise, the value will be inserted and returned.
    pub(crate) fn insert<'a>(
        &self,
        hash: u64,
        value: T,
        eq: impl Fn(T::Ref<'_>, T::Ref<'_>) -> bool,
    ) -> Result<(T::Ref<'a>, Option<T>), T> {
        let value = T::into_raw(value);
        assert!(!T::is_null(value));

        let mut index = hash as usize & (self.entries.len() - 1);

        // We do not know if the table is full already,
        // or becoming full concurrently with us.
        // So we have to stop after `capacity` iterations.
        for _ in 0..self.entries.len() {
            let entry = self.entries[index].load(Ordering::Acquire);
            let entry = if T::is_null(entry) {
                match self.entries[index].compare_exchange(
                    T::null(),
                    value,
                    // If the compare_exchange succeeds, the pointer has been inserted
                    // into the table. This means we need `Release` to ensure that the
                    // data behind the `Box<T>` is visible to other threads.
                    Ordering::Release,
                    // The code below reads the data behind the pointer, this means we
                    // need acquire so that this thread sees any previous changes to that
                    // data.
                    Ordering::Acquire,
                ) {
                    Ok(_) => {
                        self.size.fetch_add(1, Ordering::Relaxed);
                        return Ok((unsafe { T::deref(value) }, None));
                    }
                    Err(entry) => {
                        // Someone has just inserted the value into the bucket.
                        entry
                    }
                }
            } else {
                entry
            };
            if eq(unsafe { T::deref(value) }, unsafe { T::deref(entry) }) {
                return Ok(unsafe { (T::deref(entry), Some(T::from_raw(value))) });
            }
            index = (index + 1) & (self.entries.len() - 1);
        }
        Err(unsafe { T::from_raw(value) })
    }

    /// Iterate over the pointers to the entries.
    #[inline]
    pub(crate) fn iter_ptrs(&self) -> IterPtrs<'_, T> {
        IterPtrs {
            iter: self.entries.iter(),
        }
    }

    /// Iterate over the entries.
    #[inline]
    pub(crate) fn iter(&self) -> Iter<'_, T> {
        Iter {
            iter: self.iter_ptrs(),
        }
    }

    /// Drop all entries.
    pub(crate) unsafe fn drop_entries(&mut self) {
        for entry in self.iter_ptrs() {
            unsafe {
                let _drop = T::from_raw(entry);
            }
        }
    }

    #[inline]
    pub(crate) fn len(&self) -> usize {
        self.size.load(Ordering::Relaxed)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! (Almost) lock-free insertion only hashtable.
//!
//! Lookups are performed without locking (wait-free).
//! Insertions are performed with a shared lock (lock-free unless resize is needed).
//! Resizing is performed with an exclusive lock.
//!
//! Entries are never removed (until the table is dropped).

#![deny(missing_docs)]

pub mod atomic_value;
mod fixed_cap;
pub mod raw;
pub mod sharded;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! (Almost) lock-free insertion only hashtable.
//!
//! This module provides raw hashtable, which can be used to implement
//! higher-level hashtables.

use std::mem;
use std::ptr;
use std::sync::atomic::AtomicPtr;
use std::sync::atomic::Ordering;

use allocative::Allocative;
use allocative::Key;
use allocative::Visitor;
use parking_lot::RwLock;

use crate::atomic_value::AtomicValue;
use crate::fixed_cap::FixedCapTable;
pub use crate::fixed_cap::Iter;

struct CurrentTable<T: AtomicValue> {
    /// Previous tables are kept forever because there's no way to know
    /// when they are no longer in use.
    /// We double the capacity every time we resize,
    /// so previous tables allocation overhead is linear in the current capacity.
    _prev: Option<Box<CurrentTable<T>>>,
    /// Table used for insertions and lookups.
    /// When resizing/resized, this table is stored in the `prev` field,
    /// so all the pointers to the table remain valid.
    table: FixedCapTable<T>,
}

impl<T: AtomicValue + Allocative> CurrentTable<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>, current: bool) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        {
            let mut visitor =
                visitor.enter_unique(Key::new("_prev"), mem::size_of_val(&self._prev));
            if let Some(prev) = &self._prev {
                prev.visit(&mut visitor, false);
            }
            visitor.exit();
        }
        {
            let mut visitor =
                visitor.enter_unique(Key::new("table"), mem::size_of_val(&self.table));
            self.table.visit(&mut visitor, current);
            visitor.exit();
        }
        visitor.exit();
    }
}

/// (Almost) lock-free insertion only hashtable.
pub struct LockFreeRawTable<T: AtomicValue> {
    /// Shared lock for insertion, exclusive lock for resizing.
    /// No lock for lookup.
    write_lock: RwLock<()>,
    /// Current table. Readers can grab the pointer without locking
    /// and pointer remains valid even if the table is resized.
    current: AtomicPtr<CurrentTable<T>>,
}

impl<T: AtomicValue> Drop for LockFreeRawTable<T> {
    fn drop(&mut self) {
        let current = self.current.get_mut();
        if !current.is_null() {
            unsafe {
                let mut current = Box::from_raw(*current);
                // Only the current table owns the entries.
                // Previous tables store subset of the same entries, but we must not drop them.
                current.table.drop_entries();
            };
        }
    }
}

impl<T: AtomicValue> Default for LockFreeRawTable<T> {
    #[inline]
    fn default() -> LockFreeRawTable<T> {
        LockFreeRawTable::new()
    }
}

impl<T: AtomicValue> LockFreeRawTable<T> {
    /// Empty table.
    #[inline]
    pub const fn new() -> LockFreeRawTable<T> {
        LockFreeRawTable {
            write_lock: RwLock::new(()),
            current: AtomicPtr::new(ptr::null_mut()),
        }
    }

    /// Find an entry.
    #[inline]
    pub fn lookup<'a>(&'a self, hash: u64, eq: impl Fn(T::Ref<'_>) -> bool) -> Option<T::Ref<'a>> {
        let current = self.current.load(Ordering::Acquire);

        if current.is_null() {
            return None;
        }

        let current = unsafe { &*current };
        current.table.lookup(hash, eq)
    }

    /// Insert an entry.
    ///
    /// If the entry does not exist, the value is inserted
    /// and a pointer to the inserted entry is returned.
    ///
    /// Otherwise the pointer to existing entry along with the given value is returned.
    pub fn insert<'a>(
        &'a self,
        hash: u64,
        mut value: T,
        eq: impl Fn(T::Ref<'_>, T::Ref<'_>) -> bool,
        hash_fn: impl Fn(T::Ref<'_>) -> u64,
    ) -> (T::Ref<'a>, Option<T>) {
        loop {
            // Acquire shared lock.
            let guard = self.write_lock.read();

            let current = self.current.load(Ordering::Relaxed);

            if current.is_null() {
                // The table is just created. Allocate capacity and start over.
                drop(guard);
                self.resize_if_needed(|v| hash_fn(v));
                continue;
            }

            let current = unsafe { &*current };
            match current.table.insert(hash, value, |a, b| eq(a, b)) {
                Ok((referece, value)) => {
                    drop(guard);
                    // Insert was successful. However, we or other threads
                    // may have exceeded the load factor.
                    // So resize the table if needed to make lookup faster.
                    if current.table.need_resize() {
                        self.resize_if_needed(|v| hash_fn(v));
                    }
                    return (referece, value);
                }
                Err(ret_value) => {
                    drop(guard);
                    // Table is full. Resize the table and try again.
                    self.resize_if_needed(|v| hash_fn(v));
                    value = ret_value;
                }
            }
        }
    }

    #[cold]
    fn resize_if_needed(&self, hash: impl Fn(T::Ref<'_>) -> u64) {
        // Acquire exclusive lock.
        // Readers still read the old table, but new insertions will wait.
        let _guard = self.write_lock.write();

        let current_ptr = self.current.load(Ordering::Relaxed);

        if current_ptr.is_null() {
            let new_table: FixedCapTable<T> = FixedCapTable::with_capacity(16);
            let new_current = Box::new(CurrentTable {
                _prev: None,
                table: new_table,
            });
            self.current
                .store(Box::into_raw(new_current), Ordering::Release);
            return;
        }

        let current = unsafe { &*current_ptr };

        if !current.table.need_resize() {
            return;
        }

        let new_cap = current.table.capacity().checked_mul(2).unwrap();
        let mut new_table: FixedCapTable<T> = FixedCapTable::with_capacity(new_cap);

        for entry_ptr in current.table.iter_ptrs() {
            let entry = unsafe { T::deref(entry_ptr) };
            let hash = hash(entry);
            new_table.insert_unique_unchecked(hash, entry_ptr);
        }

        let new_current = Box::new(CurrentTable {
            _prev: Some(unsafe { Box::from_raw(current_ptr) }),
            table: new_table,
        });

        self.current
            .store(Box::into_raw(new_current), Ordering::Release);
    }

    /// Iterate over all entries.
    pub fn iter(&self) -> Iter<'_, T> {
        let current = self.current.load(Ordering::Acquire);

        if current.is_null() {
            return Iter::empty();
        }

        let current = unsafe { &*current };
        current.table.iter()
    }

    /// Number of entries in the table.
    #[inline]
    pub fn len(&self) -> usize {
        let current = self.current.load(Ordering::Acquire);

        if current.is_null() {
            return 0;
        }

        let current = unsafe { &*current };
        current.table.len()
    }

    /// Number of entries in the table is zero.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl<T: AtomicValue + Allocative> Allocative for LockFreeRawTable<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        {
            let mut visitor = visitor.enter_unique(
                allocative::Key::new("current"),
                mem::size_of_val(&self.current),
            );
            let current = self.current.load(Ordering::Acquire);
            if !current.is_null() {
                let current = unsafe { &*current };
                current.visit(&mut visitor, true);
            }
            visitor.exit();
        }
        visitor.exit();
    }
}

#[cfg(test)]
mod tests {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::Hash;
    use std::hash::Hasher;
    use std::num::NonZeroU32;
    use std::ptr;
    use std::ptr::NonNull;

    use crate::atomic_value::RawPtr;
    use crate::raw::LockFreeRawTable;

    fn hash<T: Hash>(key: T) -> u64 {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish()
    }

    #[allow(clippy::trivially_copy_pass_by_ref)]
    fn hash_fn(key: &u32) -> u64 {
        hash(*key)
    }

    #[test]
    fn test_simple() {
        let t = LockFreeRawTable::new();
        let v0 = t.insert(hash(1), Box::new(1), |a, b| a == b, hash_fn).0;
        assert_eq!(&1, v0);
        let v1 = t.insert(hash(1), Box::new(1), |a, b| a == b, hash_fn).0;
        assert_eq!(&1, v1);
        assert!(ptr::eq(v0, v1));

        let v2 = t.lookup(hash(1), |a| a == &1).unwrap();
        assert_eq!(&1, v2);
        assert!(ptr::eq(v0, v2));

        assert_eq!(vec![1], t.iter().cloned().collect::<Vec<_>>());
    }

    #[test]
    fn test_million() {
        let t = LockFreeRawTable::new();
        for i in 0..1000000 {
            t.insert(hash(i), Box::new(i), |a, b| a == b, hash_fn);
        }

        for i in 0..1000000 {
            let v = t.lookup(hash(i), |a| a == &i).unwrap();
            assert_eq!(&i, v);
        }

        assert_eq!(1000000, t.iter().count());
    }

    #[test]
    fn test_u32() {
        let t = LockFreeRawTable::new();
        for i in 1..10000 {
            t.insert(hash(i), NonZeroU32::new(i).unwrap(), |a, b| a == b, hash);
        }

        for i in 1..10000 {
            let v = t.lookup(hash(i), |a| a.get() == i).unwrap();
            assert_eq!(NonZeroU32::new(i).unwrap(), v);
        }
    }

    #[test]
    fn test_raw_pointers() {
        let t = LockFreeRawTable::new();
        let p = NonNull::<u32>::dangling();
        let p = RawPtr(p);
        let r = t.insert(hash(p.0), p, |a, b| a == b, hash).0;
        assert_eq!(p.0, r);
    }

    #[test]
    fn test_allocative() {
        let table = LockFreeRawTable::<Box<u16>>::new();

        for i in 0..100 {
            let (_r, inserted) = table.insert(hash(i), Box::new(i), |a, b| a == b, |v| hash(v));
            assert!(inserted.is_none());
        }

        let mut builder = allocative::FlameGraphBuilder::default();
        builder.visit_root(&table);
        let _flame_graph = builder.finish_and_write_flame_graph();
        // Proper automated test is possible but hard to maintain.
        // The best option is to print the flame graph and check resulting image.
        // println!("{}", _flame_graph);
        // At the moment of writing it looks like this: https://www.internalfb.com/intern/px/p/2Gb33

        // Alternatively, we can set up golden test.
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! Sharded (almost) lock-free hashtable.

use std::marker;

use allocative::Allocative;

use crate::atomic_value::AtomicValue;
use crate::raw::LockFreeRawTable;

/// Lock-free hashtable sharded by key hash.
#[derive(Allocative)]
pub struct ShardedLockFreeRawTable<T: AtomicValue, const SHARDS: usize> {
    shards: [LockFreeRawTable<T>; SHARDS],
}

impl<T: AtomicValue, const SHARDS: usize> ShardedLockFreeRawTable<T, SHARDS> {
    const _ASSERTIONS: () = assert!(SHARDS.is_power_of_two());

    const SHARD_BITS: usize = SHARDS.trailing_zeros() as usize;

    /// Create a new empty hashtable.
    pub const fn new() -> ShardedLockFreeRawTable<T, SHARDS> {
        struct Empty<A>(marker::PhantomData<A>);
        impl<A: AtomicValue> Empty<A> {
            #[allow(clippy::declare_interior_mutable_const)]
            const EMPTY: LockFreeRawTable<A> = LockFreeRawTable::new();
        }

        ShardedLockFreeRawTable {
            shards: [Empty::<T>::EMPTY; SHARDS],
        }
    }

    #[inline]
    fn table_for_hash(&self, hash: u64) -> &LockFreeRawTable<T> {
        // `LockFreeRawTable` uses low bits of hash, so we use high bits to select a shard.
        let shard_index = (hash >> (64 - Self::SHARD_BITS)) as usize;
        &self.shards[shard_index]
    }

    /// Find an entry.
    #[inline]
    pub fn lookup<'a>(&'a self, hash: u64, eq: impl Fn(T::Ref<'_>) -> bool) -> Option<T::Ref<'_>> {
        self.table_for_hash(hash).lookup(hash, eq)
    }

    /// Insert an entry.
    ///
    /// If the entry does not exist, the value is inserted
    /// and a pointer to the inserted entry is returned.
    ///
    /// Otherwise the pointer to existing entry along with the given value is returned.
    #[inline]
    pub fn insert<'a>(
        &'a self,
        hash: u64,
        value: T,
        eq: impl Fn(T::Ref<'_>, T::Ref<'_>) -> bool,
        hash_fn: impl Fn(T::Ref<'_>) -> u64,
    ) -> (T::Ref<'a>, Option<T>) {
        self.table_for_hash(hash).insert(hash, value, eq, hash_fn)
    }

    /// Iterate entries in unspecified order.
    #[inline]
    pub fn iter(&self) -> Iter<T, SHARDS> {
        Iter {
            table: self,
            shard: 0,
            iter: self.shards[0].iter(),
        }
    }

    /// Number of entries in the table.
    #[inline]
    pub fn len(&self) -> usize {
        self.shards.iter().map(|s| s.len()).sum()
    }

    /// Number of entries in the table is zero.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// Iterator over all entries in sharded raw table.
pub struct Iter<'a, T: AtomicValue, const SHARDS: usize> {
    table: &'a ShardedLockFreeRawTable<T, SHARDS>,
    /// Current iterator shard index.
    shard: usize,
    iter: crate::raw::Iter<'a, T>,
}

impl<'a, T: AtomicValue + 'a, const SHARDS: usize> Iterator for Iter<'a, T, SHARDS> {
    type Item = T::Ref<'a>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(next) = self.iter.next() {
                return Some(next);
            }
            if self.shard >= SHARDS - 1 {
                debug_assert!(self.shard == SHARDS - 1);
                return None;
            }
            self.shard += 1;
            self.iter = self.table.shards[self.shard].iter();
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::hash_map::DefaultHasher;
    use std::collections::HashSet;
    use std::hash::Hash;
    use std::hash::Hasher;

    use crate::sharded::ShardedLockFreeRawTable;

    #[test]
    fn test_shard_bits() {
        assert_eq!(0, ShardedLockFreeRawTable::<Box<u32>, 1>::SHARD_BITS);
        assert_eq!(1, ShardedLockFreeRawTable::<Box<u32>, 2>::SHARD_BITS);
        assert_eq!(2, ShardedLockFreeRawTable::<Box<u32>, 4>::SHARD_BITS);
        assert_eq!(3, ShardedLockFreeRawTable::<Box<u32>, 8>::SHARD_BITS);
        assert_eq!(4, ShardedLockFreeRawTable::<Box<u32>, 16>::SHARD_BITS);
    }

    fn hash(key: u32) -> u64 {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish()
    }

    #[allow(clippy::trivially_copy_pass_by_ref)]
    fn hash_fn(key: &u32) -> u64 {
        hash(*key)
    }

    #[test]
    fn test_iter() {
        let table = ShardedLockFreeRawTable::<Box<u32>, 8>::new();
        let mut expected = Vec::new();
        for i in 0..1000 {
            table.insert(hash(i), Box::new(i), |a, b| a == b, hash_fn);
            expected.push(i);
        }

        let mut collect = HashSet::new();
        for i in table.iter() {
            let inserted = collect.insert(*i);
            assert!(inserted);
        }

        let mut collect = Vec::from_iter(collect);
        collect.sort_unstable();
        assert_eq!(expected, collect);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! Lock-free vector-like data structure.

#![deny(missing_docs)]

use std::array;
use std::cell::UnsafeCell;
use std::cmp;
use std::mem;
use std::mem::MaybeUninit;
use std::ptr;
use std::slice;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

use allocative::Allocative;
use allocative::Visitor;

/// Compute the number of buckets needed for a given capacity.
/// The number of bucket is `log2(capacity) - 3`.
#[inline]
pub const fn buckets_for_max_capacity(max_capacity: usize) -> usize {
    assert!(max_capacity.is_power_of_two());
    assert!(max_capacity >= 16);
    max_capacity.trailing_zeros() as usize - (FIRST_BUCKET_CAP_EXP as usize - 1)
}

/// `Box<[MaybeUninit<T>]>`.
struct BucketAlloc<T> {
    data: *mut T,
    cap: usize,
}

impl<T> BucketAlloc<T> {
    #[inline]
    fn new(cap: usize) -> BucketAlloc<T> {
        debug_assert!(cap.is_power_of_two());
        let mut data = Vec::with_capacity(cap);
        data.resize_with(cap, MaybeUninit::uninit);
        let data = data.into_boxed_slice();
        let data: *mut [MaybeUninit<T>] = Box::into_raw(data);
        let data: *mut T = data as *mut T;
        BucketAlloc { data, cap }
    }

    #[inline]
    unsafe fn drop_data(&self, len: usize) {
        assert!(len <= self.cap);
        let data: *mut [T] = unsafe { slice::from_raw_parts_mut(self.data, len) };
        ptr::drop_in_place(data)
    }

    #[inline]
    unsafe fn from_raw_parts(data: *mut T, cap: usize) -> BucketAlloc<T> {
        debug_assert!(cap.is_power_of_two());
        BucketAlloc { data, cap }
    }
}

/// Deallocate.
impl<T> Drop for BucketAlloc<T> {
    fn drop(&mut self) {
        unsafe {
            let data: *mut MaybeUninit<T> = self.data as *mut MaybeUninit<T>;
            let data: *mut [MaybeUninit<T>] = slice::from_raw_parts_mut(data, self.cap);
            let data: Box<[MaybeUninit<T>]> = Box::from_raw(data);
            drop(data);
        }
    }
}

const LOCKED_FLAG: usize = 1 << (mem::size_of::<usize>() * 8 - 1);
/// `log2` of first bucket capacity.
const FIRST_BUCKET_CAP_EXP: u32 = 4;
/// 16.
const FIRST_BUCKET_CAP: usize = 1 << FIRST_BUCKET_CAP_EXP;

/// Maximum number of buckets.
pub const MAX_BUCKETS: usize = mem::size_of::<usize>() * 8 - FIRST_BUCKET_CAP_EXP as usize;

/// Data structure with push and get operations.
///
/// * `get` is wait-free and no locks.
/// * `push` is single-threaded.
///
/// Both operations are `O(1)`, hence the `Vec` part in name.
/// But data is not stored in contiguous memory.
// This can be parameterized by capacity when `feature(generic_const_exprs)` is stable.
// https://github.com/rust-lang/rust/issues/76560
pub struct LockFreeVec<T, const BUCKETS: usize = MAX_BUCKETS> {
    /// Buckets containing the actual data.
    /// ```ignore
    /// [0..15] // first is special
    /// [16..31]
    /// [32..63]
    /// [64..127]
    /// [128..255]
    /// [256..511]
    /// ...
    /// ```
    /// The number of initialized fields in this array is given by the size. If, for
    /// example, the size is 100 then the first four fields of this array are
    /// initialized. Those fields can then also be read without need for any
    /// synchronization.
    buckets: [UnsafeCell<*mut T>; BUCKETS],
    /// Size, and high bit is locked flag.
    /// Operations are synchronized on this field.
    size: AtomicUsize,
}

impl<T, const BUCKETS: usize> Default for LockFreeVec<T, BUCKETS> {
    #[inline]
    fn default() -> LockFreeVec<T, BUCKETS> {
        LockFreeVec::new()
    }
}

unsafe impl<T: Sync + Send, const BUCKETS: usize> Send for LockFreeVec<T, BUCKETS> {}
unsafe impl<T: Sync + Send, const BUCKETS: usize> Sync for LockFreeVec<T, BUCKETS> {}

impl<T, const BUCKETS: usize> LockFreeVec<T, BUCKETS> {
    const _ASSERTS: () = {
        assert!(BUCKETS > 0);
        assert!(BUCKETS <= MAX_BUCKETS);
        assert!(Self::max_capacity() < LOCKED_FLAG);
    };

    /// Empty.
    // This can be `const fn` when something is stabilized as const, for example:
    // * `MaybeUninit::zeroed`
    // * `[const {expr}; 10 ]`
    #[inline]
    pub fn new() -> LockFreeVec<T, BUCKETS> {
        LockFreeVec {
            buckets: array::from_fn(|_| UnsafeCell::new(ptr::null_mut())),
            size: AtomicUsize::new(0),
        }
    }

    /// Maximum capacity.
    #[inline]
    pub const fn max_capacity() -> usize {
        // Does not overflow because static assertions.
        1 << (BUCKETS + FIRST_BUCKET_CAP_EXP as usize - 1)
    }

    /// Capacity of a bucket given its index.
    #[inline]
    fn bucket_capacity(bucket: usize) -> usize {
        debug_assert!(bucket < BUCKETS);
        cmp::max(
            FIRST_BUCKET_CAP,
            1 << (bucket + FIRST_BUCKET_CAP_EXP as usize - 1),
        )
    }

    /// Unpack index into bucket index and offset in bucket.
    #[inline]
    fn unpack_index(index: usize) -> (usize, usize) {
        debug_assert!(index < Self::max_capacity());
        let bucket =
            mem::size_of::<usize>() * 8 - (index >> FIRST_BUCKET_CAP_EXP).leading_zeros() as usize;
        let index_in_bucket = index & (Self::bucket_capacity(bucket) - 1);
        (bucket, index_in_bucket)
    }

    /// Number of elements.
    ///
    /// This operation has acquire semantics.
    #[inline]
    pub fn len(&self) -> usize {
        // We want acquire because the returned value is to be used in `push_at`.
        self.size.load(Ordering::Acquire) & !LOCKED_FLAG
    }

    /// Is empty?
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Get element.
    ///
    /// Wait-free, no locks.
    #[inline]
    pub fn get(&self, index: usize) -> Option<&T> {
        let size = self.len();
        if index >= size {
            return None;
        }
        let (bucket, index_in_bucket) = Self::unpack_index(index);
        let bucket: *mut T = unsafe { *self.buckets[bucket].get() };
        Some(unsafe { &*bucket.add(index_in_bucket) })
    }

    #[inline]
    unsafe fn init_bucket(&self, bucket: usize) -> *mut T {
        let bucket_ptr_ptr: *mut *mut T = self.buckets[bucket].get();
        let bucket_ptr = *bucket_ptr_ptr;
        if !bucket_ptr.is_null() {
            return bucket_ptr;
        }
        self.init_bucket_slow(bucket)
    }

    #[cold]
    unsafe fn init_bucket_slow(&self, bucket: usize) -> *mut T {
        let bucket_ptr_ptr: *mut *mut T = self.buckets[bucket].get();
        assert!((*bucket_ptr_ptr).is_null());
        let bucket_alloc = BucketAlloc::new(Self::bucket_capacity(bucket));
        let bucket_ptr = bucket_alloc.data;
        *bucket_ptr_ptr = bucket_ptr;
        mem::forget(bucket_alloc);
        bucket_ptr
    }

    /// Push an element to the end.
    ///
    /// This operation returns error if index is not equal to the current size.
    ///
    /// Wait-free.
    #[inline]
    pub fn push_at(&self, index: usize, value: T) -> Result<(), T> {
        if index == Self::max_capacity() {
            panic!("capacity overflow: {}", Self::max_capacity());
        }

        if self
            .size
            .compare_exchange(
                index,
                index | LOCKED_FLAG,
                Ordering::Acquire,
                Ordering::Relaxed,
            )
            .is_err()
        {
            // We could retry, but we promised not to lock.
            return Err(value);
        }

        let (bucket, index_in_bucket) = Self::unpack_index(index);

        unsafe {
            let bucket_ptr = self.init_bucket(bucket);
            ptr::write(bucket_ptr.add(index_in_bucket), value);
        }

        self.size.store(index + 1, Ordering::Release);
        Ok(())
    }

    /// Iterate over all elements.
    ///
    /// Iterates over snapshot of the vector, newly pushed elements may not be visible.
    #[inline]
    pub fn iter(&self) -> Iter<T, BUCKETS> {
        let len = self.len();
        let slice = if len == 0 {
            &[]
        } else {
            unsafe {
                slice::from_raw_parts(*self.buckets[0].get(), cmp::min(len, FIRST_BUCKET_CAP))
            }
        };
        Iter {
            vec: self,
            bucket: 0,
            slice: slice.iter(),
            rem: len - slice.len(),
        }
    }
}

impl<T, const BUCKETS: usize> Drop for LockFreeVec<T, BUCKETS> {
    fn drop(&mut self) {
        unsafe {
            let mut rem_len = *self.size.get_mut();
            for (bucket, bucket_ptr_ptr) in self.buckets.iter().enumerate() {
                let bucket_ptr: *mut T = *bucket_ptr_ptr.get();
                if bucket_ptr.is_null() {
                    assert_eq!(0, rem_len);
                    break;
                }

                let bucket_len = cmp::min(rem_len, Self::bucket_capacity(bucket));
                rem_len -= bucket_len;

                let bucket_alloc: BucketAlloc<T> =
                    BucketAlloc::from_raw_parts(bucket_ptr, Self::bucket_capacity(bucket));

                bucket_alloc.drop_data(bucket_len);

                drop(bucket_alloc);
            }

            assert_eq!(0, rem_len);
        }
    }
}

impl<T: Allocative, const SHARDS: usize> Allocative for LockFreeVec<T, SHARDS> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        let mut rem_len = self.len();
        for (bucket, bucket_ptr_ptr) in self.buckets.iter().enumerate() {
            if rem_len == 0 {
                break;
            }

            let bucket_ptr: *mut T = unsafe { *bucket_ptr_ptr.get() };
            if bucket_ptr.is_null() {
                break;
            }

            let bucket_len = cmp::min(rem_len, Self::bucket_capacity(bucket));
            rem_len -= bucket_len;

            let mut visitor = visitor.enter_unique(
                allocative::Key::new("bucket"),
                mem::size_of_val(&bucket_ptr),
            );
            {
                let mut visitor = visitor.enter(
                    allocative::Key::new("capacity"),
                    mem::size_of::<T>() * Self::bucket_capacity(bucket),
                );
                for i in 0..bucket_len {
                    let value: &T = unsafe { &*bucket_ptr.add(i) };
                    visitor.visit_field(allocative::Key::new("item"), value);
                }
                visitor.exit();
            }

            visitor.exit();
        }
        visitor.exit();
    }
}

/// Iterator over all elements.
pub struct Iter<'a, T, const SHARDS: usize> {
    vec: &'a LockFreeVec<T, SHARDS>,
    /// Current bucket index.
    bucket: usize,
    /// Remaining elements in the current bucket.
    slice: slice::Iter<'a, T>,
    /// Remaining elements after the current bucket iterator.
    rem: usize,
}

impl<'a, T, const SHARDS: usize> Iterator for Iter<'a, T, SHARDS> {
    type Item = &'a T;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        match self.slice.next() {
            Some(value) => Some(value),
            None if self.rem == 0 => None,
            None => {
                self.bucket += 1;
                let bucket_len = cmp::min(
                    self.rem,
                    LockFreeVec::<T, SHARDS>::bucket_capacity(self.bucket),
                );
                let slice = unsafe {
                    slice::from_raw_parts(*self.vec.buckets[self.bucket].get(), bucket_len)
                };
                self.rem -= slice.len();
                self.slice = slice.iter();
                self.slice.next()
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let rem = self.rem + self.slice.len();
        (rem, Some(rem))
    }
}

impl<'a, T, const SHARDS: usize> ExactSizeIterator for Iter<'a, T, SHARDS> {}

impl<'a, T, const SHARDS: usize> IntoIterator for &'a LockFreeVec<T, SHARDS> {
    type Item = &'a T;
    type IntoIter = Iter<'a, T, SHARDS>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::Mutex;

    use allocative::FlameGraphBuilder;

    use crate::buckets_for_max_capacity;
    use crate::LockFreeVec;

    #[test]
    fn test_max_capacity() {
        assert_eq!(16, LockFreeVec::<u8, 1>::max_capacity());
        assert_eq!(32, LockFreeVec::<u8, 2>::max_capacity());
        assert_eq!(64, LockFreeVec::<u8, 3>::max_capacity());
        assert_eq!(128, LockFreeVec::<u8, 4>::max_capacity());
    }

    #[test]
    fn test_buckets_for_max_capacity() {
        assert_eq!(1, buckets_for_max_capacity(16));
        assert_eq!(2, buckets_for_max_capacity(32));
        assert_eq!(3, buckets_for_max_capacity(64));
        assert_eq!(4, buckets_for_max_capacity(128));
    }

    #[test]
    fn test_bucket_capacity() {
        assert_eq!(16, LockFreeVec::<u8, 1>::bucket_capacity(0));
        assert_eq!(16, LockFreeVec::<u8, 2>::bucket_capacity(1));
        assert_eq!(32, LockFreeVec::<u8, 3>::bucket_capacity(2));
        assert_eq!(64, LockFreeVec::<u8, 4>::bucket_capacity(3));
    }

    #[test]
    fn test_unpack_index() {
        // Special
        assert_eq!((0, 0), LockFreeVec::<u8, 10>::unpack_index(0));
        assert_eq!((0, 1), LockFreeVec::<u8, 10>::unpack_index(1));
        assert_eq!((0, 2), LockFreeVec::<u8, 10>::unpack_index(2));
        assert_eq!((0, 14), LockFreeVec::<u8, 10>::unpack_index(14));
        assert_eq!((0, 15), LockFreeVec::<u8, 10>::unpack_index(15));
        // Regular
        assert_eq!((1, 0), LockFreeVec::<u8, 10>::unpack_index(16));
        assert_eq!((1, 15), LockFreeVec::<u8, 10>::unpack_index(31));
        assert_eq!((2, 0), LockFreeVec::<u8, 10>::unpack_index(32));
        assert_eq!((2, 31), LockFreeVec::<u8, 10>::unpack_index(63));
        assert_eq!((3, 0), LockFreeVec::<u8, 10>::unpack_index(64));
        assert_eq!((3, 63), LockFreeVec::<u8, 10>::unpack_index(127));
        assert_eq!((4, 0), LockFreeVec::<u8, 10>::unpack_index(128));
        assert_eq!((4, 127), LockFreeVec::<u8, 10>::unpack_index(255));
        assert_eq!((5, 0), LockFreeVec::<u8, 10>::unpack_index(256));
    }

    #[test]
    fn test_simple() {
        let vec = LockFreeVec::<u32, 10>::new();
        for i in 0..100 {
            assert_eq!(i as usize, vec.len());
            vec.push_at(i as usize, i).unwrap();
        }
        for i in 0..100 {
            assert_eq!(i, *vec.get(i as usize).unwrap());
        }
    }

    #[test]
    fn test_full() {
        type V = LockFreeVec<u32, 4>;
        let vec = V::new();
        for i in 0..V::max_capacity() {
            vec.push_at(i, i as u32).unwrap();
        }
        assert_eq!(V::max_capacity(), vec.len());
    }

    #[test]
    fn test_max_buckets() {
        let v = LockFreeVec::<u32>::new();
        for i in 0..100 {
            assert_eq!(i as usize, v.len());
            v.push_at(i as usize, i).unwrap();
        }
        assert_eq!(
            (0..100).collect::<Vec<_>>(),
            v.iter().copied().collect::<Vec<_>>()
        );
    }

    #[test]
    fn test_million() {
        let vec = LockFreeVec::<String, 20>::new();
        for i in 0..1_000_000 {
            assert_eq!(i, vec.len());
            vec.push_at(i, i.to_string()).unwrap();
        }
    }

    #[test]
    fn test_drop() {
        struct RecordDrops(u32, Arc<Mutex<Vec<u32>>>);
        impl Drop for RecordDrops {
            fn drop(&mut self) {
                self.1.lock().unwrap().push(self.0);
            }
        }

        let drops = Arc::new(Mutex::new(Vec::new()));
        let vec = LockFreeVec::<RecordDrops, 20>::new();

        for i in 0..1000 {
            vec.push_at(i as usize, RecordDrops(i, drops.clone()))
                .ok()
                .unwrap();
        }

        drop(vec);

        let drops = Arc::try_unwrap(drops).unwrap().into_inner().unwrap();
        assert_eq!((0..1000).collect::<Vec<_>>(), drops);
    }

    #[test]
    fn test_iter_empty() {
        let v = LockFreeVec::<u32, 10>::new();
        assert!(v.iter().next().is_none());
    }

    #[test]
    fn test_iter_max() {
        type V = LockFreeVec<u32, 10>;
        let vec = V::new();
        for i in 0..V::max_capacity() {
            vec.push_at(i, i as u32).unwrap();
        }
        let mut iter = vec.iter().enumerate();
        assert_eq!(vec.len(), iter.len());
        while let Some((i, v)) = iter.next() {
            assert_eq!(vec.len() - i - 1, iter.len());
            assert_eq!(i, *v as usize);
        }
        assert_eq!(0, iter.len());
    }

    #[test]
    fn test_iter_partial() {
        type V = LockFreeVec<u32, 10>;
        let vec = V::new();
        for i in 0..100 {
            vec.push_at(i, i as u32).unwrap();
        }
        let mut iter = vec.iter().enumerate();
        assert_eq!(100, iter.len());
        while let Some((i, v)) = iter.next() {
            // Pushing elements does not affect the iterator.
            vec.push_at(vec.len(), 12345).unwrap();

            assert_eq!(99 - i, iter.len());
            assert_eq!(i, *v as usize);
        }
        assert_eq!(0, iter.len());
    }

    #[test]
    fn test_allocative() {
        let v = LockFreeVec::<String, 10>::new();
        for i in 0..100 {
            v.push_at(v.len(), format!("{:<3}", i)).unwrap();
        }

        let mut builder = FlameGraphBuilder::default();
        builder.visit_root(&v);
        let flame_graph = builder.finish_and_write_flame_graph();
        // Proper test is possible but expensive to maintain.
        // So print it out and eyeball it.
        // At the moment of writing it looks like this:
        // https://www.internalfb.com/intern/px/p/2GgX5
        // Alternatively we can set up a golden test.
        if false {
            println!("{}", flame_graph);
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::cell::RefCell;
use std::pin::Pin;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::sync::Weak;

use dupe::Clone_;
use dupe::Dupe;
use dupe::Dupe_;
use futures::future::Future;
use futures::future::FutureExt;
use futures::future::Shared;
use futures::task::Context;
use futures::task::Poll;
use futures::task::Waker;
use parking_lot::Mutex;
use pin_project::pin_project;
use tokio::sync::oneshot;

thread_local! {
    /// The ExecutionContext for the currently executing CancellableFuture.
    static CURRENT: RefCell<Option<Box<ExecutionContext>>> = RefCell::new(None);
}

enum State {
    /// This future has been constructed, but not polled yet.
    Pending,

    /// This future has been polled. A waker is available.
    Polled { waker: Waker },

    /// This future has already been cancelled.
    Cancelled,

    /// This future has already finished executing.
    Exited,
}

struct SharedStateData {
    state: Mutex<State>,

    /// When set, this future has been cancelled and should attempt to exit as soon as possible.
    cancelled: AtomicBool,
}

#[derive(Clone_, Dupe_)]
struct SharedState {
    inner: Arc<SharedStateData>,
}

impl SharedState {
    pub fn cancel(&self) {
        self.inner.cancelled.store(true, Ordering::SeqCst);

        let future = std::mem::replace(&mut *self.inner.state.lock(), State::Cancelled);

        match future {
            State::Pending => {
                // When the future starts, it'll see its cancellation;
            }
            State::Polled { waker } => {
                waker.wake();
            }
            State::Cancelled => {
                // We were already cancelled, no need to so again.
            }
            State::Exited => {
                // Nothing to do, that future is done.
            }
        }
    }
}

struct ExecutionContext {
    ref_count: WeakRefCount,
    shared: Arc<Mutex<ExecutionContextShared>>,
}

impl ExecutionContext {
    fn new(ref_count: WeakRefCount) -> Self {
        Self {
            ref_count,
            shared: Arc::new(Mutex::new(ExecutionContextShared {
                critical_section: None,
                structured_cancellation: None,
            })),
        }
    }

    fn enter_critical_section(&mut self) {
        let mut shared = self.shared.lock();

        let critical_section = shared
            .critical_section
            .get_or_insert_with(|| CriticalSection {
                observers: 0,
                _guard: self.ref_count.upgrade(),
            });

        critical_section.observers += 1;
    }

    fn enter_structured_cancellation(&mut self) -> CancellationObserver {
        let mut shared = self.shared.lock();

        let cancellation = shared.structured_cancellation.get_or_insert_with(|| {
            let (tx, rx) = oneshot::channel();
            StructuredCancellation {
                observers: 0,
                tx: Some(tx),
                rx: rx.shared(),
            }
        });

        cancellation.observers += 1;
        CancellationObserver {
            rx: Some(cancellation.rx.clone()),
        }
    }
}

/// In theory, this is not shared, but we can't *prevent* it statically, so we slap this behind an
/// Arc<Mutex<..>>. We only lock this on poll() if we've been cancelled, so in practice it's not
/// all that expensive.
struct ExecutionContextShared {
    critical_section: Option<CriticalSection>,
    structured_cancellation: Option<StructuredCancellation>,
}

impl ExecutionContextShared {
    /// Does this future not currently prevent its cancellation?
    fn can_exit(&self) -> bool {
        self.critical_section.is_none() && self.structured_cancellation.is_none()
    }

    fn notify_cancelled(&mut self) {
        if let Some(structured_cancellation) = self.structured_cancellation.as_mut() {
            if let Some(tx) = structured_cancellation.tx.take() {
                let _ignored = tx.send(());
            }
        }
    }

    fn exit_critical_section(&mut self) {
        let observers = {
            let critical_section = self
                .critical_section
                .as_mut()
                .expect("Cannot have more exits than enters");
            critical_section.observers -= 1;
            critical_section.observers
        };

        if observers == 0 {
            self.critical_section.take();
        }
    }

    fn exit_structured_cancellation(&mut self) {
        let observers = {
            let cancellation = self
                .structured_cancellation
                .as_mut()
                .expect("Cannot have more exits than enters");
            cancellation.observers -= 1;
            cancellation.observers
        };

        if observers == 0 {
            self.structured_cancellation.take();
        }
    }
}

struct CriticalSection {
    /// How many callers are in a critical_section (they can be nested).
    observers: usize,
    /// We keep a strong ref count in critical sections so that callers don't attempt to cancel
    /// this future if we know it will not care. It is possible for this be None if we get
    /// cancelled while polling. This is unlike what we do in StructuredCancellation where since we
    /// inform the future that it is cancelled (which will be observable), we still want callers to
    /// know if they cancelled it.
    _guard: Option<StrongRefCount>,
}

struct StructuredCancellation {
    /// How many observers we are tracking.
    observers: usize,
    tx: Option<oneshot::Sender<()>>,
    rx: Shared<oneshot::Receiver<()>>,
}

/// NOTE: this future is intended only to be polled in a consistent tokio runtime, and never moved
/// from one executor to another.
/// The general safe way of using this future is to spawn it directly via `tokio::spawn`.
#[pin_project(project = CancellableFutureProj)]
pub struct CancellableFuture<F> {
    shared: SharedState,

    /// This is notionally a `ExecutionContext` field, but we put it in an Option<Box<...>> to
    /// cheaply move it into a thread local every time we enter `poll()`. This is used for the
    /// running future to be able to access the API we expose to e.g. upgrade the refcount.
    execution: Option<Box<ExecutionContext>>,

    /// NOTE: this is duplicative of the `SharedState`, but unlike that state this is not behind a
    /// lock. This avoids us needing to grab the lock to check if we're Pending every time we poll.
    started: bool,

    #[pin]
    future: F,
}

impl<F> CancellableFutureProj<'_, F>
where
    F: Future,
{
    fn poll_inner(&mut self, cx: &mut Context<'_>) -> Poll<Option<<F as Future>::Output>> {
        let is_cancelled = self.shared.inner.cancelled.load(Ordering::SeqCst);

        if is_cancelled {
            let mut execution = self.execution.as_mut().unwrap().shared.lock();
            execution.notify_cancelled();
            if execution.can_exit() {
                return Poll::Ready(None);
            }
        }

        struct ReplaceOnDrop<'a> {
            me: &'a mut Option<Box<ExecutionContext>>,
            previous: Option<Box<ExecutionContext>>,
        }

        impl Drop for ReplaceOnDrop<'_> {
            fn drop(&mut self) {
                *self.me = CURRENT.with(|g| g.replace(self.previous.take()))
            }
        }

        let res = {
            let previous = CURRENT.with(|g| g.replace(Some(self.execution.take().unwrap())));
            let _replace = ReplaceOnDrop {
                previous,
                me: self.execution,
            };
            self.future.as_mut().poll(cx).map(Some)
        };

        // If we were using structured cancellation but just exited the critical section, then we
        // should exit now.
        if res.is_pending()
            && is_cancelled
            && self.execution.as_ref().unwrap().shared.lock().can_exit()
        {
            return Poll::Ready(None);
        }

        res
    }
}

impl<F> Future for CancellableFuture<F>
where
    F: Future,
{
    type Output = Option<<F as Future>::Output>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();

        if !*this.started {
            // we only update the Waker once at the beginning of the poll. For the same tokio
            // runtime, this is always safe and behaves correctly, as such, this future is
            // restricted to be ran on the same tokio executor and never moved from one runtime to
            // another
            take_mut::take(
                &mut *this.shared.inner.state.lock(),
                |future| match future {
                    State::Pending => State::Polled {
                        waker: cx.waker().clone(),
                    },
                    other => other,
                },
            );

            *this.started = true;
        }

        let poll = this.poll_inner(cx);

        // When we exit, release our waker to ensure we don't keep create a reference cycle for
        // this task.
        if poll.is_ready() {
            let _ignored = std::mem::replace(&mut *this.shared.inner.state.lock(), State::Exited);
        }

        poll
    }
}

struct RefCountData {
    shared: SharedState,
}

impl Drop for RefCountData {
    fn drop(&mut self) {
        self.shared.cancel();
    }
}

#[derive(Clone_, Dupe_)]
pub struct StrongRefCount {
    inner: Arc<RefCountData>,
}

impl StrongRefCount {
    pub fn downgrade(&self) -> WeakRefCount {
        WeakRefCount {
            inner: Arc::downgrade(&self.inner),
        }
    }
}

#[derive(Clone_, Dupe_)]
pub struct WeakRefCount {
    inner: Weak<RefCountData>,
}

impl WeakRefCount {
    pub fn upgrade(&self) -> Option<StrongRefCount> {
        self.inner.upgrade().map(|inner| StrongRefCount { inner })
    }
}

impl<F> CancellableFuture<F> {
    pub fn new_refcounted(future: F) -> (Self, StrongRefCount) {
        let shared = SharedState {
            inner: Arc::new(SharedStateData {
                state: Mutex::new(State::Pending),
                cancelled: AtomicBool::new(false),
            }),
        };

        let ref_count = StrongRefCount {
            inner: Arc::new(RefCountData {
                shared: shared.dupe(),
            }),
        };

        (
            CancellableFuture {
                shared,
                execution: Some(Box::new(ExecutionContext::new(ref_count.downgrade()))),
                started: false,
                future,
            },
            ref_count,
        )
    }
}

struct WithExecutionContext<F: FnOnce(&mut ExecutionContextShared)> {
    exit: Option<(F, Arc<Mutex<ExecutionContextShared>>)>,
}

impl<F> WithExecutionContext<F>
where
    F: FnOnce(&mut ExecutionContextShared),
{
    fn enter<Enter, T>(enter: Enter, exit: F) -> (T, Self)
    where
        Enter: FnOnce(&mut ExecutionContext) -> T,
        T: Default,
    {
        let (t, exit) = CURRENT.with(|g| match g.borrow_mut().as_mut() {
            Some(ctx) => {
                let t = enter(ctx);
                (t, Some((exit, ctx.shared.dupe())))
            }
            None => (T::default(), None),
        });

        (t, Self { exit })
    }

    async fn exit(mut self) {
        let mut should_yield = false;

        if let Some((exit, ctx)) = self.exit.take() {
            let mut ctx = ctx.lock();
            exit(&mut ctx);
            should_yield = ctx.can_exit();
        }

        // If the current future should actually be cancelled now, we try to return control to it
        // immediately to allow cancellation to kick in faster.
        if should_yield {
            tokio::task::yield_now().await;
        }
    }
}

impl<F> Drop for WithExecutionContext<F>
where
    F: for<'a> FnOnce(&'a mut ExecutionContextShared),
{
    fn drop(&mut self) {
        if let Some((exit, ctx)) = self.exit.take() {
            exit(&mut ctx.lock())
        }

        // TODO: make it a soft_error if the current future isn't the one that entered the critical
        // section
    }
}

/// Enter a critical section during which the current future (if any) should not be dropped. If the
/// future was not cancelled before entering the critical section, it becomes non-cancellable
/// during the critical section. If it *was* cancelled before entering the critical section (i.e.
/// the last ref was dropped during `poll`), then the future is allowed to continue executing until
/// this future resolves.
pub(crate) fn critical_section<F, Fut>(make: F) -> impl Future<Output = <Fut as Future>::Output>
where
    F: FnOnce() -> Fut,
    Fut: Future,
{
    let ((), guard) = WithExecutionContext::enter(
        ExecutionContext::enter_critical_section,
        ExecutionContextShared::exit_critical_section,
    );

    let fut = make();

    async move {
        let res = fut.await;
        guard.exit().await;
        res
    }
}

#[derive(Clone, Default)]
#[pin_project]
pub struct CancellationObserver {
    #[pin]
    pub(crate) rx: Option<Shared<oneshot::Receiver<()>>>,
}

impl Dupe for CancellationObserver {}

impl Future for CancellationObserver {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        match this.rx.as_pin_mut() {
            Some(rx) => rx.poll(cx).map(|_| ()),
            None => Poll::Pending,
        }
    }
}

/// Enter a structured cancellation section. The caller receives a CancellationObserver. The
/// CancellationObserver is a future that resolves when cancellation is requested (or when this
/// section exits).
pub(crate) fn with_structured_cancellation<F, Fut>(
    make: F,
) -> impl Future<Output = <Fut as Future>::Output>
where
    Fut: Future,
    F: FnOnce(CancellationObserver) -> Fut,
{
    let (obs, guard) = WithExecutionContext::enter(
        ExecutionContext::enter_structured_cancellation,
        ExecutionContextShared::exit_structured_cancellation,
    );

    let fut = make(obs);

    async move {
        let res = fut.await;
        guard.exit().await;
        res
    }
}

pub struct DisableCancellationGuard {
    /// This might be None if we are not in a CancellableFuture.
    pub(crate) _guard: Option<StrongRefCount>,
}

/// Obtain a StrongRefCount for the current task. This will return None if the task *is* within a
/// CancellableFuture but has already been cancelled.
pub(crate) fn try_to_disable_cancellation() -> Option<DisableCancellationGuard> {
    CURRENT.with(
        |g| match g.borrow().as_ref().map(|g| g.ref_count.upgrade()) {
            Some(Some(g)) => Some(DisableCancellationGuard { _guard: Some(g) }),
            Some(None) => None,
            None => Some(DisableCancellationGuard { _guard: None }),
        },
    )
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use assert_matches::assert_matches;

    use super::*;

    struct MaybePanicOnDrop {
        panic: bool,
    }

    impl Drop for MaybePanicOnDrop {
        fn drop(&mut self) {
            if self.panic {
                panic!()
            }
        }
    }

    #[tokio::test]
    async fn test_ready() {
        let (fut, _guard) = CancellableFuture::new_refcounted(futures::future::ready(()));
        assert_matches!(futures::poll!(fut), Poll::Ready(Some(())));
    }

    #[tokio::test]
    async fn test_cancel() {
        let (fut, guard) = CancellableFuture::new_refcounted(futures::future::pending::<()>());
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        drop(guard);
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_wakeup() {
        let (fut, guard) = CancellableFuture::new_refcounted(futures::future::pending::<()>());
        let task = tokio::task::spawn(fut);
        futures::pin_mut!(task);

        assert_matches!(
            tokio::time::timeout(Duration::from_millis(100), &mut task).await,
            Err(..)
        );

        drop(guard);

        assert_matches!(
            tokio::time::timeout(Duration::from_millis(100), &mut task).await,
            Ok(Ok(None))
        );
    }

    #[tokio::test]
    async fn test_is_dropped() {
        let dropped = Arc::new(Mutex::new(false));

        struct SetOnDrop {
            dropped: Arc<Mutex<bool>>,
        }

        impl Drop for SetOnDrop {
            fn drop(&mut self) {
                *self.dropped.lock() = true;
            }
        }

        impl Future for SetOnDrop {
            type Output = ();

            fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
                Poll::Ready(())
            }
        }

        let (fut, _guard) = CancellableFuture::new_refcounted(SetOnDrop {
            dropped: dropped.dupe(),
        });

        let task = tokio::task::spawn(fut);

        task.await.unwrap();
        assert!(*dropped.lock());
    }

    #[tokio::test]
    async fn test_critical_section() {
        let (fut, guard) = CancellableFuture::new_refcounted(async {
            {
                critical_section(tokio::task::yield_now).await;
            }
            futures::future::pending::<()>().await
        });
        futures::pin_mut!(fut);

        // We reach the first yield. At this point there are 2 guards: ours and the one held via
        // critical_section().
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        // Drop our guard, then poll again. Cancellation is checked, *then* the guard in the future
        // is dropped, so at this point we proceed to the pending() step after havin cancelled the
        // future (we would get notified for wakeup if we weren't manually polling).
        drop(guard);
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        // Poll again, this time we don't enter the future's poll because it is cancelled.
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_nested_critical_section() {
        let (fut, guard) = CancellableFuture::new_refcounted(async {
            {
                critical_section(|| async move { tokio::task::yield_now().await }).await;
            }
            futures::future::pending::<()>().await
        });
        futures::pin_mut!(fut);

        // We reach the first yield.
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        drop(guard);
        fut.await;
    }

    #[tokio::test]
    async fn test_critical_section_cancelled_during_poll() {
        let guard_slot = Mutex::new(None);

        let (fut, guard) = CancellableFuture::new_refcounted(async {
            {
                guard_slot
                    .lock()
                    .take()
                    .expect("Expected the guard to be here by now");

                critical_section(|| async {
                    let mut panic = MaybePanicOnDrop { panic: true };
                    tokio::task::yield_now().await;
                    panic.panic = false;
                })
                .await;
            }
            futures::future::pending::<()>().await
        });
        futures::pin_mut!(fut);

        *guard_slot.lock() = Some(guard);

        // Run the future. It'll drop the guard (and cancel itself) after entering the critical
        // section while it's being polled, but it'll proceed to the end.
        fut.await;
    }

    // Cases to test:
    // - Basic
    // - Reentrant
    // - Cancel when exiting critical section (with no further wakeups)

    #[tokio::test]
    async fn test_structured_cancellation_notifies() {
        let (fut, guard) = CancellableFuture::new_refcounted(async {
            with_structured_cancellation(|observer| observer).await;
        });
        futures::pin_mut!(fut);

        // Proceed all the way to awaiting the observer
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        // Drop our guard. At this point we'll cancel, and notify the observer.
        drop(guard);
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(..));
    }

    #[tokio::test]
    async fn test_structured_cancellation_is_blocking() {
        let (fut, guard) = CancellableFuture::new_refcounted(async {
            with_structured_cancellation(|_observer| async move {
                let mut panic = MaybePanicOnDrop { panic: true };
                tokio::task::yield_now().await;
                panic.panic = false;
            })
            .await;
        });
        futures::pin_mut!(fut);

        // Proceed all the way to the first pending.
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        // Drop our guard. We should resume and disarm the guard.
        drop(guard);
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(..));
    }

    #[tokio::test]
    async fn test_structured_cancellation_cancels_on_exit() {
        let (fut, guard) = CancellableFuture::new_refcounted(async {
            with_structured_cancellation(|observer| observer).await;
            futures::future::pending::<()>().await
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        drop(guard);
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    // This is a bit of an implementation detail.
    #[tokio::test]
    async fn test_structured_cancellation_returns_to_executor() {
        let (fut, guard) = CancellableFuture::new_refcounted(async {
            with_structured_cancellation(|observer| observer).await
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        drop(guard);
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_structured_cancellation_is_reentrant() {
        let (fut, guard) = CancellableFuture::new_refcounted(async {
            with_structured_cancellation(|o1| async move {
                with_structured_cancellation(|o2| async move {
                    o2.await;
                    o1.await;
                })
                .await;
            })
            .await;
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        drop(guard);
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(..));
    }

    #[tokio::test]
    async fn test_structured_cancellation_can_be_reentered() {
        let (fut, guard) = CancellableFuture::new_refcounted(async {
            with_structured_cancellation(|_o1| async move {}).await;
            with_structured_cancellation(|o2| async move {
                o2.await;
            })
            .await;
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        drop(guard);
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(..));
    }

    #[tokio::test]
    async fn test_structured_cancellation_works_after_cancel() {
        let (fut, guard) = CancellableFuture::new_refcounted(async {
            with_structured_cancellation(|_o1| async move {
                tokio::task::yield_now().await;
                // At this point we'll get cancelled.
                with_structured_cancellation(|o2| async move {
                    o2.await;
                })
                .await;
            })
            .await;
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        drop(guard);
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn noop_drop_is_allowed() {
        let section = critical_section(futures::future::pending::<()>);

        let (fut, _guard) = CancellableFuture::new_refcounted(async {
            drop(section); // Drop it within an ExecutionContext
        });

        fut.await;
    }

    #[allow(clippy::async_yields_async)]
    #[tokio::test]
    async fn drop_outside_of_task_is_allowed() {
        let (fut, _guard) = CancellableFuture::new_refcounted(async {
            critical_section(futures::future::pending::<()>)
        });

        let section = fut.await;

        drop(section)
    }

    #[tokio::test]
    async fn drop_outside_of_task_takes_effect() {
        let (fut, guard) = CancellableFuture::new_refcounted(async {
            let section = critical_section(futures::future::pending::<()>);
            std::thread::spawn(move || drop(section));
            loop {
                tokio::task::yield_now().await;
            }
        });
        futures::pin_mut!(fut);

        // Advance to the yield then cance.
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);
        drop(guard);

        assert_matches!(
            tokio::time::timeout(Duration::from_millis(1000), fut).await,
            Ok(None)
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! Defines a future with explicit cancellation

pub mod future;

use std::future::Future;

use dupe::Dupe;
use futures::FutureExt;
use once_cell::sync::Lazy;

use crate::cancellable_future::critical_section;
use crate::cancellable_future::try_to_disable_cancellation;
use crate::cancellable_future::with_structured_cancellation;
use crate::cancellable_future::CancellationObserver;
use crate::cancellable_future::DisableCancellationGuard;
use crate::cancellation::future::CriticalSectionGuard;
use crate::cancellation::future::ExecutionContext;

static INSTANCE: Lazy<CancellationContext> =
    Lazy::new(|| CancellationContext(CancellationContextInner::ThreadLocal));

/// Context available to the function running inside the future to control and manage it's own
/// cancellation
pub struct CancellationContext<'a>(CancellationContextInner<'a>);

impl<'a> CancellationContext<'a> {
    /// TODO replace with real initialization
    pub fn todo() -> Self {
        CancellationContext(CancellationContextInner::ThreadLocal)
    }

    pub fn testing() -> &'a Self {
        &INSTANCE
    }

    /// Create a new context from a thread that is never canceled
    pub fn never_cancelled() -> &'static Self {
        &INSTANCE
    }

    /// Enter a critical section during which the current future (if supports explicit cancellation)
    /// should not be dropped. If the future was not cancelled before entering the critical section,
    /// it becomes non-cancellable during the critical section. If it *was* cancelled before
    /// entering the critical section (i.e. the last ref was dropped during `poll`), then the
    /// future is allowed to continue executing until this future resolves.
    pub fn critical_section<F, Fut>(
        &'a self,
        make: F,
    ) -> impl Future<Output = <Fut as Future>::Output> + 'a
    where
        F: FnOnce() -> Fut + 'a,
        Fut: Future + 'a,
    {
        // placeholder to just delegate to existing critical_section code for now until we migrate
        // the callsites
        self.0.critical_section(make)
    }

    /// Enter a structured cancellation section. The caller receives a CancellationObserver. The
    /// CancellationObserver is a future that resolves when cancellation is requested (or when this
    /// section exits).
    pub fn with_structured_cancellation<F, Fut>(
        &'a self,
        make: F,
    ) -> impl Future<Output = <Fut as Future>::Output> + 'a
    where
        Fut: Future + 'a,
        F: FnOnce(CancellationObserver) -> Fut + 'a,
    {
        self.0.with_structured_cancellation(make)
    }

    /// For CancellableFutures futures, obtain a StrongRefCount for the current task and prevent
    /// cancellation while the guard is held.
    ///
    /// For ExplicitCancellationFutures, disables the cancellation of the future from here on out
    /// so that it will never end with canceled.
    ///
    /// This will return None if the task *is* within a cancellable future but has already been cancelled.
    pub fn try_to_disable_cancellation(&self) -> Option<DisableCancellationGuard> {
        self.0.try_to_disable_cancellation()
    }
}

/// Context available to only explicitly cancellable futures to manage their own cancellation
pub struct ExplicitCancellationContext {
    inner: ExecutionContext,
}

/// When held, prevents cancellation of the current explicitly cancellable future.
/// When no longer held (unless drop or cancellations are disabled was not ran), cancellations will be allowed.
/// `disable_cancellations`ing this will leak the guard, causing cancellations to be disabled forever, regardless
/// of whether or not the task is already cancelled.
pub struct IgnoreCancellationGuard<'a> {
    guard: CriticalSectionGuard<'a>,
    observer: CancellationObserver,
}

impl<'a> IgnoreCancellationGuard<'a> {
    pub fn cancellation_observer(&self) -> CancellationObserver {
        self.observer.dupe()
    }

    /// Allow cancellations again, but unlike dropping it, also checks if we should be cancelled
    /// right now, at this specific await point.
    pub async fn allow_cancellations_again(self) {
        if self.guard.exit_prevent_cancellation() {
            // If the current future should actually be cancelled now, we try to return
            // control to it immediately to allow cancellation to kick in faster.
            tokio::task::yield_now().await;
        }
    }

    /// Ignores cancellations forever if we are not already cancelled. Otherwise, begin allowing
    /// cancellations again
    pub fn keep_going_on_cancellations_if_not_cancelled(self) -> Option<DisableCancellationGuard> {
        if self.guard.try_to_disable_cancellation() {
            Some(DisableCancellationGuard { _guard: None })
        } else {
            None
        }
    }
}

impl ExplicitCancellationContext {
    /// Ignore cancellations while 'PreventCancellation' is held
    pub fn begin_ignore_cancellation(&self) -> IgnoreCancellationGuard {
        let (observer, guard) = self.inner.enter_structured_cancellation();

        IgnoreCancellationGuard { guard, observer }
    }

    /// Enter a critical section during which the current future (if supports explicit cancellation)
    /// should not be dropped. If the future was not cancelled before entering the critical section,
    /// it becomes non-cancellable during the critical section. If it *was* cancelled before
    /// entering the critical section (i.e. the last ref was dropped during `poll`), then the
    /// future is allowed to continue executing until this future resolves.
    pub fn critical_section<'a, F, Fut>(
        &'a self,
        make: F,
    ) -> impl Future<Output = <Fut as Future>::Output> + 'a
    where
        F: FnOnce() -> Fut + 'a,
        Fut: Future + 'a,
    {
        let guard = self.begin_ignore_cancellation();
        async move {
            let r = make().await;

            guard.allow_cancellations_again().await;

            r
        }
    }

    /// Enter a structured cancellation section. The caller receives a CancellationObserver. The
    /// CancellationObserver is a future that resolves when cancellation is requested (or when this
    /// section exits).
    pub fn with_structured_cancellation<'a, F, Fut>(
        &'a self,
        make: F,
    ) -> impl Future<Output = <Fut as Future>::Output> + 'a
    where
        Fut: Future + 'a,
        F: FnOnce(CancellationObserver) -> Fut + 'a,
    {
        let guard = self.begin_ignore_cancellation();

        async move {
            let r = make(guard.cancellation_observer()).await;

            guard.allow_cancellations_again().await;

            r
        }
    }

    /// Tries to ignore the cancellation for this future from here on
    pub fn try_to_keep_going_on_cancellation(&self) -> Option<DisableCancellationGuard> {
        self.begin_ignore_cancellation()
            .keep_going_on_cancellations_if_not_cancelled()
    }

    pub fn into_compatible(&self) -> CancellationContext {
        CancellationContext(CancellationContextInner::Explicit(self))
    }

    pub fn testing<'a>() -> &'a Self {
        static INSTANCE: Lazy<ExplicitCancellationContext> =
            Lazy::new(|| ExplicitCancellationContext {
                inner: ExecutionContext::testing(),
            });

        &INSTANCE
    }
}

enum CancellationContextInner<'a> {
    /// The old thread local based implementation
    ThreadLocal,
    /// The cancellation context passed explicitly
    Explicit(&'a ExplicitCancellationContext),
}

impl<'a> CancellationContextInner<'a> {
    /// Enter a critical section during which the current future (if supports explicit cancellation)
    /// should not be dropped. If the future was not cancelled before entering the critical section,
    /// it becomes non-cancellable during the critical section. If it *was* cancelled before
    /// entering the critical section (i.e. the last ref was dropped during `poll`), then the
    /// future is allowed to continue executing until this future resolves.
    pub fn critical_section<F, Fut>(
        &'a self,
        make: F,
    ) -> impl Future<Output = <Fut as Future>::Output> + 'a
    where
        F: FnOnce() -> Fut + 'a,
        Fut: Future + 'a,
    {
        match self {
            CancellationContextInner::ThreadLocal => critical_section(make).left_future(),
            CancellationContextInner::Explicit(context) => {
                context.critical_section(make).right_future()
            }
        }
    }

    /// Enter a structured cancellation section. The caller receives a CancellationObserver. The
    /// CancellationObserver is a future that resolves when cancellation is requested (or when this
    /// section exits).
    pub fn with_structured_cancellation<F, Fut>(
        &'a self,
        make: F,
    ) -> impl Future<Output = <Fut as Future>::Output> + 'a
    where
        Fut: Future + 'a,
        F: FnOnce(CancellationObserver) -> Fut + 'a,
    {
        match self {
            CancellationContextInner::ThreadLocal => {
                with_structured_cancellation(make).left_future()
            }
            CancellationContextInner::Explicit(context) => {
                context.with_structured_cancellation(make).right_future()
            }
        }
    }

    /// Obtain a StrongRefCount for the current task. This will return None if the task *is* within a
    /// CancellableFuture but has already been cancelled.
    pub fn try_to_disable_cancellation(&self) -> Option<DisableCancellationGuard> {
        match self {
            CancellationContextInner::ThreadLocal => try_to_disable_cancellation(),
            CancellationContextInner::Explicit(context) => {
                context.try_to_keep_going_on_cancellation()
            }
        }
    }
}