Add checked arithmetic functions to addresses, pages, and frames

src/addr.rs

@@ -168,6 +168,50 @@ impl VirtAddr {
         self.0 == 0
     }
 
+    #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))]
+    /// Add an offset to a virtual address.
+    ///
+    /// This function performs normal arithmetic addition and doesn't jump the
+    /// address gap. If you're looking for a successor operation that jumps the
+    /// address gap, use [`Step::forward`].
+    ///
+    /// This function returns `None` on overflow or if the result is not a
+    /// canonical address.
+    #[inline]
+    pub const fn checked_add(self, rhs: u64) -> Option<Self> {
+        let Some(addr) = self.0.checked_add(rhs) else {
+            return None;
+        };
+
+        let Ok(addr) = Self::try_new(addr) else {
+            return None;
+        };
+
+        Some(addr)
+    }
+
+    #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))]
+    /// Subtract an offset from a virtual address.
+    ///
+    /// This function performs normal arithmetic subtraction and doesn't jump
+    /// the address gap. If you're looking for a predecessor operation that
+    /// jumps the address gap, use [`Step::backward`].
+    ///
+    /// This function returns `None` on overflow or if the result is not a
+    /// canonical address.
+    #[inline]
+    pub const fn checked_sub(self, rhs: u64) -> Option<Self> {
+        let Some(addr) = self.0.checked_sub(rhs) else {
+            return None;
+        };
+
+        let Ok(addr) = Self::try_new(addr) else {
+            return None;
+        };
+
+        Some(addr)
+    }
+
     /// Aligns the virtual address upwards to the given alignment.
     ///
     /// See the `align_up` function for more information.
@@ -379,41 +423,21 @@ impl fmt::Pointer for VirtAddr {
 impl Add<u64> for VirtAddr {
     type Output = Self;
 
-    #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))]
     /// Add an offset to a virtual address.
     ///
-    /// This function performs normal arithmetic addition and doesn't jump the
-    /// address gap. If you're looking for a successor operation that jumps the
-    /// address gap, use [`Step::forward`].
-    ///
-    /// # Panics
-    ///
-    /// This function will panic on overflow or if the result is not a
-    /// canonical address.
+    /// This function panics when [`VirtAddr::checked_add`] returns `None`.
     #[inline]
     #[track_caller]
     fn add(self, rhs: u64) -> Self::Output {
-        VirtAddr::try_new(
-            self.0
-                .checked_add(rhs)
-                .expect("attempt to add with overflow"),
-        )
-        .expect("attempt to add resulted in non-canonical virtual address")
+        self.checked_add(rhs)
+            .expect("attempt to add with overflow or resulted in non-canonical virtual address")
     }
 }
 
 impl AddAssign<u64> for VirtAddr {
-    #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))]
     /// Add an offset to a virtual address.
     ///
-    /// This function performs normal arithmetic addition and doesn't jump the
-    /// address gap. If you're looking for a successor operation that jumps the
-    /// address gap, use [`Step::forward`].
-    ///
-    /// # Panics
-    ///
-    /// This function will panic on overflow or if the result is not a
-    /// canonical address.
+    /// This function panics when [`VirtAddr::checked_add`] returns `None`.
     #[inline]
     #[track_caller]
     fn add_assign(&mut self, rhs: u64) {
@@ -424,41 +448,22 @@ impl AddAssign<u64> for VirtAddr {
 impl Sub<u64> for VirtAddr {
     type Output = Self;
 
-    #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))]
     /// Subtract an offset from a virtual address.
     ///
-    /// This function performs normal arithmetic subtraction and doesn't jump
-    /// the address gap. If you're looking for a predecessor operation that
-    /// jumps the address gap, use [`Step::backward`].
-    ///
-    /// # Panics
-    ///
-    /// This function will panic on overflow or if the result is not a
-    /// canonical address.
+    /// This function panics when [`VirtAddr::checked_sub`] returns `None`.
     #[inline]
     #[track_caller]
     fn sub(self, rhs: u64) -> Self::Output {
-        VirtAddr::try_new(
-            self.0
-                .checked_sub(rhs)
-                .expect("attempt to subtract with overflow"),
+        self.checked_sub(rhs).expect(
+            "attempt to subtract with overflow or resulted in non-canonical virtual address",
         )
-        .expect("attempt to subtract resulted in non-canonical virtual address")
     }
 }
 
 impl SubAssign<u64> for VirtAddr {
-    #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))]
     /// Subtract an offset from a virtual address.
     ///
-    /// This function performs normal arithmetic subtraction and doesn't jump
-    /// the address gap. If you're looking for a predecessor operation that
-    /// jumps the address gap, use [`Step::backward`].
-    ///
-    /// # Panics
-    ///
-    /// This function will panic on overflow or if the result is not a
-    /// canonical address.
+    /// This function panics when [`VirtAddr::checked_sub`] returns `None`.
     #[inline]
     #[track_caller]
     fn sub_assign(&mut self, rhs: u64) {
@@ -618,6 +623,42 @@ impl PhysAddr {
         self.0 == 0
     }
 
+    /// Add an offset to a physical address.
+    ///
+    /// This function returns `None` on overflow or if the result is not a
+    /// valid physical address.
+    #[inline]
+    #[const_fn(cfg(not(feature = "memory_encryption")))]
+    pub const fn checked_add(self, rhs: u64) -> Option<Self> {
+        let Some(addr) = self.0.checked_add(rhs) else {
+            return None;
+        };
+
+        let Ok(addr) = Self::try_new(addr) else {
+            return None;
+        };
+
+        Some(addr)
+    }
+
+    /// Subtract an offset from a physical address.
+    ///
+    /// This function returns `None` on overflow or if the result is not a
+    /// valid physical address.
+    #[inline]
+    #[const_fn(cfg(not(feature = "memory_encryption")))]
+    pub const fn checked_sub(self, rhs: u64) -> Option<Self> {
+        let Some(addr) = self.0.checked_sub(rhs) else {
+            return None;
+        };
+
+        let Ok(addr) = Self::try_new(addr) else {
+            return None;
+        };
+
+        Some(addr)
+    }
+
     /// Aligns the physical address upwards to the given alignment.
     ///
     /// See the `align_up` function for more information.
@@ -717,7 +758,8 @@ impl Add<u64> for PhysAddr {
     #[inline]
     #[track_caller]
     fn add(self, rhs: u64) -> Self::Output {
-        PhysAddr::new(self.0.checked_add(rhs).unwrap())
+        self.checked_add(rhs)
+            .expect("attempt to add with overflow or resulted in invalid physical address")
     }
 }
 
@@ -734,7 +776,8 @@ impl Sub<u64> for PhysAddr {
     #[inline]
     #[track_caller]
     fn sub(self, rhs: u64) -> Self::Output {
-        PhysAddr::new(self.0.checked_sub(rhs).unwrap())
+        self.checked_sub(rhs)
+            .expect("attempt to subtract with overflow or resulted in invalid physical address")
     }
 }
 

src/structures/paging/frame.rs

@@ -155,6 +155,42 @@ impl<S: PageSize> PhysFrame<S> {
         self.start_address.as_u64() / S::SIZE
     }
 
+    /// Add an offset to a frame.
+    ///
+    /// This function returns `None` on overflow or if the result does not have a
+    /// valid physical address.
+    #[inline]
+    #[const_fn(cfg(not(feature = "memory_encryption")))]
+    pub const fn checked_add(self, rhs: u64) -> Option<Self> {
+        let Some(offset) = rhs.checked_mul(S::SIZE) else {
+            return None;
+        };
+
+        let Some(addr) = self.start_address().checked_add(offset) else {
+            return None;
+        };
+
+        Some(Self::containing_address(addr))
+    }
+
+    /// Subtract an offset from a frame.
+    ///
+    /// This function returns `None` on overflow or if the result does not have a
+    /// valid physical address.
+    #[inline]
+    #[const_fn(cfg(not(feature = "memory_encryption")))]
+    pub const fn checked_sub(self, rhs: u64) -> Option<Self> {
+        let Some(offset) = rhs.checked_mul(S::SIZE) else {
+            return None;
+        };
+
+        let Some(addr) = self.start_address().checked_sub(offset) else {
+            return None;
+        };
+
+        Some(Self::containing_address(addr))
+    }
+
     /// Returns a range of frames, exclusive `end`.
     #[inline]
     pub const fn range(start: PhysFrame<S>, end: PhysFrame<S>) -> PhysFrameRange<S> {
@@ -186,7 +222,8 @@ impl<S: PageSize> Add<u64> for PhysFrame<S> {
     #[inline]
     #[track_caller]
     fn add(self, rhs: u64) -> Self::Output {
-        PhysFrame::containing_address(self.start_address() + rhs * S::SIZE)
+        self.checked_add(rhs)
+            .expect("attempt to add with overflow or resulted in invalid physical address")
     }
 }
 
@@ -203,7 +240,8 @@ impl<S: PageSize> Sub<u64> for PhysFrame<S> {
     #[inline]
     #[track_caller]
     fn sub(self, rhs: u64) -> Self::Output {
-        PhysFrame::containing_address(self.start_address() - rhs * S::SIZE)
+        self.checked_sub(rhs)
+            .expect("attempt to subtract with overflow or resulted in invalid physical address")
     }
 }
 

src/structures/paging/mapper/mapped_page_table/iter.rs

@@ -4,6 +4,8 @@
 
 use core::ops::Add;
 
+use const_fn::const_fn;
+
 use super::{MappedPageTable, PageTableFrameMapping, PageTableWalkError, PageTableWalker};
 use crate::structures::paging::{
     Page, PageSize, PageTable, PageTableFlags, PageTableIndex, PhysFrame, Size1GiB, Size2MiB,
@@ -23,15 +25,32 @@ pub struct MappedPage<S: PageSize = Size4KiB> {
     pub flags: PageTableFlags,
 }
 
+impl<S: PageSize> MappedPage<S> {
+    #[const_fn(cfg(not(feature = "memory_encryption")))]
+    pub const fn checked_add(self, rhs: u64) -> Option<Self> {
+        let Some(page) = self.page.checked_add(rhs) else {
+            return None;
+        };
+
+        let Some(frame) = self.frame.checked_add(rhs) else {
+            return None;
+        };
+
+        Some(Self {
+            page,
+            frame,
+            flags: self.flags,
+        })
+    }
+}
+
 impl<S: PageSize> Add<u64> for MappedPage<S> {
     type Output = Self;
 
+    #[track_caller]
     fn add(self, rhs: u64) -> Self::Output {
-        Self {
-            page: self.page + rhs,
-            frame: self.frame + rhs,
-            flags: self.flags,
-        }
+        self.checked_add(rhs)
+            .expect("attempt to add with overflow or resulted in invalid mapped page")
     }
 }
 
@@ -48,15 +67,42 @@ pub enum MappedPageItem {
     Size1GiB(MappedPage<Size1GiB>),
 }
 
+impl MappedPageItem {
+    #[const_fn(cfg(not(feature = "memory_encryption")))]
+    pub const fn checked_add(self, rhs: u64) -> Option<Self> {
+        match self {
+            Self::Size4KiB(mapped_page) => {
+                let Some(page) = mapped_page.checked_add(rhs) else {
+                    return None;
+                };
+
+                Some(Self::Size4KiB(page))
+            }
+            Self::Size2MiB(mapped_page) => {
+                let Some(page) = mapped_page.checked_add(rhs) else {
+                    return None;
+                };
+
+                Some(Self::Size2MiB(page))
+            }
+            Self::Size1GiB(mapped_page) => {
+                let Some(page) = mapped_page.checked_add(rhs) else {
+                    return None;
+                };
+
+                Some(Self::Size1GiB(page))
+            }
+        }
+    }
+}
+
 impl Add<u64> for MappedPageItem {
     type Output = Self;
 
+    #[track_caller]
     fn add(self, rhs: u64) -> Self::Output {
-        match self {
-            Self::Size4KiB(mapped_page) => Self::Size4KiB(mapped_page + rhs),
-            Self::Size2MiB(mapped_page) => Self::Size2MiB(mapped_page + rhs),
-            Self::Size1GiB(mapped_page) => Self::Size1GiB(mapped_page + rhs),
-        }
+        self.checked_add(rhs)
+            .expect("attempt to add with overflow or resulted in invalid mapped page item")
     }
 }