fix Page Step impl on 32-bit platforms

Previously, we scaled count as a usize, not a u64. This is bad because
stepping forward by 0x10_0000 is perfectly fine, yet we were always
rejecting this because 0x10_0000 * 0x1000 doesn't fit in usize on
32-bit platforms. Similarly, we would fail to return a step count above
or equal to 0x10_0000 because the call to
<VirtAddr as Step>::steps_between would fail.
Never scale usize values, instead, always scale u64 values. To make
this easier to work with, this patch adds variants of the Step
functions to VirtAddr that take and return u64 instead of usize.
This patch also adds some regression tests.

Author: Freax13

src/addr.rs

@@ -240,30 +240,43 @@ impl VirtAddr {
     }
 
     // FIXME: Move this into the `Step` impl, once `Step` is stabilized.
-    #[cfg(any(feature = "instructions", feature = "step_trait"))]
+    #[cfg(feature = "step_trait")]
     pub(crate) fn steps_between_impl(start: &Self, end: &Self) -> (usize, Option<usize>) {
-        let mut steps = if let Some(steps) = end.0.checked_sub(start.0) {
-            steps
+        if let Some(steps) = Self::steps_between_u64(start, end) {
+            let steps = usize::try_from(steps).ok();
+            (steps.unwrap_or(!0), steps)
         } else {
-            return (0, None);
-        };
+            (0, None)
+        }
+    }
+
+    /// An implementation of steps_between that returns u64. Note that this
+    /// function always returns the exact bound, so it doesn't need to return a
+    /// lower and upper bound like steps_between does.
+    #[cfg(any(feature = "instructions", feature = "step_trait"))]
+    pub(crate) fn steps_between_u64(start: &Self, end: &Self) -> Option<u64> {
+        let mut steps = end.0.checked_sub(start.0)?;
 
         // Mask away extra bits that appear while jumping the gap.
         steps &= 0xffff_ffff_ffff;
 
-        let steps = usize::try_from(steps).ok();
-        (steps.unwrap_or(!0), steps)
+        Some(steps)
     }
 
     // FIXME: Move this into the `Step` impl, once `Step` is stabilized.
     #[inline]
     pub(crate) fn forward_checked_impl(start: Self, count: usize) -> Option<Self> {
-        let offset = u64::try_from(count).ok()?;
-        if offset > ADDRESS_SPACE_SIZE {
+        Self::forward_checked_u64(start, u64::try_from(count).ok()?)
+    }
+
+    /// An implementation of forward_checked that takes u64 instead of usize.
+    #[inline]
+    pub(crate) fn forward_checked_u64(start: Self, count: u64) -> Option<Self> {
+        if count > ADDRESS_SPACE_SIZE {
             return None;
         }
 
-        let mut addr = start.0.checked_add(offset)?;
+        let mut addr = start.0.checked_add(count)?;
 
         match addr.get_bits(47..) {
             0x1 => {
@@ -279,6 +292,31 @@ impl VirtAddr {
 
         Some(unsafe { Self::new_unsafe(addr) })
     }
+
+    /// An implementation of backward_checked that takes u64 instead of usize.
+    #[cfg(feature = "step_trait")]
+    #[inline]
+    pub(crate) fn backward_checked_u64(start: Self, count: u64) -> Option<Self> {
+        if count > ADDRESS_SPACE_SIZE {
+            return None;
+        }
+
+        let mut addr = start.0.checked_sub(count)?;
+
+        match addr.get_bits(47..) {
+            0x1fffe => {
+                // Jump the gap by sign extending the 47th bit.
+                addr.set_bits(47.., 0);
+            }
+            0x1fffd => {
+                // Address underflow
+                return None;
+            }
+            _ => {}
+        }
+
+        Some(unsafe { Self::new_unsafe(addr) })
+    }
 }
 
 impl fmt::Debug for VirtAddr {
@@ -376,26 +414,7 @@ impl Step for VirtAddr {
 
     #[inline]
     fn backward_checked(start: Self, count: usize) -> Option<Self> {
-        let offset = u64::try_from(count).ok()?;
-        if offset > ADDRESS_SPACE_SIZE {
-            return None;
-        }
-
-        let mut addr = start.0.checked_sub(offset)?;
-
-        match addr.get_bits(47..) {
-            0x1fffe => {
-                // Jump the gap by sign extending the 47th bit.
-                addr.set_bits(47.., 0);
-            }
-            0x1fffd => {
-                // Address underflow
-                return None;
-            }
-            _ => {}
-        }
-
-        Some(unsafe { Self::new_unsafe(addr) })
+        Self::backward_checked_u64(start, u64::try_from(count).ok()?)
     }
 }
 
@@ -779,6 +798,18 @@ mod tests {
             ),
             (0, None)
         );
+        // Make sure that we handle `steps * PAGE_SIZE > u32::MAX`
+        // correctly on 32-bit targets.
+        #[cfg(target_pointer_width = "64")]
+        assert_eq!(
+            Step::steps_between(&VirtAddr(0), &VirtAddr(0x1_0000_0000)),
+            (0x1_0000_0000, Some(0x1_0000_0000))
+        );
+        #[cfg(not(target_pointer_width = "64"))]
+        assert_eq!(
+            Step::steps_between(&VirtAddr(0), &VirtAddr(0x1_0000_0000)),
+            (!0, None)
+        );
     }
 
     #[test]

src/structures/paging/page.rs

@@ -161,17 +161,26 @@ impl<S: PageSize> Page<S> {
     // FIXME: Move this into the `Step` impl, once `Step` is stabilized.
     #[cfg(any(feature = "instructions", feature = "step_trait"))]
     pub(crate) fn steps_between_impl(start: &Self, end: &Self) -> (usize, Option<usize>) {
-        let (lower, upper) = VirtAddr::steps_between_impl(&start.start_address, &end.start_address);
-        let lower = lower / S::SIZE as usize;
-        let upper = upper.map(|steps| steps / S::SIZE as usize);
-        (lower, upper)
+        use core::convert::TryFrom;
+
+        if let Some(steps) =
+            VirtAddr::steps_between_u64(&start.start_address(), &end.start_address())
+        {
+            let steps = steps / S::SIZE;
+            let steps = usize::try_from(steps).ok();
+            (steps.unwrap_or(!0), steps)
+        } else {
+            (0, None)
+        }
     }
 
     // FIXME: Move this into the `Step` impl, once `Step` is stabilized.
     #[cfg(any(feature = "instructions", feature = "step_trait"))]
     pub(crate) fn forward_checked_impl(start: Self, count: usize) -> Option<Self> {
-        let count = count.checked_mul(S::SIZE as usize)?;
-        let start_address = VirtAddr::forward_checked_impl(start.start_address, count)?;
+        use core::convert::TryFrom;
+
+        let count = u64::try_from(count).ok()?.checked_mul(S::SIZE)?;
+        let start_address = VirtAddr::forward_checked_u64(start.start_address, count)?;
         Some(Self {
             start_address,
             size: PhantomData,
@@ -304,8 +313,10 @@ impl<S: PageSize> Step for Page<S> {
     }
 
     fn backward_checked(start: Self, count: usize) -> Option<Self> {
-        let count = count.checked_mul(S::SIZE as usize)?;
-        let start_address = Step::backward_checked(start.start_address, count)?;
+        use core::convert::TryFrom;
+
+        let count = u64::try_from(count).ok()?.checked_mul(S::SIZE)?;
+        let start_address = VirtAddr::backward_checked_u64(start.start_address, count)?;
         Some(Self {
             start_address,
             size: PhantomData,
@@ -531,4 +542,106 @@ mod tests {
         let range_inclusive = PageRangeInclusive { start, end };
         assert_eq!(range_inclusive.len(), 51);
     }
+
+    #[test]
+    #[cfg(feature = "step_trait")]
+    fn page_step_forward() {
+        let test_cases = [
+            (0, 0, Some(0)),
+            (0, 1, Some(0x1000)),
+            (0x1000, 1, Some(0x2000)),
+            (0x7fff_ffff_f000, 1, Some(0xffff_8000_0000_0000)),
+            (0xffff_8000_0000_0000, 1, Some(0xffff_8000_0000_1000)),
+            (0xffff_ffff_ffff_f000, 1, None),
+            #[cfg(target_pointer_width = "64")]
+            (0x7fff_ffff_f000, 0x1_2345_6789, Some(0xffff_9234_5678_8000)),
+            #[cfg(target_pointer_width = "64")]
+            (0x7fff_ffff_f000, 0x8_0000_0000, Some(0xffff_ffff_ffff_f000)),
+            #[cfg(target_pointer_width = "64")]
+            (0x7fff_fff0_0000, 0x8_0000_00ff, Some(0xffff_ffff_ffff_f000)),
+            #[cfg(target_pointer_width = "64")]
+            (0x7fff_fff0_0000, 0x8_0000_0100, None),
+            #[cfg(target_pointer_width = "64")]
+            (0x7fff_ffff_f000, 0x8_0000_0001, None),
+            // Make sure that we handle `steps * PAGE_SIZE > u32::MAX`
+            // correctly on 32-bit targets.
+            (0, 0x10_0000, Some(0x1_0000_0000)),
+        ];
+        for (start, count, result) in test_cases {
+            let start = Page::<Size4KiB>::from_start_address(VirtAddr::new(start)).unwrap();
+            let result = result
+                .map(|result| Page::<Size4KiB>::from_start_address(VirtAddr::new(result)).unwrap());
+            assert_eq!(Step::forward_checked(start, count), result);
+        }
+    }
+
+    #[test]
+    #[cfg(feature = "step_trait")]
+    fn page_step_backwards() {
+        let test_cases = [
+            (0, 0, Some(0)),
+            (0, 1, None),
+            (0x1000, 1, Some(0)),
+            (0xffff_8000_0000_0000, 1, Some(0x7fff_ffff_f000)),
+            (0xffff_8000_0000_1000, 1, Some(0xffff_8000_0000_0000)),
+            #[cfg(target_pointer_width = "64")]
+            (0xffff_9234_5678_8000, 0x1_2345_6789, Some(0x7fff_ffff_f000)),
+            #[cfg(target_pointer_width = "64")]
+            (0xffff_8000_0000_0000, 0x8_0000_0000, Some(0)),
+            #[cfg(target_pointer_width = "64")]
+            (0xffff_8000_0000_0000, 0x7_ffff_ff01, Some(0xff000)),
+            #[cfg(target_pointer_width = "64")]
+            (0xffff_8000_0000_0000, 0x8_0000_0001, None),
+            // Make sure that we handle `steps * PAGE_SIZE > u32::MAX`
+            // correctly on 32-bit targets.
+            (0x1_0000_0000, 0x10_0000, Some(0)),
+        ];
+        for (start, count, result) in test_cases {
+            let start = Page::<Size4KiB>::from_start_address(VirtAddr::new(start)).unwrap();
+            let result = result
+                .map(|result| Page::<Size4KiB>::from_start_address(VirtAddr::new(result)).unwrap());
+            assert_eq!(Step::backward_checked(start, count), result);
+        }
+    }
+
+    #[test]
+    #[cfg(feature = "step_trait")]
+    fn page_steps_between() {
+        let test_cases = [
+            (0, 0, 0, Some(0)),
+            (0, 0x1000, 1, Some(1)),
+            (0x1000, 0, 0, None),
+            (0x1000, 0x1000, 0, Some(0)),
+            (0x7fff_ffff_f000, 0xffff_8000_0000_0000, 1, Some(1)),
+            (0xffff_8000_0000_0000, 0x7fff_ffff_f000, 0, None),
+            (0xffff_8000_0000_0000, 0xffff_8000_0000_0000, 0, Some(0)),
+            (0xffff_8000_0000_0000, 0xffff_8000_0000_1000, 1, Some(1)),
+            (0xffff_8000_0000_1000, 0xffff_8000_0000_0000, 0, None),
+            (0xffff_8000_0000_1000, 0xffff_8000_0000_1000, 0, Some(0)),
+            // Make sure that we handle `steps > u32::MAX` correctly on 32-bit
+            // targets.
+            (
+                0x0000_0000_0000,
+                0x0001_0000_0000,
+                0x10_0000,
+                Some(0x10_0000),
+            ),
+            // The returned bounds are different when `steps` doesn't fit in
+            // into `usize`.
+            #[cfg(target_pointer_width = "64")]
+            (
+                0x0000_0000_0000,
+                0x1000_0000_0000,
+                0x1_0000_0000,
+                Some(0x1_0000_0000),
+            ),
+            #[cfg(not(target_pointer_width = "64"))]
+            (0x0000_0000_0000, 0x1000_0000_0000, !0, None),
+        ];
+        for (start, end, lower, upper) in test_cases {
+            let start = Page::<Size4KiB>::from_start_address(VirtAddr::new(start)).unwrap();
+            let end = Page::from_start_address(VirtAddr::new(end)).unwrap();
+            assert_eq!(Step::steps_between(&start, &end), (lower, upper));
+        }
+    }
 }