Skip to main content

ctutils/traits/
ct_select.rs

1use crate::{Choice, CtAssign, CtAssignSlice};
2use core::{
3    cmp,
4    num::{
5        NonZeroI8, NonZeroI16, NonZeroI32, NonZeroI64, NonZeroI128, NonZeroIsize, NonZeroU8,
6        NonZeroU16, NonZeroU32, NonZeroU64, NonZeroU128, NonZeroUsize,
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`], [`NonZeroI128`]
19/// - [`NonZeroU8`], [`NonZeroU16`], [`NonZeroU32`], [`NonZeroU64`], [`NonZeroU128`],, [`NonZeroUsize`]
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    NonZeroIsize,
113    NonZeroU8,
114    NonZeroU16,
115    NonZeroU32,
116    NonZeroU64,
117    NonZeroU128,
118    NonZeroUsize,
119    cmp::Ordering
120);
121
122#[cfg(feature = "subtle")]
123impl CtSelect for subtle::Choice {
124    #[inline]
125    fn ct_select(&self, other: &Self, choice: Choice) -> Self {
126        Choice::from(*self)
127            .ct_select(&Choice::from(*other), choice)
128            .into()
129    }
130}
131
132#[cfg(feature = "subtle")]
133impl<T> CtSelect for subtle::CtOption<T>
134where
135    T: CtSelect + Default + subtle::ConditionallySelectable,
136{
137    #[inline]
138    fn ct_select(&self, other: &Self, choice: Choice) -> Self {
139        CtOption::from(*self)
140            .ct_select(&CtOption::from(*other), choice)
141            .into()
142    }
143}
144
145#[cfg(feature = "alloc")]
146mod alloc {
147    use super::CtSelectUsingCtAssign;
148    use crate::{CtAssign, CtAssignSlice};
149    use ::alloc::{boxed::Box, vec::Vec};
150
151    impl<T: Clone + CtAssign> CtSelectUsingCtAssign for Box<T> {}
152
153    #[cfg(feature = "alloc")]
154    impl<T> CtSelectUsingCtAssign for Box<[T]> where T: Clone + CtAssignSlice {}
155
156    #[cfg(feature = "alloc")]
157    impl<T: Clone + CtAssignSlice> CtSelectUsingCtAssign for Vec<T> {}
158}
159
160#[cfg(test)]
161mod tests {
162    use super::{Choice, CtSelect, cmp};
163
164    macro_rules! ct_select_test_unsigned {
165        ($ty:ty, $name:ident) => {
166            #[test]
167            fn $name() {
168                let a: $ty = 1;
169                let b: $ty = 2;
170                assert_eq!(a.ct_select(&b, Choice::FALSE), a);
171                assert_eq!(a.ct_select(&b, Choice::TRUE), b);
172            }
173        };
174    }
175
176    macro_rules! ct_select_test_signed {
177        ($ty:ty, $name:ident) => {
178            #[test]
179            fn $name() {
180                let a: $ty = 1;
181                let b: $ty = -2;
182                assert_eq!(a.ct_select(&b, Choice::FALSE), a);
183                assert_eq!(a.ct_select(&b, Choice::TRUE), b);
184            }
185        };
186    }
187
188    ct_select_test_unsigned!(u8, u8_ct_select);
189    ct_select_test_unsigned!(u16, u16_ct_select);
190    ct_select_test_unsigned!(u32, u32_ct_select);
191    ct_select_test_unsigned!(u64, u64_ct_select);
192    ct_select_test_unsigned!(u128, u128_ct_select);
193    ct_select_test_unsigned!(usize, usize_ct_select);
194
195    ct_select_test_signed!(i8, i8_ct_select);
196    ct_select_test_signed!(i16, i16_ct_select);
197    ct_select_test_signed!(i32, i32_ct_select);
198    ct_select_test_signed!(i64, i64_ct_select);
199    ct_select_test_signed!(i128, i128_ct_select);
200    ct_select_test_signed!(isize, isize_ct_select);
201
202    #[test]
203    fn ordering_ct_select() {
204        let a = cmp::Ordering::Less;
205        let b = cmp::Ordering::Greater;
206        assert_eq!(a.ct_select(&b, Choice::FALSE), a);
207        assert_eq!(a.ct_select(&b, Choice::TRUE), b);
208    }
209}