kernel/
ptr.rs

1// SPDX-License-Identifier: GPL-2.0
2
3//! Types and functions to work with pointers and addresses.
4
5use core::fmt::Debug;
6use core::mem::align_of;
7use core::num::NonZero;
8
9use crate::build_assert;
10
11/// Type representing an alignment, which is always a power of two.
12///
13/// It is used to validate that a given value is a valid alignment, and to perform masking and
14/// alignment operations.
15///
16/// This is a temporary substitute for the [`Alignment`] nightly type from the standard library,
17/// and to be eventually replaced by it.
18///
19/// [`Alignment`]: https://github.com/rust-lang/rust/issues/102070
20///
21/// # Invariants
22///
23/// An alignment is always a power of two.
24#[repr(transparent)]
25#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
26pub struct Alignment(NonZero<usize>);
27
28impl Alignment {
29    /// Validates that `ALIGN` is a power of two at build-time, and returns an [`Alignment`] of the
30    /// same value.
31    ///
32    /// A build error is triggered if `ALIGN` is not a power of two.
33    ///
34    /// # Examples
35    ///
36    /// ```
37    /// use kernel::ptr::Alignment;
38    ///
39    /// let v = Alignment::new::<16>();
40    /// assert_eq!(v.as_usize(), 16);
41    /// ```
42    #[inline(always)]
43    pub const fn new<const ALIGN: usize>() -> Self {
44        build_assert!(
45            ALIGN.is_power_of_two(),
46            "Provided alignment is not a power of two."
47        );
48
49        // INVARIANT: `align` is a power of two.
50        // SAFETY: `align` is a power of two, and thus non-zero.
51        Self(unsafe { NonZero::new_unchecked(ALIGN) })
52    }
53
54    /// Validates that `align` is a power of two at runtime, and returns an
55    /// [`Alignment`] of the same value.
56    ///
57    /// Returns [`None`] if `align` is not a power of two.
58    ///
59    /// # Examples
60    ///
61    /// ```
62    /// use kernel::ptr::Alignment;
63    ///
64    /// assert_eq!(Alignment::new_checked(16), Some(Alignment::new::<16>()));
65    /// assert_eq!(Alignment::new_checked(15), None);
66    /// assert_eq!(Alignment::new_checked(1), Some(Alignment::new::<1>()));
67    /// assert_eq!(Alignment::new_checked(0), None);
68    /// ```
69    #[inline(always)]
70    pub const fn new_checked(align: usize) -> Option<Self> {
71        if align.is_power_of_two() {
72            // INVARIANT: `align` is a power of two.
73            // SAFETY: `align` is a power of two, and thus non-zero.
74            Some(Self(unsafe { NonZero::new_unchecked(align) }))
75        } else {
76            None
77        }
78    }
79
80    /// Returns the alignment of `T`.
81    ///
82    /// This is equivalent to [`align_of`], but with the return value provided as an [`Alignment`].
83    #[inline(always)]
84    pub const fn of<T>() -> Self {
85        #![allow(clippy::incompatible_msrv)]
86        // This cannot panic since alignments are always powers of two.
87        //
88        // We unfortunately cannot use `new` as it would require the `generic_const_exprs` feature.
89        const { Alignment::new_checked(align_of::<T>()).unwrap() }
90    }
91
92    /// Returns this alignment as a [`usize`].
93    ///
94    /// It is guaranteed to be a power of two.
95    ///
96    /// # Examples
97    ///
98    /// ```
99    /// use kernel::ptr::Alignment;
100    ///
101    /// assert_eq!(Alignment::new::<16>().as_usize(), 16);
102    /// ```
103    #[inline(always)]
104    pub const fn as_usize(self) -> usize {
105        self.as_nonzero().get()
106    }
107
108    /// Returns this alignment as a [`NonZero`].
109    ///
110    /// It is guaranteed to be a power of two.
111    ///
112    /// # Examples
113    ///
114    /// ```
115    /// use kernel::ptr::Alignment;
116    ///
117    /// assert_eq!(Alignment::new::<16>().as_nonzero().get(), 16);
118    /// ```
119    #[inline(always)]
120    pub const fn as_nonzero(self) -> NonZero<usize> {
121        // Allow the compiler to know that the value is indeed a power of two. This can help
122        // optimize some operations down the line, like e.g. replacing divisions by bit shifts.
123        if !self.0.is_power_of_two() {
124            // SAFETY: Per the invariants, `self.0` is always a power of two so this block will
125            // never be reached.
126            unsafe { core::hint::unreachable_unchecked() }
127        }
128        self.0
129    }
130
131    /// Returns the base-2 logarithm of the alignment.
132    ///
133    /// # Examples
134    ///
135    /// ```
136    /// use kernel::ptr::Alignment;
137    ///
138    /// assert_eq!(Alignment::of::<u8>().log2(), 0);
139    /// assert_eq!(Alignment::new::<16>().log2(), 4);
140    /// ```
141    #[inline(always)]
142    pub const fn log2(self) -> u32 {
143        self.0.ilog2()
144    }
145
146    /// Returns the mask for this alignment.
147    ///
148    /// This is equivalent to `!(self.as_usize() - 1)`.
149    ///
150    /// # Examples
151    ///
152    /// ```
153    /// use kernel::ptr::Alignment;
154    ///
155    /// assert_eq!(Alignment::new::<0x10>().mask(), !0xf);
156    /// ```
157    #[inline(always)]
158    pub const fn mask(self) -> usize {
159        // No underflow can occur as the alignment is guaranteed to be a power of two, and thus is
160        // non-zero.
161        !(self.as_usize() - 1)
162    }
163}
164
165/// Trait for items that can be aligned against an [`Alignment`].
166pub trait Alignable: Sized {
167    /// Aligns `self` down to `alignment`.
168    ///
169    /// # Examples
170    ///
171    /// ```
172    /// use kernel::ptr::{Alignable, Alignment};
173    ///
174    /// assert_eq!(0x2f_usize.align_down(Alignment::new::<0x10>()), 0x20);
175    /// assert_eq!(0x30usize.align_down(Alignment::new::<0x10>()), 0x30);
176    /// assert_eq!(0xf0u8.align_down(Alignment::new::<0x1000>()), 0x0);
177    /// ```
178    fn align_down(self, alignment: Alignment) -> Self;
179
180    /// Aligns `self` up to `alignment`, returning `None` if aligning would result in an overflow.
181    ///
182    /// # Examples
183    ///
184    /// ```
185    /// use kernel::ptr::{Alignable, Alignment};
186    ///
187    /// assert_eq!(0x4fusize.align_up(Alignment::new::<0x10>()), Some(0x50));
188    /// assert_eq!(0x40usize.align_up(Alignment::new::<0x10>()), Some(0x40));
189    /// assert_eq!(0x0usize.align_up(Alignment::new::<0x10>()), Some(0x0));
190    /// assert_eq!(u8::MAX.align_up(Alignment::new::<0x10>()), None);
191    /// assert_eq!(0x10u8.align_up(Alignment::new::<0x100>()), None);
192    /// assert_eq!(0x0u8.align_up(Alignment::new::<0x100>()), Some(0x0));
193    /// ```
194    fn align_up(self, alignment: Alignment) -> Option<Self>;
195}
196
197/// Implement [`Alignable`] for unsigned integer types.
198macro_rules! impl_alignable_uint {
199    ($($t:ty),*) => {
200        $(
201        impl Alignable for $t {
202            #[inline(always)]
203            fn align_down(self, alignment: Alignment) -> Self {
204                // The operands of `&` need to be of the same type so convert the alignment to
205                // `Self`. This means we need to compute the mask ourselves.
206                ::core::num::NonZero::<Self>::try_from(alignment.as_nonzero())
207                    .map(|align| self & !(align.get() - 1))
208                    // An alignment larger than `Self` always aligns down to `0`.
209                    .unwrap_or(0)
210            }
211
212            #[inline(always)]
213            fn align_up(self, alignment: Alignment) -> Option<Self> {
214                let aligned_down = self.align_down(alignment);
215                if self == aligned_down {
216                    Some(aligned_down)
217                } else {
218                    Self::try_from(alignment.as_usize())
219                        .ok()
220                        .and_then(|align| aligned_down.checked_add(align))
221                }
222            }
223        }
224        )*
225    };
226}
227
228impl_alignable_uint!(u8, u16, u32, u64, usize);