Skip to main content

ctutils/traits/
ct_assign.rs

1use crate::Choice;
2use cmov::Cmov;
3use core::{
4    cmp,
5    num::{
6        NonZeroI8, NonZeroI16, NonZeroI32, NonZeroI64, NonZeroI128, NonZeroIsize, NonZeroU8,
7        NonZeroU16, NonZeroU32, NonZeroU64, NonZeroU128, NonZeroUsize,
8    },
9};
10
11#[cfg(feature = "subtle")]
12use crate::CtSelect;
13
14#[cfg(doc)]
15use core::num::NonZero;
16
17/// Constant-time conditional assignment: assign a given value to another based on a [`Choice`].
18///
19/// This crate provides built-in implementations for the following types:
20/// - [`i8`], [`i16`], [`i32`], [`i64`], [`i128`], [`isize`]
21/// - [`u8`], [`u16`], [`u32`], [`u64`], [`u128`], [`usize`]
22/// - [`NonZeroI8`], [`NonZeroI16`], [`NonZeroI32`], [`NonZeroI64`], [`NonZeroI128`], [`NonZeroI128`]
23/// - [`NonZeroU8`], [`NonZeroU16`], [`NonZeroU32`], [`NonZeroU64`], [`NonZeroU128`],, [`NonZeroUsize`]
24/// - [`cmp::Ordering`]
25/// - [`Choice`]
26/// - `[T]` and `[T; N]` where `T` impls [`CtAssignSlice`], which the previously mentioned
27///   types all do.
28pub trait CtAssign<Rhs: ?Sized = Self> {
29    /// Conditionally assign `src` to `self` if `choice` is [`Choice::TRUE`].
30    fn ct_assign(&mut self, src: &Rhs, choice: Choice);
31}
32
33/// Implementing this trait enables use of the [`CtAssign`] trait for `[T]` where `T` is the
34/// `Self` type implementing the trait, via a blanket impl.
35///
36/// It needs to be a separate trait from [`CtAssign`] because we need to be able to impl
37/// [`CtAssign`] for `[T]` which is `?Sized`.
38pub trait CtAssignSlice: CtAssign + Sized {
39    /// Conditionally assign `src` to `dst` if `choice` is [`Choice::TRUE`], or leave it unchanged
40    /// for [`Choice::FALSE`].
41    fn ct_assign_slice(dst: &mut [Self], src: &[Self], choice: Choice) {
42        assert_eq!(
43            dst.len(),
44            src.len(),
45            "source slice length ({}) does not match destination slice length ({})",
46            src.len(),
47            dst.len()
48        );
49
50        for (a, b) in dst.iter_mut().zip(src) {
51            a.ct_assign(b, choice);
52        }
53    }
54}
55
56impl<T: CtAssignSlice> CtAssign for [T] {
57    fn ct_assign(&mut self, src: &[T], choice: Choice) {
58        T::ct_assign_slice(self, src, choice);
59    }
60}
61
62/// Impl `CtAssign` using the `cmov::Cmov` trait
63macro_rules! impl_ct_assign_with_cmov {
64    ( $($ty:ty),+ ) => {
65        $(
66            impl CtAssign for $ty {
67                #[inline]
68                fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
69                    self.cmovnz(rhs, choice.into());
70                }
71            }
72        )+
73    };
74}
75
76/// Impl `CtAssign` and `CtAssignSlice` using the `cmov::Cmov` trait
77macro_rules! impl_ct_assign_slice_with_cmov {
78    ( $($ty:ty),+ ) => {
79        $(
80            impl_ct_assign_with_cmov!($ty);
81
82            impl CtAssignSlice for $ty {
83                #[inline]
84                fn ct_assign_slice(dst: &mut [Self], src: &[Self], choice: Choice) {
85                    dst.cmovnz(src, choice.into());
86                }
87            }
88        )+
89    };
90}
91
92// NOTE: impls `CtAssign` and `CtAssignSlice`
93impl_ct_assign_slice_with_cmov!(
94    i8,
95    i16,
96    i32,
97    i64,
98    i128,
99    u8,
100    u16,
101    u32,
102    u64,
103    u128,
104    NonZeroI8,
105    NonZeroI16,
106    NonZeroI32,
107    NonZeroI64,
108    NonZeroI128,
109    NonZeroIsize,
110    NonZeroU8,
111    NonZeroU16,
112    NonZeroU32,
113    NonZeroU64,
114    NonZeroU128,
115    NonZeroUsize,
116    cmp::Ordering
117);
118
119impl_ct_assign_with_cmov!(isize, usize);
120impl CtAssignSlice for isize {}
121impl CtAssignSlice for usize {}
122
123impl<T, const N: usize> CtAssign for [T; N]
124where
125    T: CtAssignSlice,
126{
127    #[inline]
128    fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
129        self.as_mut_slice().ct_assign(rhs, choice);
130    }
131}
132
133impl<T, const N: usize> CtAssignSlice for [T; N] where T: CtAssignSlice {}
134
135#[cfg(feature = "subtle")]
136impl CtAssign for subtle::Choice {
137    #[inline]
138    fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
139        *self = Self::ct_select(self, rhs, choice);
140    }
141}
142
143#[cfg(feature = "subtle")]
144impl<T> CtAssign for subtle::CtOption<T>
145where
146    T: Default + subtle::ConditionallySelectable,
147{
148    #[inline]
149    fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
150        use subtle::ConditionallySelectable as _;
151        self.conditional_assign(rhs, choice.into());
152    }
153}
154#[cfg(feature = "alloc")]
155mod alloc {
156    use super::{Choice, CtAssign, CtAssignSlice};
157    use ::alloc::{boxed::Box, vec::Vec};
158
159    impl<T> CtAssign for Box<T>
160    where
161        T: CtAssign,
162    {
163        #[inline]
164        #[track_caller]
165        fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
166            (**self).ct_assign(rhs, choice);
167        }
168    }
169
170    impl<T> CtAssign for Box<[T]>
171    where
172        T: CtAssignSlice,
173    {
174        #[inline]
175        #[track_caller]
176        fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
177            self.ct_assign(&**rhs, choice);
178        }
179    }
180
181    impl<T> CtAssign<[T]> for Box<[T]>
182    where
183        T: CtAssignSlice,
184    {
185        #[inline]
186        #[track_caller]
187        fn ct_assign(&mut self, rhs: &[T], choice: Choice) {
188            (**self).ct_assign(rhs, choice);
189        }
190    }
191
192    impl<T> CtAssign for Vec<T>
193    where
194        T: CtAssignSlice,
195    {
196        #[inline]
197        #[track_caller]
198        fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
199            self.ct_assign(rhs.as_slice(), choice);
200        }
201    }
202
203    impl<T> CtAssign<[T]> for Vec<T>
204    where
205        T: CtAssignSlice,
206    {
207        #[inline]
208        #[track_caller]
209        fn ct_assign(&mut self, rhs: &[T], choice: Choice) {
210            self.as_mut_slice().ct_assign(rhs, choice);
211        }
212    }
213}