implement Step for VirtAddr

Author: Freax13

Cargo.toml

@@ -33,9 +33,10 @@ volatile = "0.4.4"
 [features]
 default = [ "nightly", "instructions" ]
 instructions = []
-nightly = [ "const_fn", "abi_x86_interrupt", "doc_cfg" ]
+nightly = [ "const_fn", "step_trait", "abi_x86_interrupt", "doc_cfg" ]
 abi_x86_interrupt = []
 const_fn = []
+step_trait = []
 doc_cfg = []
 
 # These features are no longer used and only there for backwards compatibility.

src/addr.rs

@@ -1,6 +1,10 @@
 //! Physical and virtual addresses manipulation
 
+#[cfg(feature = "step_trait")]
+use core::convert::TryFrom;
 use core::fmt;
+#[cfg(feature = "step_trait")]
+use core::iter::Step;
 use core::ops::{Add, AddAssign, Sub, SubAssign};
 
 use crate::structures::paging::page_table::PageTableLevel;
@@ -322,6 +326,44 @@ impl Sub<VirtAddr> for VirtAddr {
     }
 }
 
+#[cfg(feature = "step_trait")]
+impl Step for VirtAddr {
+    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
+        let mut steps = end.0.checked_sub(start.0)?;
+
+        // Check if we jumped the gap.
+        if end.0.get_bit(47) && !start.0.get_bit(47) {
+            steps -= 0xffff_0000_0000_0000;
+        }
+
+        usize::try_from(steps).ok()
+    }
+
+    fn forward_checked(start: Self, count: usize) -> Option<Self> {
+        let offset = u64::try_from(count).ok()?;
+        let mut addr = start.0.checked_add(offset)?;
+
+        // Jump the gap by sign extending the 47th bit.
+        if addr.get_bits(47..) == 0x1 {
+            addr.set_bits(47.., 0x1ffff);
+        }
+
+        Some(Self(addr))
+    }
+
+    fn backward_checked(start: Self, count: usize) -> Option<Self> {
+        let offset = u64::try_from(count).ok()?;
+        let mut addr = start.0.checked_sub(offset)?;
+
+        // Jump the gap by sign extending the 47th bit.
+        if addr.get_bits(47..) == 0x1fffe {
+            addr.set_bits(47.., 0);
+        }
+
+        Some(Self(addr))
+    }
+}
+
 /// A passed `u64` was not a valid physical address.
 ///
 /// This means that bits 52 to 64 were not all null.
@@ -577,6 +619,76 @@ mod tests {
         assert_eq!(VirtAddr::new_truncate(123 << 47), VirtAddr(0xfffff << 47));
     }
 
+    #[test]
+    #[cfg(feature = "step_trait")]
+    fn virtaddr_step() {
+        assert_eq!(Step::forward(VirtAddr(0), 0), VirtAddr(0));
+        assert_eq!(Step::forward(VirtAddr(0), 1), VirtAddr(1));
+        assert_eq!(
+            Step::forward(VirtAddr(0x7fff_ffff_ffff), 1),
+            VirtAddr(0xffff_8000_0000_0000)
+        );
+        assert_eq!(
+            Step::forward(VirtAddr(0xffff_8000_0000_0000), 1),
+            VirtAddr(0xffff_8000_0000_0001)
+        );
+        assert_eq!(
+            Step::forward_checked(VirtAddr(0xffff_ffff_ffff_ffff), 1),
+            None
+        );
+
+        assert_eq!(Step::backward(VirtAddr(0), 0), VirtAddr(0));
+        assert_eq!(Step::backward_checked(VirtAddr(0), 1), None);
+        assert_eq!(Step::backward(VirtAddr(1), 1), VirtAddr(0));
+        assert_eq!(
+            Step::backward(VirtAddr(0xffff_8000_0000_0000), 1),
+            VirtAddr(0x7fff_ffff_ffff)
+        );
+        assert_eq!(
+            Step::backward(VirtAddr(0xffff_8000_0000_0001), 1),
+            VirtAddr(0xffff_8000_0000_0000)
+        );
+
+        assert_eq!(Step::steps_between(&VirtAddr(0), &VirtAddr(0)), Some(0));
+        assert_eq!(Step::steps_between(&VirtAddr(0), &VirtAddr(1)), Some(1));
+        assert_eq!(Step::steps_between(&VirtAddr(1), &VirtAddr(0)), None);
+        assert_eq!(
+            Step::steps_between(
+                &VirtAddr(0x7fff_ffff_ffff),
+                &VirtAddr(0xffff_8000_0000_0000)
+            ),
+            Some(1)
+        );
+        assert_eq!(
+            Step::steps_between(
+                &VirtAddr(0xffff_8000_0000_0000),
+                &VirtAddr(0x7fff_ffff_ffff)
+            ),
+            None
+        );
+        assert_eq!(
+            Step::steps_between(
+                &VirtAddr(0xffff_8000_0000_0000),
+                &VirtAddr(0xffff_8000_0000_0000)
+            ),
+            Some(0)
+        );
+        assert_eq!(
+            Step::steps_between(
+                &VirtAddr(0xffff_8000_0000_0000),
+                &VirtAddr(0xffff_8000_0000_0001)
+            ),
+            Some(1)
+        );
+        assert_eq!(
+            Step::steps_between(
+                &VirtAddr(0xffff_8000_0000_0001),
+                &VirtAddr(0xffff_8000_0000_0000)
+            ),
+            None
+        );
+    }
+
     #[test]
     pub fn test_align_up() {
         // align 1

src/lib.rs

@@ -6,6 +6,7 @@
 #![cfg_attr(feature = "const_fn", feature(const_fn_fn_ptr_basics))] // IDT new()
 #![cfg_attr(feature = "const_fn", feature(const_fn_trait_bound))] // PageSize marker trait
 #![cfg_attr(feature = "abi_x86_interrupt", feature(abi_x86_interrupt))]
+#![cfg_attr(feature = "step_trait", feature(step_trait))]
 #![cfg_attr(feature = "doc_cfg", feature(doc_cfg))]
 #![warn(missing_docs)]
 #![deny(missing_debug_implementations)]