//! Macro definitions.

use core::slice;

/// Generates a mask the width of the input type if the input value is non-zero.
///
/// Uses `core::hint::black_box` to coerce our desired codegen based on real-world observations
/// of the assembly generated by Rust/LLVM.
///
/// See also:
/// - CVE-2026-23519
/// - RustCrypto/utils#1332
macro_rules! masknz {
    ($value:tt : $int:ident) => {{
        let mut value: $int = $value;
        value |= value.wrapping_neg(); // has MSB `1` if non-zero, `0` if zero

        // use `black_box` to obscure we're computing a 1-bit value
        core::hint::black_box(
            value >> ($int::BITS - 1), // Extract MSB
        )
        .wrapping_neg() // Generate $int::MAX mask if `black_box` outputs `1`
    }};
}

// Uses 64-bit words on 64-bit targets, 32-bit everywhere else
#[cfg(not(target_pointer_width = "64"))]
pub(crate) type Word = u32;
#[cfg(target_pointer_width = "64")]
pub(crate) type Word = u64;
pub(crate) const WORD_SIZE: usize = size_of::<Word>();
const _: () = debug_assert!(size_of::<usize>() <= WORD_SIZE, "unexpected word size");

/// Rust core `[T]::as_chunks` vendored because of its 1.88 MSRV.
/// TODO(tarcieri): use upstream function when we bump MSRV
#[inline]
#[track_caller]
#[must_use]
#[allow(clippy::integer_division_remainder_used)]
pub(crate) fn slice_as_chunks<T, const N: usize>(slice: &[T]) -> (&[[T; N]], &[T]) {
    assert!(N != 0, "chunk size must be non-zero");
    let len_rounded_down = slice.len() / N * N;
    // SAFETY: The rounded-down value is always the same or smaller than the
    // original length, and thus must be in-bounds of the slice.
    let (multiple_of_n, remainder) = unsafe { slice.split_at_unchecked(len_rounded_down) };
    // SAFETY: We already panicked for zero, and ensured by construction
    // that the length of the subslice is a multiple of N.
    let array_slice = unsafe { slice_as_chunks_unchecked(multiple_of_n) };
    (array_slice, remainder)
}

/// Rust core `[T]::as_chunks_mut` vendored because of its 1.88 MSRV.
/// TODO(tarcieri): use upstream function when we bump MSRV
#[inline]
#[track_caller]
#[must_use]
#[allow(clippy::integer_division_remainder_used)]
pub(crate) fn slice_as_chunks_mut<T, const N: usize>(slice: &mut [T]) -> (&mut [[T; N]], &mut [T]) {
    assert!(N != 0, "chunk size must be non-zero");
    let len_rounded_down = slice.len() / N * N;
    // SAFETY: The rounded-down value is always the same or smaller than the
    // original length, and thus must be in-bounds of the slice.
    let (multiple_of_n, remainder) = unsafe { slice.split_at_mut_unchecked(len_rounded_down) };
    // SAFETY: We already panicked for zero, and ensured by construction
    // that the length of the subslice is a multiple of N.
    let array_slice = unsafe { slice_as_chunks_unchecked_mut(multiple_of_n) };
    (array_slice, remainder)
}

/// Rust core `[T]::as_chunks_unchecked` vendored because of its 1.88 MSRV.
/// TODO(tarcieri): use upstream function when we bump MSRV
#[inline]
#[must_use]
#[track_caller]
#[allow(clippy::integer_division_remainder_used)]
unsafe fn slice_as_chunks_unchecked<T, const N: usize>(slice: &[T]) -> &[[T; N]] {
    // SAFETY: Caller must guarantee that `N` is nonzero and exactly divides the slice length
    const { debug_assert!(N != 0) };
    debug_assert_eq!(slice.len() % N, 0);
    let new_len = slice.len() / N;

    // SAFETY: We cast a slice of `new_len * N` elements into
    // a slice of `new_len` many `N` elements chunks.
    unsafe { slice::from_raw_parts(slice.as_ptr().cast(), new_len) }
}

/// Rust core `[T]::as_chunks_unchecked_mut` vendored because of its 1.88 MSRV.
/// TODO(tarcieri): use upstream function when we bump MSRV
#[inline]
#[must_use]
#[track_caller]
#[allow(clippy::integer_division_remainder_used)]
unsafe fn slice_as_chunks_unchecked_mut<T, const N: usize>(slice: &mut [T]) -> &mut [[T; N]] {
    // SAFETY: Caller must guarantee that `N` is nonzero and exactly divides the slice length
    const { debug_assert!(N != 0) };
    debug_assert_eq!(slice.len() % N, 0);
    let new_len = slice.len() / N;

    // SAFETY: We cast a slice of `new_len * N` elements into
    // a slice of `new_len` many `N` elements chunks.
    unsafe { slice::from_raw_parts_mut(slice.as_mut_ptr().cast(), new_len) }
}

#[cfg(test)]
mod tests {
    // Spot check up to a given limit
    const TEST_LIMIT: u8 = 128;

    macro_rules! masknz_test {
        ( $($name:ident : $int:ident),+ ) => {
            $(
                #[test]
                fn $name() {
                    assert_eq!(masknz!(0: $int), 0);

                    // Test lower values
                    for i in 1..=$int::from(TEST_LIMIT) {
                        assert_eq!(masknz!(i: $int), $int::MAX);
                    }

                    // Test upper values
                    for i in ($int::MAX - $int::from(TEST_LIMIT))..=$int::MAX {
                        assert_eq!(masknz!(i: $int), $int::MAX);
                    }
                }
            )+
        }
    }

    // Ensure the macro works with any types we might use it with (we only use u8, u32, and u64)
    masknz_test!(
        masknz_u8: u8,
        masknz_u16: u16,
        masknz_u32: u32,
        masknz_u64: u64,
        masknz_u128: u128
    );
}
