//! Memory management internals for the bucket array

//!

//! The solver needs sparse arrays of sorting buckets as temporary space. To

//! minimize peak memory usage, we need to dynamically reallocate portions

//! of that space between usages. A naive solution could use a separate Vec for

//! each array of bucket keys/values. A quick benchmark confirms that this is

//! between 2% and 10% worse than if the solver could use a single block of

//! memory with predetermined packed layout.

//!

//! Using a static layout presents some challenges in initialization, since the

//! entire struct is now too large to fit on the stack reliably. Creating an

//! instance using Box::new() would cause a stack overflow.

//!

//! We can solve several of these problems by orienting the data structure

//! around MaybeUninit memory: No initialization needed, we can easily switch

//! between space layouts using a union without risk of transmuting live values,

//! and we avoid an additional redundant zeroing of memory.

//!

//! This module provides a trait, [`Uninit`], for memory that we can safely

//! treat as uninitialized when it isn't being used. We track that use via a

//! mutable reference, so that gives us mutually exclusive access enforced at

//! compile-time.

//!

//! That uninitialized memory can be used to build backing layouts that are

//! made from instances of [`BucketArrayMemory`]. The whole backing layout is

//! then instantiated using freshly allocated uninitialized memory.

//!

//! A [`BucketArrayMemory`] can be temporarily, using a mutable reference,

//! wrapped into a [`BucketArray`] or half of a [`BucketArrayPair`].

//! Internally, a [`BucketState`] tracks how many items are in each bucket.

//! The only supported write operation is appending to an underfull bucket.

//! Once initialized by such a write, bucket items may be randomly accessed

//! via the methods on the [`BucketArray`] or [`BucketArrayPair`].

//!

//! It's critical for memory safety that we only read from a [`MaybeUninit`]

//! that has definitely been initialized. The static lifetimes of the mutable

//! reference and our `counts` array ensure we always start out with zeroed

//! counts and fully assumed-uninitialized memory. Our write method must be

//! certain to never increment the counter without actually writing to the

//! [`MaybeUninit`]. See [`BucketState::insert`].

// We need to allow this warning because we conditionally make some private

// functions public, but their documentation links to private types.

#![allow(rustdoc::private_intra_doc_links)]

use num_traits::{One, Zero};

use std::alloc;

use std::mem::MaybeUninit;

use std::ops::{Add, Range};

/// Marker trait for types that are normally assumed uninitialized

///

/// # Safety

///

/// By implementing this trait, it's a guarantee that uninitialized memory

/// may be transmuted into this type safely. It implies that the type is `Copy`.

/// It should contain no bits except for [`MaybeUninit`] fields. Structs and

/// arrays made entirely of [`MaybeUninit`] are fine.

///

/// This memory is always assumed to be uninitialized unless we hold a mutable

/// reference that's associated with information about specific fields that

/// were initialized during the reference's lifetime.

#[cfg_attr(feature = "bucket-array", visibility::make(pub))]

pub(crate) unsafe trait Uninit: Copy {

    /// Allocate new uninitialized memory, returning a new Box.

}

/// Backing memory for a single key or value bucket array

///

/// Describes `N` buckets which each hold at most `M` items of type `T`.

///

/// Implements [`Uninit`]. Structs and unions made from `BucketArrayMemory`

/// can be soundly marked as [`Uninit`].

#[cfg_attr(feature = "bucket-array", visibility::make(pub))]

#[derive(Copy, Clone)]

pub(crate) struct BucketArrayMemory<

    // Number of buckets

    const N: usize,

    // Number of item slots in each bucket

    const M: usize,

    // Item type

    T: Copy,

>([[MaybeUninit<T>; M]; N]);

/// Implements the [`Uninit`] trait. This memory is always assumed to be

/// uninitialized unless we hold a mutable reference paired with bucket

/// usage information.

///

/// SAFETY: We can implement [`Uninit`] for [`BucketArrayMemory`]

///         because it's `Copy` and it contains only `MaybeUninit` bits.

unsafe impl<const N: usize, const M: usize, T: Copy> Uninit for BucketArrayMemory<N, M, T> {}

/// Types that can be used as a count of items in a bucket

#[cfg_attr(feature = "bucket-array", visibility::make(pub))]

pub(crate) trait Count: Copy + Zero + One + Into<usize> + Add<Self, Output = Self> {}

impl<T: Copy + Zero + One + Into<usize> + Add<Self, Output = Self>> Count for T {}

/// Common implementation for key/value and value-only bucket arrays

///

/// Tracks the number of items in each bucket. This is used by [`BucketArray`]

/// and [`BucketArrayPair`] implementations to safely access only initialized

/// items in the [`BucketArrayMemory`].

struct BucketState<

    // Number of buckets

    const N: usize,

    // Maximum number of items per bucket

    const CAP: usize,

    // Type for bucket item counts

    C: Count,

> {

    /// Number of initialized items in each bucket

    ///

    /// Each bucket B's valid item range would be `(0 .. counts[B])`

    counts: [C; N],

}

impl<const N: usize, const CAP: usize, C: Count> BucketState<N, CAP, C> {

    /// Create a new counter store.

    ///

    /// This will happen inside the lifetime of our mutable reference

    /// to the backing store memory.

    /// Append a new item to a specific bucket using a writer callback.

    ///

    /// The writer is invoked with an item index, after checking

    /// bucket capacity but before marking the new item as written.

    ///

    /// The writer must unconditionally write to the index it's given, in

    /// each of the backing memories covered by this state tracker.

    #[inline(always)]

        } else {

        }

    /// Look up the valid item range for a particular bucket.

    ///

    /// Panics if the bucket index is out of range. Item indices inside the

    /// returned range are initialized, and any outside may be uninitialized.

    #[inline(always)]

}

/// Concrete binding between one [`BucketState`] and one [`BucketArrayMemory`]

#[cfg_attr(feature = "bucket-array", visibility::make(pub))]

pub(crate) struct BucketArray<

    // Lifetime for mutable reference to the backing memory

    'a,

    // Number of buckets

    const N: usize,

    // Maximum number of items per bucket

    const CAP: usize,

    // Type for bucket item counts

    C: Count,

    // Item type

    A: Copy,

> {

    /// Reference to external backing memory for type A

    mem: &'a mut BucketArrayMemory<N, CAP, A>,

    /// Tracking for which items are in use within each bucket

    state: BucketState<N, CAP, C>,

}

impl<'a, const N: usize, const CAP: usize, C: Count, A: Copy> BucketArray<'a, N, CAP, C, A> {

    /// A new [`BucketArray`] wraps a new [`BucketState`] and some possibly-recycled [`BucketArrayMemory`]

    /// Look up the valid item range for a particular bucket.

    ///

    /// Panics if the bucket index is out of range.

    /// Look up the value of one item in one bucket.

    ///

    /// Panics if the indices are out of range.

        // SAFETY: This requires that our [`BucketState`] instance accurately

        //         represents which fields in [`mem`] have been initialized.

    /// Append a new item to a bucket.

    ///

    /// If the bucket is full, returns `Err(())` and makes no changes.

    #[cfg_attr(

        feature = "bucket-array",

}

/// Concrete binding between one [`BucketState`] and a pair of [`BucketArrayMemory`]

#[cfg_attr(feature = "bucket-array", visibility::make(pub))]

pub(crate) struct BucketArrayPair<

    // Lifetime for mutable reference to the first backing memory

    'a,

    // Lifetime for mutable reference to the second backing memory

    'b,

    // Number of buckets

    const N: usize,

    // Maximum number of items per bucket

    const CAP: usize,

    // Type for bucket item counts

    C: Count,

    // Type for items in the first backing memory

    A: Copy,

    // Type for items in the second backing memory

    B: Copy,

> {

    /// Reference to external backing memory for type `A`

    mem_a: &'a mut BucketArrayMemory<N, CAP, A>,

    /// Reference to external backing memory for type `B`

    mem_b: &'b mut BucketArrayMemory<N, CAP, B>,

    /// Tracking for which items are in use within each bucket

    state: BucketState<N, CAP, C>,

}

impl<'a, 'b, const N: usize, const CAP: usize, C: Count, A: Copy, B: Copy>

    BucketArrayPair<'a, 'b, N, CAP, C, A, B>

{

    /// A new [`BucketArray`] wraps a new [`BucketState`] and two [`BucketArrayMemory`]

    /// Look up the valid item range for a particular bucket.

    ///

    /// Panics if the bucket index is out of range.

    /// Look up the first value for one item in one bucket.

    ///

    /// Panics if the indices are out of range.

        // SAFETY: This requires that our [`BucketState`] instance accurately

        //         represents which fields in [`mem`] have been initialized.

    /// Look up the second value for one item in one bucket.

    ///

    /// Panics if the indices are out of range.

        // SAFETY: This requires that our [`BucketState`] instance accurately

        //         represents which fields in [`mem`] have been initialized.

    /// Append a new item pair to a bucket.

    ///

    /// If the bucket is full, returns Err(()) and makes no changes.

    #[cfg_attr(

        feature = "bucket-array",

    /// Transfer the [`BucketState`] to a new single [`BucketArray`],

    /// keeping the second half and dropping the first.

}