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, NonZeroU8, NonZeroU16,
7        NonZeroU32, NonZeroU64, NonZeroU128,
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`]
23/// - [`NonZeroU8`], [`NonZeroU16`], [`NonZeroU32`], [`NonZeroU64`], [`NonZeroU128`]
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    NonZeroU8,
110    NonZeroU16,
111    NonZeroU32,
112    NonZeroU64,
113    NonZeroU128,
114    cmp::Ordering
115);
116
117impl_ct_assign_with_cmov!(isize, usize);
118impl CtAssignSlice for isize {}
119impl CtAssignSlice for usize {}
120
121impl<T, const N: usize> CtAssign for [T; N]
122where
123    T: CtAssignSlice,
124{
125    #[inline]
126    fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
127        self.as_mut_slice().ct_assign(rhs, choice);
128    }
129}
130
131impl<T, const N: usize> CtAssignSlice for [T; N] where T: CtAssignSlice {}
132
133#[cfg(feature = "subtle")]
134impl CtAssign for subtle::Choice {
135    #[inline]
136    fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
137        *self = Self::ct_select(self, rhs, choice);
138    }
139}
140
141#[cfg(feature = "subtle")]
142impl<T> CtAssign for subtle::CtOption<T>
143where
144    T: Default + subtle::ConditionallySelectable,
145{
146    #[inline]
147    fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
148        use subtle::ConditionallySelectable as _;
149        self.conditional_assign(rhs, choice.into());
150    }
151}
152#[cfg(feature = "alloc")]
153mod alloc {
154    use super::{Choice, CtAssign, CtAssignSlice};
155    use ::alloc::{boxed::Box, vec::Vec};
156
157    impl<T> CtAssign for Box<T>
158    where
159        T: CtAssign,
160    {
161        #[inline]
162        #[track_caller]
163        fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
164            (**self).ct_assign(rhs, choice);
165        }
166    }
167
168    impl<T> CtAssign for Box<[T]>
169    where
170        T: CtAssignSlice,
171    {
172        #[inline]
173        #[track_caller]
174        fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
175            self.ct_assign(&**rhs, choice);
176        }
177    }
178
179    impl<T> CtAssign<[T]> for Box<[T]>
180    where
181        T: CtAssignSlice,
182    {
183        #[inline]
184        #[track_caller]
185        fn ct_assign(&mut self, rhs: &[T], choice: Choice) {
186            (**self).ct_assign(rhs, choice);
187        }
188    }
189
190    impl<T> CtAssign for Vec<T>
191    where
192        T: CtAssignSlice,
193    {
194        #[inline]
195        #[track_caller]
196        fn ct_assign(&mut self, rhs: &Self, choice: Choice) {
197            self.ct_assign(rhs.as_slice(), choice);
198        }
199    }
200
201    impl<T> CtAssign<[T]> for Vec<T>
202    where
203        T: CtAssignSlice,
204    {
205        #[inline]
206        #[track_caller]
207        fn ct_assign(&mut self, rhs: &[T], choice: Choice) {
208            self.as_mut_slice().ct_assign(rhs, choice);
209        }
210    }
211}