ctutils/traits/
ct_select.rs

1use crate::{Choice, CtAssign, CtAssignSlice};
2use core::{
3    cmp,
4    num::{
5        NonZeroI8, NonZeroI16, NonZeroI32, NonZeroI64, NonZeroI128, NonZeroU8, NonZeroU16,
6        NonZeroU32, NonZeroU64, NonZeroU128,
7    },
8};
9
10#[cfg(feature = "subtle")]
11use crate::CtOption;
12
13/// Constant-time selection: choose between two values based on a given [`Choice`].
14///
15/// This crate provides built-in implementations for the following types:
16/// - [`i8`], [`i16`], [`i32`], [`i64`], [`i128`], [`isize`]
17/// - [`u8`], [`u16`], [`u32`], [`u64`], [`u128`], [`usize`]
18/// - [`NonZeroI8`], [`NonZeroI16`], [`NonZeroI32`], [`NonZeroI64`], [`NonZeroI128`]
19/// - [`NonZeroU8`], [`NonZeroU16`], [`NonZeroU32`], [`NonZeroU64`], [`NonZeroU128`]
20/// - [`cmp::Ordering`]
21/// - [`Choice`]
22/// - `[T; N]` where `T` impls [`CtSelectArray`], which the previously mentioned types all do,
23///   as well as any type which impls [`Clone`] + [`CtAssignSlice`] + [`CtSelect`].
24pub trait CtSelect: Sized {
25    /// Select between `self` and `other` based on `choice`, returning a copy of the value.
26    ///
27    /// # Returns
28    /// - `self` if `choice` is [`Choice::FALSE`].
29    /// - `other` if `choice` is [`Choice::TRUE`].
30    #[must_use]
31    fn ct_select(&self, other: &Self, choice: Choice) -> Self;
32
33    /// Conditionally swap `self` and `other` if `choice` is [`Choice::TRUE`].
34    fn ct_swap(&mut self, other: &mut Self, choice: Choice) {
35        let tmp = self.ct_select(other, choice);
36        *other = Self::ct_select(other, self, choice);
37        *self = tmp;
38    }
39}
40
41/// Implementing this trait enables use of the [`CtSelect`] trait to construct `[T; N]` where `T`
42/// is the `Self` type implementing the trait, via a blanket impl.
43///
44/// All types which impl [`Clone`] + [`CtAssignSlice`] + [`CtSelect`] will receive a blanket impl
45/// of this trait and thus also be usable with the [`CtSelect`] impl for `[T; N]`.
46pub trait CtSelectArray<const N: usize>: CtSelect + Sized {
47    /// Select between `a` and `b` in constant-time based on `choice`.
48    #[must_use]
49    fn ct_select_array(a: &[Self; N], b: &[Self; N], choice: Choice) -> [Self; N] {
50        core::array::from_fn(|i| Self::ct_select(&a[i], &b[i], choice))
51    }
52}
53
54impl<T, const N: usize> CtSelect for [T; N]
55where
56    T: CtSelectArray<N>,
57{
58    #[inline]
59    fn ct_select(&self, other: &Self, choice: Choice) -> Self {
60        T::ct_select_array(self, other, choice)
61    }
62}
63
64impl<T, const N: usize> CtSelectArray<N> for T
65where
66    T: Clone + CtAssignSlice + CtSelect,
67{
68    #[inline]
69    fn ct_select_array(a: &[Self; N], b: &[Self; N], choice: Choice) -> [Self; N] {
70        let mut ret = a.clone();
71        ret.ct_assign(b, choice);
72        ret
73    }
74}
75
76/// Marker trait which enables a blanket impl of [`CtSelect`] for types which also impl
77/// [`Clone`] + [`CtAssign`].
78pub trait CtSelectUsingCtAssign: Clone + CtAssign {}
79
80impl<T: CtSelectUsingCtAssign> CtSelect for T {
81    #[inline]
82    fn ct_select(&self, other: &Self, choice: Choice) -> Self {
83        let mut ret = self.clone();
84        ret.ct_assign(other, choice);
85        ret
86    }
87}
88
89/// Macro to write impls of `CtSelectUsingCtAssign`.
90macro_rules! impl_ct_select_with_ct_assign {
91    ( $($ty:ty),+ ) => { $(impl CtSelectUsingCtAssign for $ty {})+ };
92}
93
94impl_ct_select_with_ct_assign!(
95    i8,
96    i16,
97    i32,
98    i64,
99    i128,
100    isize,
101    u8,
102    u16,
103    u32,
104    u64,
105    u128,
106    usize,
107    NonZeroI8,
108    NonZeroI16,
109    NonZeroI32,
110    NonZeroI64,
111    NonZeroI128,
112    NonZeroU8,
113    NonZeroU16,
114    NonZeroU32,
115    NonZeroU64,
116    NonZeroU128,
117    cmp::Ordering
118);
119
120#[cfg(feature = "subtle")]
121impl CtSelect for subtle::Choice {
122    #[inline]
123    fn ct_select(&self, other: &Self, choice: Choice) -> Self {
124        Choice::from(*self)
125            .ct_select(&Choice::from(*other), choice)
126            .into()
127    }
128}
129
130#[cfg(feature = "subtle")]
131impl<T> CtSelect for subtle::CtOption<T>
132where
133    T: CtSelect + Default + subtle::ConditionallySelectable,
134{
135    #[inline]
136    fn ct_select(&self, other: &Self, choice: Choice) -> Self {
137        CtOption::from(*self)
138            .ct_select(&CtOption::from(*other), choice)
139            .into()
140    }
141}
142
143#[cfg(feature = "alloc")]
144mod alloc {
145    use super::CtSelectUsingCtAssign;
146    use crate::{CtAssign, CtAssignSlice};
147    use ::alloc::{boxed::Box, vec::Vec};
148
149    impl<T: Clone + CtAssign> CtSelectUsingCtAssign for Box<T> {}
150
151    #[cfg(feature = "alloc")]
152    impl<T> CtSelectUsingCtAssign for Box<[T]> where T: Clone + CtAssignSlice {}
153
154    #[cfg(feature = "alloc")]
155    impl<T: Clone + CtAssignSlice> CtSelectUsingCtAssign for Vec<T> {}
156}
157
158#[cfg(test)]
159mod tests {
160    use super::{Choice, CtSelect, cmp};
161
162    macro_rules! ct_select_test {
163        ($ty:ty, $name:ident) => {
164            #[test]
165            fn $name() {
166                let a: $ty = 1;
167                let b: $ty = 2;
168                assert_eq!(a.ct_select(&b, Choice::FALSE), a);
169                assert_eq!(a.ct_select(&b, Choice::TRUE), b);
170            }
171        };
172    }
173
174    ct_select_test!(u8, u8_ct_select);
175    ct_select_test!(u16, u16_ct_select);
176    ct_select_test!(u32, u32_ct_select);
177    ct_select_test!(u64, u64_ct_select);
178    ct_select_test!(u128, u128_ct_select);
179
180    #[test]
181    fn ordering_ct_select() {
182        let a = cmp::Ordering::Less;
183        let b = cmp::Ordering::Greater;
184        assert_eq!(a.ct_select(&b, Choice::FALSE), a);
185        assert_eq!(a.ct_select(&b, Choice::TRUE), b);
186    }
187}