diskann-quantization 0.51.0

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

//! Number types with limited dynamic range.

use std::{fmt::Debug, num::NonZeroUsize};

use thiserror::Error;

/// A number type that must be greater than zero.
#[derive(Debug, Clone, Copy, PartialEq)]
#[repr(transparent)]
pub struct Positive<T>(T)
where
    T: PartialOrd + Default + Debug;

#[derive(Debug, Clone, Copy, Error)]
#[error("value {:?} is not greater than {:?} (its default value)", .0, T::default())]
pub struct NotPositiveError<T: Debug + Default>(T);

impl<T> Positive<T>
where
    T: PartialOrd + Default + Debug,
{
    /// Create a new `Positive` if the given value is greater than 0 (`T::default()`);
    pub fn new(value: T) -> Result<Self, NotPositiveError<T>> {
        if value > T::default() {
            Ok(Self(value))
        } else {
            Err(NotPositiveError(value))
        }
    }

    /// Create a new `Positive` without checking whether the value is greater than 0.
    ///
    /// # Safety
    ///
    /// The value must be greater than `T::default()`.
    pub const unsafe fn new_unchecked(value: T) -> Self {
        Self(value)
    }

    /// Consume `self` and return the inner value.
    pub fn into_inner(self) -> T {
        self.0
    }
}

// SAFETY: 1.0 is positive.
pub(crate) const POSITIVE_ONE_F32: Positive<f32> = unsafe { Positive::new_unchecked(1.0) };

#[derive(Debug, Clone, Copy, PartialEq)]
#[repr(transparent)]
pub struct PowerOfTwo(NonZeroUsize);

#[derive(Debug, Clone, Copy, Error)]
#[error("value {0} must be a power of two")]
#[non_exhaustive]
pub struct NotPowerOfTwo(usize);

macro_rules! constants {
    (($pow:ident, $value:ident) => $shift:literal) => {
        pub const $pow: Self = match Self::new(1 << $shift) {
            Ok(v) => v,
            Err(_) => panic!("not a power of two"),
        };

        pub const $value: Self = Self::$pow;
    };
    ($(($pow:ident, $value:ident) => $shift:literal),+ $(,)?) => {
        $(constants!(($pow, $value) => $shift);)+
    };
}

impl PowerOfTwo {
    constants! {
        (P0, V1) => 0,
        (P1, V2) => 1,
        (P2, V4) => 2,
        (P3, V8) => 3,
        (P4, V16) => 4,
        (P5, V32) => 5,
        (P6, V64) => 6,
        (P7, V128) => 7,
        (P8, V256) => 8,
        (P9, V512) => 9,
        (P10, V1024) => 10,
        (P11, V2048) => 11,
        (P12, V4096) => 12,
        (P13, V8192) => 13,
        (P14, V16384) => 14,
        (P15, V32768) => 15,
        (P16, V65536) => 16,
        (P17, V131072) => 17,
        (P18, V262144) => 18,
        (P19, V524288) => 19,
        (P20, V1048576) => 20,
        (P21, V2097152) => 21,
        (P22, V4194304) => 22,
        (P23, V8388608) => 23,
        (P24, V16777216) => 24,
        (P25, V33554432) => 25,
        (P26, V67108864) => 26,
        (P27, V134217728) => 27,
        (P28, V268435456) => 28,
        (P29, V536870912) => 29,
        (P30, V1073741824) => 30,
        (P31, V2147483648) => 31,
    }
}

#[cfg(target_pointer_width = "64")]
impl PowerOfTwo {
    constants! {
        (P32, V4294967296) => 32,
        (P33, V8589934592) => 33,
        (P34, V17179869184) => 34,
        (P35, V34359738368) => 35,
        (P36, V68719476736) => 36,
        (P37, V137438953472) => 37,
        (P38, V274877906944) => 38,
        (P39, V549755813888) => 39,
        (P40, V1099511627776) => 40,
        (P41, V2199023255552) => 41,
        (P42, V4398046511104) => 42,
        (P43, V8796093022208) => 43,
        (P44, V17592186044416) => 44,
        (P45, V35184372088832) => 45,
        (P46, V70368744177664) => 46,
        (P47, V140737488355328) => 47,
        (P48, V281474976710656) => 48,
        (P49, V562949953421312) => 49,
        (P50, V1125899906842624) => 50,
        (P51, V2251799813685248) => 51,
        (P52, V4503599627370496) => 52,
        (P53, V9007199254740992) => 53,
        (P54, V18014398509481984) => 54,
        (P55, V36028797018963968) => 55,
        (P56, V72057594037927936) => 56,
        (P57, V144115188075855872) => 57,
        (P58, V288230376151711744) => 58,
        (P59, V576460752303423488) => 59,
        (P60, V1152921504606846976) => 60,
        (P61, V2305843009213693952) => 61,
        (P62, V4611686018427387904) => 62,
        (P63, V9223372036854775808) => 63,
    }
}

impl PowerOfTwo {
    /// Create a new `PowerOfTwo` if the given value is greater a power of two.
    pub const fn new(value: usize) -> Result<Self, NotPowerOfTwo> {
        let v = match NonZeroUsize::new(value) {
            Some(value) => value,
            None => return Err(NotPowerOfTwo(value)),
        };
        if v.is_power_of_two() {
            // Safety: We just checked.
            Ok(unsafe { Self::new_unchecked(v) })
        } else {
            Err(NotPowerOfTwo(value))
        }
    }

    /// Return the smallest power of two greater than or equal to `value`. If the next
    /// power of two is greater than `usize::MAX`, `None` is returned.
    pub const fn next(value: usize) -> Option<Self> {
        // Note: use `match` instead of `Option::map` for `const`-compatibility.
        match value.checked_next_power_of_two() {
            // SAFETY: We trust the implementation of `usize::checked_next_power_of_two` since:
            //
            // * 0 can never be a power of two and thus cannot be returned.
            // * If it succeeds, the returned value should be a power of two.
            Some(v) => Some(unsafe { Self::new_unchecked(NonZeroUsize::new_unchecked(v)) }),
            None => None,
        }
    }

    /// Create a new `PowerOfTwo` without checking whether the value is a power of two.
    ///
    /// # Safety
    ///
    /// The value must be a power of two.
    pub const unsafe fn new_unchecked(value: NonZeroUsize) -> Self {
        Self(value)
    }

    /// Consume `self` and return the inner value.
    pub const fn into_inner(self) -> NonZeroUsize {
        self.0
    }

    /// Consume `self` and return the inner value as a `usize`.
    pub const fn raw(self) -> usize {
        self.0.get()
    }

    /// Construct `self` from the alignment in `layout`.
    pub const fn from_align(layout: &std::alloc::Layout) -> Self {
        // SAFETY: Alignment is guaranteed to be a power of two:
        // - <https://doc.rust-lang.org/beta/std/alloc/struct.Layout.html#method.align>
        unsafe { Self::new_unchecked(NonZeroUsize::new_unchecked(layout.align())) }
    }

    /// Return the alignment of `T` as a power of two.
    pub const fn alignment_of<T>() -> Self {
        // SAFETY: Alignment is guaranteed to be a power of two:
        // - <https://doc.rust-lang.org/beta/std/alloc/struct.Layout.html#method.align>
        unsafe { Self::new_unchecked(NonZeroUsize::new_unchecked(std::mem::align_of::<T>())) }
    }

    /// Compute the operation `lhs % self`.
    ///
    /// # Note
    ///
    /// The argument order of this function is reversed from the typical `align_offset`
    /// method in the standard library.
    pub const fn arg_mod(self, lhs: usize) -> usize {
        lhs & (self.raw() - 1)
    }

    /// Compute the amount `x` that would have to be added to `lhs` so the quantity
    /// `lhs + x` is a multiple of `self`.
    ///
    /// # Note
    ///
    /// The argument order of this function is reversed from the typical `align_offset`
    /// method in the standard library.
    pub const fn arg_align_offset(self, lhs: usize) -> usize {
        let m = self.arg_mod(lhs);
        if m == 0 { 0 } else { self.raw() - m }
    }

    /// Calculate the smallest value greater than or equal to `lhs` that is a multiple of
    /// `self`. Return `None` if the operation would result in an overflow.
    pub const fn arg_checked_next_multiple_of(self, lhs: usize) -> Option<usize> {
        let offset = self.arg_align_offset(lhs);
        lhs.checked_add(offset)
    }
}

impl From<PowerOfTwo> for usize {
    #[inline(always)]
    fn from(value: PowerOfTwo) -> Self {
        value.raw()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn format_not_positive_error<T>(value: T) -> String
    where
        T: Debug + Default,
    {
        format!(
            "value {:?} is not greater than {:?} (its default value)",
            value,
            T::default(),
        )
    }

    #[test]
    fn test_positive_f32() {
        let x = Positive::<f32>::new(1.0);
        assert!(x.is_ok());
        let x = x.unwrap();
        assert_eq!(x.into_inner(), 1.0);

        // Using 0 should return an error.
        let x = Positive::<f32>::new(0.0);
        assert!(x.is_err());
        assert_eq!(
            x.unwrap_err().to_string(),
            format_not_positive_error::<f32>(0.0)
        );

        // Using -1 should return an error.
        let x = Positive::<f32>::new(-1.0);
        assert!(x.is_err());
        assert_eq!(
            x.unwrap_err().to_string(),
            format_not_positive_error::<f32>(-1.0)
        );

        // SAFETY: 1.0 is greater than zero.
        let x = unsafe { Positive::<f32>::new_unchecked(1.0) };
        assert_eq!(x.into_inner(), 1.0);
    }

    #[test]
    fn test_positive_i64() {
        let x = Positive::<i64>::new(1);
        assert!(x.is_ok());
        let x = x.unwrap();
        assert_eq!(x.into_inner(), 1);

        // Using 0 should return an error.
        let x = Positive::<i64>::new(0);
        assert!(x.is_err());
        assert_eq!(
            x.unwrap_err().to_string(),
            format_not_positive_error::<i64>(0)
        );

        // Using -1 should return an error.
        let x = Positive::<i64>::new(-1);
        assert!(x.is_err());
        assert_eq!(
            x.unwrap_err().to_string(),
            format_not_positive_error::<i64>(-1)
        );

        // SAFETY: 1 is greater than zero.
        let x = unsafe { Positive::<i64>::new_unchecked(1) };
        assert_eq!(x.into_inner(), 1);
    }

    #[test]
    fn test_power_of_two() {
        assert!(PowerOfTwo::new(0).is_err());
        assert_eq!(PowerOfTwo::next(0).unwrap(), PowerOfTwo::new(1).unwrap());
        for i in 0..63 {
            let base = 2usize.pow(i);
            let p = PowerOfTwo::new(base).unwrap();
            assert_eq!(p.into_inner().get(), base);
            assert_eq!(p.raw(), base);
            assert_eq!(PowerOfTwo::new(base).unwrap().raw(), base);
            assert_eq!(<_ as Into<usize>>::into(p), base);

            if i != 1 {
                assert!(PowerOfTwo::new(base - 1).is_err(), "failed for i = {}", i);
                assert_eq!(PowerOfTwo::next(base - 1).unwrap().raw(), base);
            }

            if i != 0 {
                assert!(PowerOfTwo::new(base + 1).is_err(), "failed for i = {}", i);
            }

            assert_eq!(p.arg_mod(0), 0);
            assert_eq!(p.arg_mod(p.raw()), 0);

            assert_eq!(p.arg_align_offset(0), 0);
            assert_eq!(p.arg_align_offset(base), 0);

            assert_eq!(p.arg_checked_next_multiple_of(0), Some(0));
            assert_eq!(p.arg_checked_next_multiple_of(base), Some(base));

            assert_eq!(p.arg_checked_next_multiple_of(1), Some(base));

            if i > 1 {
                assert_eq!(p.arg_mod(base + 1), 1);
                assert_eq!(p.arg_mod(2 * base - 1), base - 1);

                assert_eq!(p.arg_align_offset(base + 1), base - 1);
                assert_eq!(p.arg_align_offset(2 * base - 1), 1);

                assert_eq!(p.arg_checked_next_multiple_of(base + 1), Some(2 * base));
                assert_eq!(p.arg_checked_next_multiple_of(2 * base - 1), Some(2 * base));
            }
        }

        assert!(PowerOfTwo::next(2usize.pow(63) + 1).is_none());
        assert!(PowerOfTwo::next(usize::MAX).is_none());
    }
}