implement DoubleEndedIterator for ranges + implement more iterator methods

src/addr.rs

@@ -259,7 +259,6 @@ impl VirtAddr {
     /// 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)?;
 
@@ -488,6 +487,13 @@ impl Step for VirtAddr {
     }
 }
 
+#[cfg(kani)]
+impl kani::Arbitrary for VirtAddr {
+    fn any() -> Self {
+        Self::new_truncate(kani::any())
+    }
+}
+
 /// A passed `u64` was not a valid physical address.
 ///
 /// This means that bits 52 to 64 were not all null.
@@ -993,10 +999,8 @@ mod proofs {
     // step starting from any address.
     #[kani::proof]
     fn forward_base_case() {
-        let start_raw: u64 = kani::any();
-        let Ok(start) = VirtAddr::try_new(start_raw) else {
-            return;
-        };
+        let start = kani::any::<VirtAddr>();
+        let start_raw = start.as_u64();
 
         // Adding 0 to any address should always yield the same address.
         let same = Step::forward(start, 0);
@@ -1030,10 +1034,7 @@ mod proofs {
     // same as taking one combined large step.
     #[kani::proof]
     fn forward_induction_step() {
-        let start_raw: u64 = kani::any();
-        let Ok(start) = VirtAddr::try_new(start_raw) else {
-            return;
-        };
+        let start = kani::any::<VirtAddr>();
 
         let count1: usize = kani::any();
         let count2: usize = kani::any();
@@ -1060,10 +1061,7 @@ mod proofs {
     // for all inputs for which `forward_checked` succeeds.
     #[kani::proof]
     fn forward_implies_backward() {
-        let start_raw: u64 = kani::any();
-        let Ok(start) = VirtAddr::try_new(start_raw) else {
-            return;
-        };
+        let start = kani::any::<VirtAddr>();
         let count: usize = kani::any();
 
         // If `forward_checked` succeeds...
@@ -1080,10 +1078,7 @@ mod proofs {
     // succeeds, `forward` succeeds as well.
     #[kani::proof]
     fn backward_implies_forward() {
-        let end_raw: u64 = kani::any();
-        let Ok(end) = VirtAddr::try_new(end_raw) else {
-            return;
-        };
+        let end = kani::any::<VirtAddr>();
         let count: usize = kani::any();
 
         // If `backward_checked` succeeds...
@@ -1105,10 +1100,7 @@ mod proofs {
     // `steps_between` for all inputs for which `forward_checked` succeeds.
     #[kani::proof]
     fn forward_implies_steps_between() {
-        let start: u64 = kani::any();
-        let Ok(start) = VirtAddr::try_new(start) else {
-            return;
-        };
+        let start = kani::any::<VirtAddr>();
         let count: usize = kani::any();
 
         // If `forward_checked` succeeds...
@@ -1124,14 +1116,8 @@ mod proofs {
     // succeeds, `forward` succeeds as well.
     #[kani::proof]
     fn steps_between_implies_forward() {
-        let start: u64 = kani::any();
-        let Ok(start) = VirtAddr::try_new(start) else {
-            return;
-        };
-        let end: u64 = kani::any();
-        let Ok(end) = VirtAddr::try_new(end) else {
-            return;
-        };
+        let start = kani::any::<VirtAddr>();
+        let end = kani::any::<VirtAddr>();
 
         // If `steps_between` succeeds...
         let Some(count) = Step::steps_between(&start, &end).1 else {

src/structures/paging/frame.rs

@@ -3,6 +3,7 @@
 use super::page::AddressNotAligned;
 use crate::structures::paging::page::{PageSize, Size4KiB};
 use crate::PhysAddr;
+use core::convert::TryFrom;
 use core::fmt;
 use core::marker::PhantomData;
 use core::ops::{Add, AddAssign, Sub, SubAssign};
@@ -179,6 +180,71 @@ impl<S: PageSize> Iterator for PhysFrameRange<S> {
             None
         }
     }
+
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        if self.is_empty() {
+            return None;
+        }
+
+        // Convert to `u64`. If the value doesn't fit just use `u64::MAX`.
+        // `self.len()` is guaranteed to be smaller than the real value and
+        // `u64::MAX` anyway, so it doesn't make a difference.
+        let n = u64::try_from(n).unwrap_or(u64::MAX);
+
+        // Handling `n >= self.len()` is a bit more complicated because we
+        // can't just add `n` to `self.start` (it might overflow). Handle this
+        // by doing two steps, `self.len()-1` and `1`. This should return
+        // `None`.
+        if n >= self.len() {
+            self.nth(self.len() as usize - 1)?;
+            return self.next();
+        }
+
+        self.start += n;
+        self.next()
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.len();
+        usize::try_from(len)
+            .map(|len| (len, Some(len)))
+            .unwrap_or((usize::MAX, None))
+    }
+}
+
+impl<S: PageSize> DoubleEndedIterator for PhysFrameRange<S> {
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        if self.start < self.end {
+            self.end -= 1;
+            Some(self.end)
+        } else {
+            None
+        }
+    }
+
+    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
+        if self.is_empty() {
+            return None;
+        }
+
+        // Convert to `u64`. If the value doesn't fit just use `u64::MAX`.
+        // `self.len()` is guaranteed to be smaller than the real value and
+        // `u64::MAX` anyway, so it doesn't make a difference.
+        let n = u64::try_from(n).unwrap_or(u64::MAX);
+
+        // Handling `n >= self.len()` is a bit more complicated because we
+        // can't just subtract `n` to `self.end` (it might overflow). Handle
+        // this by doing two steps, `self.len()-1` and `1`. This should return
+        // `None`.
+        if n >= self.len() {
+            self.nth_back(self.len() as usize - 1)?;
+            return self.next_back();
+        }
+
+        self.end -= n;
+        self.next_back()
+    }
 }
 
 impl<S: PageSize> fmt::Debug for PhysFrameRange<S> {
@@ -237,6 +303,71 @@ impl<S: PageSize> Iterator for PhysFrameRangeInclusive<S> {
             None
         }
     }
+
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        if self.is_empty() {
+            return None;
+        }
+
+        // Convert to `u64`. If the value doesn't fit just use `u64::MAX`.
+        // `self.len()` is guaranteed to be smaller than the real value and
+        // `u64::MAX` anyway, so it doesn't make a difference.
+        let n = u64::try_from(n).unwrap_or(u64::MAX);
+
+        // Handling `n >= self.len()` is a bit more complicated because we
+        // can't just add `n` to `self.start` (it might overflow). Handle this
+        // by doing two steps, `self.len()-1` and `1`. This should return
+        // `None`.
+        if n >= self.len() {
+            self.nth(self.len() as usize - 1)?;
+            return self.next();
+        }
+
+        self.start += n;
+        self.next()
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.len();
+        usize::try_from(len)
+            .map(|len| (len, Some(len)))
+            .unwrap_or((usize::MAX, None))
+    }
+}
+
+impl<S: PageSize> DoubleEndedIterator for PhysFrameRangeInclusive<S> {
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        if self.start <= self.end {
+            self.end -= 1;
+            Some(self.end)
+        } else {
+            None
+        }
+    }
+
+    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
+        if self.is_empty() {
+            return None;
+        }
+
+        // Convert to `u64`. If the value doesn't fit just use `u64::MAX`.
+        // `self.len()` is guaranteed to be smaller than the real value and
+        // `u64::MAX` anyway, so it doesn't make a difference.
+        let n = u64::try_from(n).unwrap_or(u64::MAX);
+
+        // Handling `n >= self.len()` is a bit more complicated because we
+        // can't just subtract `n` to `self.end` (it might overflow). Handle
+        // this by doing two steps, `self.len()-1` and `1`. This should return
+        // `None`.
+        if n >= self.len() {
+            self.nth_back(self.len() as usize - 1)?;
+            return self.next_back();
+        }
+
+        self.end -= n;
+        self.next_back()
+    }
 }
 
 impl<S: PageSize> fmt::Debug for PhysFrameRangeInclusive<S> {