ptr.rs

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