ct_select.rs

use crate::{Choice, CtAssign};
use core::{
    cmp,
    num::{
        NonZeroI8, NonZeroI16, NonZeroI32, NonZeroI64, NonZeroI128, NonZeroU8, NonZeroU16,
        NonZeroU32, NonZeroU64, NonZeroU128,
    },
};

#[cfg(feature = "subtle")]
use crate::CtOption;
#[cfg(feature = "alloc")]
use {
    crate::CtAssignSlice,
    alloc::{boxed::Box, vec::Vec},
};

/// Constant-time selection: choose between two values based on a given [`Choice`].
///
/// This crate provides built-in implementations for the following types:
/// - [`i8`], [`i16`], [`i32`], [`i64`], [`i128`], [`isize`]
/// - [`u8`], [`u16`], [`u32`], [`u64`], [`u128`], [`usize`]
/// - [`NonZeroI8`], [`NonZeroI16`], [`NonZeroI32`], [`NonZeroI64`], [`NonZeroI128`]
/// - [`NonZeroU8`], [`NonZeroU16`], [`NonZeroU32`], [`NonZeroU64`], [`NonZeroU128`]
/// - [`cmp::Ordering`]
/// - [`Choice`]
/// - `[T; N]` where `T` impls [`CtSelectArray`], which the previously mentioned types all do,
///   as well as any type which impls [`Clone`] + [`CtAssignSlice`] + [`CtSelect`].
pub trait CtSelect: Sized {
    /// Select between `self` and `other` based on `choice`, returning a copy of the value.
    ///
    /// # Returns
    /// - `self` if `choice` is [`Choice::FALSE`].
    /// - `other` if `choice` is [`Choice::TRUE`].
    #[must_use]
    fn ct_select(&self, other: &Self, choice: Choice) -> Self;

    /// Conditionally swap `self` and `other` if `choice` is [`Choice::TRUE`].
    fn ct_swap(&mut self, other: &mut Self, choice: Choice) {
        let tmp = self.ct_select(other, choice);
        *other = Self::ct_select(other, self, choice);
        *self = tmp;
    }
}

/// Implementing this trait enables use of the [`CtSelect`] trait to construct `[T; N]` where `T`
/// is the `Self` type implementing the trait, via a blanket impl.
///
/// All types which impl [`Clone`] + [`CtAssignSlice`] + [`CtSelect`] will receive a blanket impl
/// of this trait and thus also be usable with the [`CtSelect`] impl for `[T; N]`.
pub trait CtSelectArray<const N: usize>: CtSelect + Sized {
    /// Select between `a` and `b` in constant-time based on `choice`.
    #[must_use]
    fn ct_select_array(a: &[Self; N], b: &[Self; N], choice: Choice) -> [Self; N] {
        core::array::from_fn(|i| Self::ct_select(&a[i], &b[i], choice))
    }
}

impl<T, const N: usize> CtSelect for [T; N]
where
    T: CtSelectArray<N>,
{
    #[inline]
    fn ct_select(&self, other: &Self, choice: Choice) -> Self {
        T::ct_select_array(self, other, choice)
    }
}

impl<T, const N: usize> CtSelectArray<N> for T
where
    T: Clone + CtSelect,
    [T]: CtAssign,
{
    #[inline]
    fn ct_select_array(a: &[Self; N], b: &[Self; N], choice: Choice) -> [Self; N] {
        let mut ret = a.clone();
        ret.ct_assign(b, choice);
        ret
    }
}

/// Marker trait which enables a blanket impl of [`CtSelect`] for types which also impl
/// [`Clone`] + [`CtAssign`].
pub trait CtSelectUsingCtAssign: Clone + CtAssign {}

impl<T: CtSelectUsingCtAssign> CtSelect for T {
    #[inline]
    fn ct_select(&self, other: &Self, choice: Choice) -> Self {
        let mut ret = self.clone();
        ret.ct_assign(other, choice);
        ret
    }
}

/// Macro to write impls of `CtSelectUsingCtAssign`.
macro_rules! impl_ct_select_with_ct_assign {
    ( $($ty:ty),+ ) => { $(impl CtSelectUsingCtAssign for $ty {})+ };
}

impl_ct_select_with_ct_assign!(
    i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize
);

/// Impl `CtSelect` for `NonZero<T>` by calling the `CtSelect` impl for `T`.
macro_rules! impl_ct_select_for_nonzero_integer {
    ( $($nzint:ident),+ ) => {
        $(
             impl CtSelect for $nzint {
                #[inline]
                fn ct_select(&self, rhs: &Self, choice: Choice) -> Self {
                    let n = self.get().ct_select(&rhs.get(), choice);

                    // SAFETY: we are constructing `NonZero` from a value we obtained from
                    // `NonZero::get`, which ensures it's non-zero.
                    #[allow(unsafe_code)]
                    unsafe { $nzint::new_unchecked(n) }
                }
            }
        )+
    };
}

impl_ct_select_for_nonzero_integer!(
    NonZeroI8,
    NonZeroI16,
    NonZeroI32,
    NonZeroI64,
    NonZeroI128,
    NonZeroU8,
    NonZeroU16,
    NonZeroU32,
    NonZeroU64,
    NonZeroU128
);

impl CtSelect for cmp::Ordering {
    fn ct_select(&self, other: &Self, choice: Choice) -> Self {
        // `Ordering` is `#[repr(i8)]` where:
        //
        // - `Less` => -1
        // - `Equal` => 0
        // - `Greater` => 1
        //
        // Given this, it's possible to operate on orderings as if they're `i8`, which allows us to
        // use the `CtSelect` impl on `i8` to select between them.
        let ret = (*self as i8).ct_select(&(*other as i8), choice);

        // SAFETY: `Ordering` is `#[repr(i8)]` and `ret` has been assigned to
        // a value which was originally a valid `Ordering` then cast to `i8`
        #[allow(trivial_casts, unsafe_code)]
        unsafe {
            *(&raw const ret).cast::<Self>()
        }
    }
}

#[cfg(feature = "alloc")]
impl<T: Clone + CtAssign> CtSelectUsingCtAssign for Box<T> {}

#[cfg(feature = "alloc")]
impl<T> CtSelectUsingCtAssign for Box<[T]>
where
    T: Clone,
    [T]: CtAssign,
{
}

#[cfg(feature = "alloc")]
impl<T: Clone + CtAssignSlice> CtSelectUsingCtAssign for Vec<T> {}

#[cfg(feature = "subtle")]
impl CtSelect for subtle::Choice {
    #[inline]
    fn ct_select(&self, other: &Self, choice: Choice) -> Self {
        Choice::from(*self)
            .ct_select(&Choice::from(*other), choice)
            .into()
    }
}

#[cfg(feature = "subtle")]
impl<T> CtSelect for subtle::CtOption<T>
where
    T: CtSelect + Default + subtle::ConditionallySelectable,
{
    #[inline]
    fn ct_select(&self, other: &Self, choice: Choice) -> Self {
        CtOption::from(*self)
            .ct_select(&CtOption::from(*other), choice)
            .into()
    }
}

#[cfg(test)]
mod tests {
    use super::{Choice, CtSelect, cmp};

    macro_rules! ct_select_test {
        ($ty:ty, $name:ident) => {
            #[test]
            fn $name() {
                let a: $ty = 1;
                let b: $ty = 2;
                assert_eq!(a.ct_select(&b, Choice::FALSE), a);
                assert_eq!(a.ct_select(&b, Choice::TRUE), b);
            }
        };
    }

    ct_select_test!(u8, u8_ct_select);
    ct_select_test!(u16, u16_ct_select);
    ct_select_test!(u32, u32_ct_select);
    ct_select_test!(u64, u64_ct_select);
    ct_select_test!(u128, u128_ct_select);

    #[test]
    fn ordering_ct_select() {
        let a = cmp::Ordering::Less;
        let b = cmp::Ordering::Greater;
        assert_eq!(a.ct_select(&b, Choice::FALSE), a);
        assert_eq!(a.ct_select(&b, Choice::TRUE), b);
    }
}