Skip to main content

heapless/
len_type.rs

1use core::{
2    fmt::{Debug, Display},
3    mem,
4    ops::{Add, AddAssign, Sub, SubAssign},
5};
6
7#[allow(non_camel_case_types)]
8pub enum TypeEnum {
9    u8,
10    u16,
11    u32,
12    usize,
13}
14
15#[cfg(feature = "zeroize")]
16use zeroize::Zeroize;
17
18pub trait Sealed:
19    Send
20    + Sync
21    + Copy
22    + Display
23    + Debug
24    + PartialEq
25    + Add<Output = Self>
26    + AddAssign
27    + Sub<Output = Self>
28    + SubAssign
29    + PartialOrd
30    + TryFrom<usize, Error: Debug>
31    + TryInto<usize, Error: Debug>
32{
33    /// The zero value of the integer type.
34    const ZERO: Self;
35    /// The one value of the integer type.
36    const MAX: Self;
37    /// The maximum value of this type, as a `usize`.
38    const MAX_USIZE: usize;
39    /// This type as an enum.
40    const TYPE: TypeEnum;
41
42    /// The one value of the integer type.
43    ///
44    /// It's a function instead of constant because we want to have implementation which panics for
45    /// type `ZeroLenType`
46    fn one() -> Self;
47
48    /// An infallible conversion from `usize` to `LenT`.
49    #[inline]
50    fn from_usize(val: usize) -> Self {
51        val.try_into().unwrap()
52    }
53
54    /// An infallible conversion from `LenT` to `usize`.
55    #[inline]
56    fn into_usize(self) -> usize {
57        self.try_into().unwrap()
58    }
59
60    /// Converts `LenT` into `Some(usize)`, unless it's `Self::MAX`, where it returns `None`.
61    #[inline]
62    fn to_non_max(self) -> Option<usize> {
63        if self == Self::MAX {
64            None
65        } else {
66            Some(self.into_usize())
67        }
68    }
69}
70
71macro_rules! impl_lentype {
72    ($($(#[$meta:meta])* $LenT:ident),*) => {$(
73        $(#[$meta])*
74        impl Sealed for $LenT {
75            const ZERO: Self = 0;
76            const MAX: Self = Self::MAX;
77            const MAX_USIZE: usize = Self::MAX as _;
78            const TYPE: TypeEnum = TypeEnum::$LenT;
79
80            fn one() -> Self {
81                1
82            }
83        }
84
85        $(#[$meta])*
86        impl LenType for $LenT {}
87    )*}
88}
89
90/// A sealed trait representing a valid type to use as a length for a container.
91///
92/// This cannot be implemented in user code, and is restricted to `u8`, `u16`, `u32`, and `usize`.
93///
94/// When the `zeroize` feature is enabled, this trait requires the `Zeroize` trait.
95#[cfg(feature = "zeroize")]
96pub trait LenType: Sealed + Zeroize {}
97
98/// A sealed trait representing a valid type to use as a length for a container.
99///
100/// This cannot be implemented in user code, and is restricted to `u8`, `u16`, `u32`, and `usize`.
101#[cfg(not(feature = "zeroize"))]
102pub trait LenType: Sealed {}
103
104impl_lentype!(
105    u8,
106    u16,
107    #[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
108    u32,
109    usize
110);
111
112pub const fn check_capacity_fits<LenT: LenType, const N: usize>() {
113    assert!(LenT::MAX_USIZE >= N, "The capacity is larger than `LenT` can hold, increase the size of `LenT` or reduce the capacity");
114}
115
116/// Const cast from [`usize`] to [`LenType`] with `as`.
117#[inline]
118pub const fn as_len_type<L: LenType>(n: usize) -> L {
119    // SAFETY: transmute is safe since after cast we cast to the same type.
120    unsafe {
121        // ALWAYS compiletime switch.
122        match L::TYPE {
123            // transmute_copy, instead of transmute - because `L`
124            // is a "dependent type".
125            TypeEnum::u8 => mem::transmute_copy(&(n as u8)),
126            TypeEnum::u16 => mem::transmute_copy(&(n as u16)),
127            TypeEnum::u32 => mem::transmute_copy(&(n as u32)),
128            TypeEnum::usize => mem::transmute_copy(&n),
129        }
130    }
131}
132
133/// Checked cast to [`LenType`].
134///
135/// # Panic
136///
137/// Panics if `n` is outside of `L` range.
138#[inline]
139pub const fn to_len_type<L: LenType>(n: usize) -> L {
140    try_to_len_type(n).unwrap()
141}
142
143/// Checked cast to [`LenType`].
144///
145/// Returns `None` if `n` is outside of `L` range.
146#[inline]
147pub const fn try_to_len_type<L: LenType>(n: usize) -> Option<L> {
148    if n > L::MAX_USIZE {
149        return None;
150    }
151    Some(as_len_type(n))
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn test_len_cast() {
160        // 1. Check constness
161        const {
162            assert!(to_len_type::<u8>(150) == 150);
163            assert!(to_len_type::<u16>(15_000) == 15_000);
164            assert!(to_len_type::<u32>(1_500_000) == 1_500_000);
165            assert!(to_len_type::<usize>(usize::MAX) == usize::MAX);
166        }
167        // 2. Check correctness
168        fn check<T: LenType>() {
169            const COUNT: usize = 100;
170            for i in 0..COUNT {
171                let n = i * (T::MAX_USIZE / COUNT);
172                assert_eq!(to_len_type::<T>(n).into_usize(), n);
173            }
174        }
175        check::<u8>();
176        check::<u16>();
177        check::<u32>();
178        check::<usize>();
179    }
180}