Add physical and virtual address types

Create wrapper types for physical and virtual address to limit available
operations on addresses and to be able to require explicit address types
for given function parameters.

Signed-off-by: Imre Kis <imre.kis@arm.com>
Change-Id: Iaef5ab1af24fc153d959d79404b3827d9c85bf53
diff --git a/src/region.rs b/src/region.rs
index 34428a2..83c54ab 100644
--- a/src/region.rs
+++ b/src/region.rs
@@ -9,6 +9,7 @@
 use log::debug;
 
 use super::{
+    address::{PhysicalAddress, VirtualAddress},
     page_pool::{PagePool, Pages},
     region_pool::Region,
 };
@@ -22,12 +23,12 @@
 pub enum PhysicalRegion {
     Unused,
     Allocated(PagePool, Pages),
-    PhysicalAddress(usize),
+    PhysicalAddress(PhysicalAddress),
 }
 
 impl PhysicalRegion {
     /// Get physical memory address
-    fn get_pa(&self) -> usize {
+    fn get_pa(&self) -> PhysicalAddress {
         match self {
             PhysicalRegion::Unused => panic!("Unused area has no PA"),
             PhysicalRegion::Allocated(_page_pool, pages) => pages.get_pa(),
@@ -40,24 +41,24 @@
 ///
 /// A virtual memory region has a virtual address, a length and a physical region.
 pub struct VirtualRegion {
-    va: usize,
+    va: VirtualAddress,
     length: usize,
     physical_region: PhysicalRegion,
 }
 
 impl VirtualRegion {
     /// Create new virtual memory region without a physical region
-    pub fn new(va: usize, length: usize) -> Self {
+    pub fn new(va: VirtualAddress, length: usize) -> Self {
         Self::new_from_fields(va, length, PhysicalRegion::Unused)
     }
 
     /// Create virtual region with points to a given physical address
-    pub fn new_with_pa(pa: usize, va: usize, length: usize) -> Self {
+    pub fn new_with_pa(pa: PhysicalAddress, va: VirtualAddress, length: usize) -> Self {
         Self::new_from_fields(va, length, PhysicalRegion::PhysicalAddress(pa))
     }
 
     /// Create virtual region by defining all the fields of the object
-    fn new_from_fields(va: usize, length: usize, physical_region: PhysicalRegion) -> Self {
+    fn new_from_fields(va: VirtualAddress, length: usize, physical_region: PhysicalRegion) -> Self {
         Self {
             va,
             length,
@@ -66,29 +67,29 @@
     }
 
     /// Get the base address of the linked physical region
-    pub fn get_pa(&self) -> usize {
+    pub fn get_pa(&self) -> PhysicalAddress {
         self.physical_region.get_pa()
     }
 
     /// Get physical address for a virtual address
-    pub fn get_pa_for_va(&self, va: usize) -> usize {
-        let offset = va.checked_sub(self.va).unwrap();
+    pub fn get_pa_for_va(&self, va: VirtualAddress) -> PhysicalAddress {
+        let offset = va.diff(self.va).unwrap();
 
         assert!(offset < self.length);
-        self.get_pa().checked_add(offset).unwrap()
+        self.get_pa().add_offset(offset).unwrap()
     }
 }
 
 impl Region for VirtualRegion {
     type Resource = PhysicalRegion;
-    type Base = usize;
+    type Base = VirtualAddress;
     type Length = usize;
 
-    fn base(&self) -> usize {
+    fn base(&self) -> Self::Base {
         self.va
     }
 
-    fn length(&self) -> usize {
+    fn length(&self) -> Self::Length {
         self.length
     }
 
@@ -96,9 +97,9 @@
         !matches!(self.physical_region, PhysicalRegion::Unused)
     }
 
-    fn contains(&self, base: usize, length: usize) -> bool {
+    fn contains(&self, base: Self::Base, length: Self::Length) -> bool {
         if let (Some(end), Some(self_end)) =
-            (base.checked_add(length), self.va.checked_add(self.length))
+            (base.add_offset(length), self.va.add_offset(self.length))
         {
             self.va <= base && end <= self_end
         } else {
@@ -108,7 +109,7 @@
 
     fn try_append(&mut self, other: &Self) -> bool {
         if let (Some(self_end), Some(new_length)) = (
-            self.va.checked_add(self.length),
+            self.va.add_offset(self.length),
             self.length.checked_add(other.length),
         ) {
             if self_end == other.va {
@@ -125,7 +126,7 @@
                     ) => {
                         // Used ranges can be only merged if the PA doesn't overflow and it is
                         // consecutive
-                        if let Some(self_end_pa) = self_pa.checked_add(self.length) {
+                        if let Some(self_end_pa) = self_pa.add_offset(self.length) {
                             if self_end_pa == *other_pa {
                                 self.length = new_length;
                                 return true;
@@ -146,8 +147,8 @@
 
     fn create_split(
         &self,
-        base: usize,
-        length: usize,
+        base: Self::Base,
+        length: Self::Length,
         resource: Option<Self::Resource>,
     ) -> (Self, Vec<Self>) {
         assert!(self.contains(base, length));
@@ -159,15 +160,15 @@
 
             let mut res = Vec::new();
             if self.va != base {
-                res.push(Self::new(self.va, base.checked_sub(self.va).unwrap()));
+                res.push(Self::new(self.va, base.diff(self.va).unwrap()));
             }
 
             res.push(Self::new_from_fields(base, length, physical_region));
 
-            let end = base.checked_add(length).unwrap();
-            let self_end = self.va.checked_add(self.length).unwrap();
+            let end = base.add_offset(length).unwrap();
+            let self_end = self.va.add_offset(self.length).unwrap();
             if end != self_end {
-                res.push(Self::new(end, self_end.checked_sub(end).unwrap()));
+                res.push(Self::new(end, self_end.diff(end).unwrap()));
             }
 
             (
@@ -190,23 +191,23 @@
 
                 res.push(Self::new_from_fields(
                     self.va,
-                    base.checked_sub(self.va).unwrap(),
+                    base.diff(self.va).unwrap(),
                     physical_region,
                 ));
             }
 
             res.push(Self::new(base, length));
 
-            let end = base.checked_add(length).unwrap();
-            let self_end = self.va.checked_add(self.length).unwrap();
+            let end = base.add_offset(length).unwrap();
+            let self_end = self.va.add_offset(self.length).unwrap();
             if end != self_end {
                 let physical_region = match &self.physical_region {
                     PhysicalRegion::Allocated(_page_pool, _pages) => {
                         todo!("Implement Pages::split");
                     }
                     PhysicalRegion::PhysicalAddress(pa) => {
-                        let offset = end.checked_sub(self.va).unwrap();
-                        PhysicalRegion::PhysicalAddress(pa.checked_add(offset).unwrap())
+                        let offset = end.diff(self.va).unwrap();
+                        PhysicalRegion::PhysicalAddress(pa.add_offset(offset).unwrap())
                     }
                     _ => {
                         panic!("Splitting unused region by other unused")
@@ -215,7 +216,7 @@
 
                 res.push(Self::new_from_fields(
                     end,
-                    self_end.checked_sub(end).unwrap(),
+                    self_end.diff(end).unwrap(),
                     physical_region,
                 ));
             }
@@ -235,9 +236,9 @@
 
         if let PhysicalRegion::Allocated(page_pool, pages) = physical_region {
             debug!(
-                "Dropping physical region with pages: PA={:#010x} VA={:#010x}",
-                pages.get_pa(),
-                self.base(),
+                "Dropping physical region with pages: PA={:#08x} VA={:#08x}",
+                pages.get_pa().0,
+                self.base().0,
             );
 
             page_pool.release_pages(pages).unwrap();
@@ -247,8 +248,8 @@
 
 #[cfg(test)]
 mod tests {
-    use super::super::page_pool::PagePoolArea;
     use super::*;
+    use crate::page_pool::PagePoolArea;
 
     #[test]
     #[should_panic]
@@ -259,12 +260,14 @@
 
     #[test]
     fn test_physical_region() {
-        const PA: usize = 0x0123_4567_89ab_cdef;
+        const PA: PhysicalAddress = PhysicalAddress(0x0123_4567_89ab_cdef);
         const LENGTH: usize = 0x8000_0000_0000;
 
         static PAGE_POOL_AREA: PagePoolArea<16> = PagePoolArea::new();
-        let region =
-            PhysicalRegion::Allocated(PagePool::new(&PAGE_POOL_AREA), Pages::new(PA, LENGTH, true));
+        let region = PhysicalRegion::Allocated(
+            PagePool::new(&PAGE_POOL_AREA),
+            Pages::new(PA.0, LENGTH, true),
+        );
         assert_eq!(PA, region.get_pa());
 
         let region = PhysicalRegion::PhysicalAddress(PA);
@@ -273,8 +276,8 @@
 
     #[test]
     fn test_virtual_region() {
-        const VA: usize = 0x0123_4567_89ab_cdef;
-        const PA: usize = 0xfedc_ba98_7654_3210;
+        const VA: VirtualAddress = VirtualAddress(0x0123_4567_89ab_cdef);
+        const PA: PhysicalAddress = PhysicalAddress(0xfedc_ba98_7654_3210);
         const LENGTH: usize = 0x8000_0000_0000;
 
         let region = VirtualRegion::new(VA, LENGTH);
@@ -300,36 +303,45 @@
 
     #[test]
     fn test_virtual_region_get_pa_for_va() {
-        let region =
-            VirtualRegion::new_with_pa(0x8000_0000_0000_0000, 0x4000_0000_0000_0000, 0x1000);
-        assert_eq!(
-            0x8000_0000_0000_0000,
-            region.get_pa_for_va(0x4000_0000_0000_0000)
+        let region = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x8000_0000_0000_0000),
+            VirtualAddress(0x4000_0000_0000_0000),
+            0x1000,
         );
         assert_eq!(
-            0x8000_0000_0000_0001,
-            region.get_pa_for_va(0x4000_0000_0000_0001)
+            PhysicalAddress(0x8000_0000_0000_0000),
+            region.get_pa_for_va(VirtualAddress(0x4000_0000_0000_0000))
         );
         assert_eq!(
-            0x8000_0000_0000_0fff,
-            region.get_pa_for_va(0x4000_0000_0000_0fff)
+            PhysicalAddress(0x8000_0000_0000_0001),
+            region.get_pa_for_va(VirtualAddress(0x4000_0000_0000_0001))
+        );
+        assert_eq!(
+            PhysicalAddress(0x8000_0000_0000_0fff),
+            region.get_pa_for_va(VirtualAddress(0x4000_0000_0000_0fff))
         );
     }
 
     #[test]
     #[should_panic]
     fn test_virtual_region_get_pa_for_va_low_va() {
-        let region =
-            VirtualRegion::new_with_pa(0x8000_0000_0000_0000, 0x4000_0000_0000_0000, 0x1000);
-        region.get_pa_for_va(0x3fff_ffff_ffff_ffff);
+        let region = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x8000_0000_0000_0000),
+            VirtualAddress(0x4000_0000_0000_0000),
+            0x1000,
+        );
+        region.get_pa_for_va(VirtualAddress(0x3fff_ffff_ffff_ffff));
     }
 
     #[test]
     #[should_panic]
     fn test_virtual_region_get_pa_for_va_high_va() {
-        let region =
-            VirtualRegion::new_with_pa(0x8000_0000_0000_0000, 0x4000_0000_0000_0000, 0x1000);
-        region.get_pa_for_va(0x4000_0000_0000_1000);
+        let region = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x8000_0000_0000_0000),
+            VirtualAddress(0x4000_0000_0000_0000),
+            0x1000,
+        );
+        region.get_pa_for_va(VirtualAddress(0x4000_0000_0000_1000));
     }
 
     #[test]
@@ -337,120 +349,151 @@
         const VA: usize = 0x8000_0000_0000_0000;
         const LENGTH: usize = 0x8000_0000_0000;
 
-        let region_overflow_end = VirtualRegion::new(0x8000_0000_0000_0000, 0x8000_0000_0000_0000);
-        assert!(!region_overflow_end.contains(0x8000_0000_0000_0000, 1));
+        let region_overflow_end =
+            VirtualRegion::new(VirtualAddress(0x8000_0000_0000_0000), 0x8000_0000_0000_0000);
+        assert!(!region_overflow_end.contains(VirtualAddress(0x8000_0000_0000_0000), 1));
 
-        let region = VirtualRegion::new(0x4000_0000_0000_0000, 0x8000_0000_0000_0000);
-        assert!(!region.contains(0x8000_0000_0000_0000, 0x8000_0000_0000_0000));
+        let region =
+            VirtualRegion::new(VirtualAddress(0x4000_0000_0000_0000), 0x8000_0000_0000_0000);
+        assert!(!region.contains(VirtualAddress(0x8000_0000_0000_0000), 0x8000_0000_0000_0000));
 
-        assert!(!region.contains(0x4000_0000_0000_0000, 0x8000_0000_0000_0001));
-        assert!(!region.contains(0x3fff_ffff_ffff_ffff, 0x8000_0000_0000_0000));
-        assert!(region.contains(0x4000_0000_0000_0000, 0x8000_0000_0000_0000));
-        assert!(region.contains(0x4000_0000_0000_0000, 0x7fff_ffff_ffff_ffff));
-        assert!(region.contains(0x4000_0000_0000_0001, 0x7fff_ffff_ffff_ffff));
+        assert!(!region.contains(VirtualAddress(0x4000_0000_0000_0000), 0x8000_0000_0000_0001));
+        assert!(!region.contains(VirtualAddress(0x3fff_ffff_ffff_ffff), 0x8000_0000_0000_0000));
+        assert!(region.contains(VirtualAddress(0x4000_0000_0000_0000), 0x8000_0000_0000_0000));
+        assert!(region.contains(VirtualAddress(0x4000_0000_0000_0000), 0x7fff_ffff_ffff_ffff));
+        assert!(region.contains(VirtualAddress(0x4000_0000_0000_0001), 0x7fff_ffff_ffff_ffff));
     }
 
     #[test]
     fn test_virtual_region_try_append() {
         // Both unused
-        let mut region_unused0 = VirtualRegion::new(0x4000_0000, 0x1000);
-        let mut region_unused1 = VirtualRegion::new(0x4000_1000, 0x1000);
+        let mut region_unused0 = VirtualRegion::new(VirtualAddress(0x4000_0000), 0x1000);
+        let mut region_unused1 = VirtualRegion::new(VirtualAddress(0x4000_1000), 0x1000);
 
         assert!(!region_unused1.try_append(&region_unused0));
-        assert_eq!(0x4000_0000, region_unused0.va);
+        assert_eq!(VirtualAddress(0x4000_0000), region_unused0.va);
         assert_eq!(0x1000, region_unused0.length);
-        assert_eq!(0x4000_1000, region_unused1.va);
+        assert_eq!(VirtualAddress(0x4000_1000), region_unused1.va);
         assert_eq!(0x1000, region_unused1.length);
 
         assert!(region_unused0.try_append(&region_unused1));
-        assert_eq!(0x4000_0000, region_unused0.va);
+        assert_eq!(VirtualAddress(0x4000_0000), region_unused0.va);
         assert_eq!(0x2000, region_unused0.length);
-        assert_eq!(0x4000_1000, region_unused1.va);
+        assert_eq!(VirtualAddress(0x4000_1000), region_unused1.va);
         assert_eq!(0x1000, region_unused1.length);
 
         // Unused and PA region
-        let mut region_unused = VirtualRegion::new(0x4000_0000, 0x1000);
-        let region_physical = VirtualRegion::new_with_pa(0x8000_0000, 0x4000_1000, 0x1000);
+        let mut region_unused = VirtualRegion::new(VirtualAddress(0x4000_0000), 0x1000);
+        let region_physical = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x8000_0000),
+            VirtualAddress(0x4000_1000),
+            0x1000,
+        );
         assert!(!region_unused.try_append(&region_physical));
-        assert_eq!(0x4000_0000, region_unused.va);
+        assert_eq!(VirtualAddress(0x4000_0000), region_unused.va);
         assert_eq!(0x1000, region_unused.length);
-        assert_eq!(0x4000_1000, region_physical.va);
+        assert_eq!(VirtualAddress(0x4000_1000), region_physical.va);
         assert_eq!(0x1000, region_physical.length);
 
         // Both PA regions but non-consecutive PA ranges
-        let mut region_physical0 = VirtualRegion::new_with_pa(0x8000_0000, 0x4000_0000, 0x1000);
-        let region_physical1 = VirtualRegion::new_with_pa(0x9000_0000, 0x4000_1000, 0x1000);
+        let mut region_physical0 = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x8000_0000),
+            VirtualAddress(0x4000_0000),
+            0x1000,
+        );
+        let region_physical1 = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x9000_0000),
+            VirtualAddress(0x4000_1000),
+            0x1000,
+        );
         assert!(!region_physical0.try_append(&region_physical1));
-        assert_eq!(0x4000_0000, region_physical0.va);
+        assert_eq!(VirtualAddress(0x4000_0000), region_physical0.va);
         assert_eq!(0x1000, region_physical0.length);
-        assert_eq!(0x4000_1000, region_physical1.va);
+        assert_eq!(VirtualAddress(0x4000_1000), region_physical1.va);
         assert_eq!(0x1000, region_physical1.length);
 
         // Both PA regions with consecutive PA ranges
-        let mut region_physical0 = VirtualRegion::new_with_pa(0x8000_0000, 0x4000_0000, 0x1000);
-        let region_physical1 = VirtualRegion::new_with_pa(0x8000_1000, 0x4000_1000, 0x1000);
+        let mut region_physical0 = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x8000_0000),
+            VirtualAddress(0x4000_0000),
+            0x1000,
+        );
+        let region_physical1 = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x8000_1000),
+            VirtualAddress(0x4000_1000),
+            0x1000,
+        );
         assert!(region_physical0.try_append(&region_physical1));
-        assert_eq!(0x4000_0000, region_physical0.va);
+        assert_eq!(VirtualAddress(0x4000_0000), region_physical0.va);
         assert_eq!(0x2000, region_physical0.length);
-        assert_eq!(0x4000_1000, region_physical1.va);
+        assert_eq!(VirtualAddress(0x4000_1000), region_physical1.va);
         assert_eq!(0x1000, region_physical1.length);
 
         // VA overflow
-        let mut region_unused0 = VirtualRegion::new(0x8000_0000_0000_0000, 0x8000_0000_0000_0000);
-        let mut region_unused1 = VirtualRegion::new(0x4000_1000, 0x1000);
+        let mut region_unused0: VirtualRegion =
+            VirtualRegion::new(VirtualAddress(0x8000_0000_0000_0000), 0x8000_0000_0000_0000);
+        let mut region_unused1 = VirtualRegion::new(VirtualAddress(0x4000_1000), 0x1000);
 
         assert!(!region_unused0.try_append(&region_unused1));
-        assert_eq!(0x8000_0000_0000_0000, region_unused0.va);
+        assert_eq!(VirtualAddress(0x8000_0000_0000_0000), region_unused0.va);
         assert_eq!(0x8000_0000_0000_0000, region_unused0.length);
-        assert_eq!(0x4000_1000, region_unused1.va);
+        assert_eq!(VirtualAddress(0x4000_1000), region_unused1.va);
         assert_eq!(0x1000, region_unused1.length);
 
         assert!(!region_unused1.try_append(&region_unused0));
-        assert_eq!(0x8000_0000_0000_0000, region_unused0.va);
+        assert_eq!(VirtualAddress(0x8000_0000_0000_0000), region_unused0.va);
         assert_eq!(0x8000_0000_0000_0000, region_unused0.length);
-        assert_eq!(0x4000_1000, region_unused1.va);
+        assert_eq!(VirtualAddress(0x4000_1000), region_unused1.va);
         assert_eq!(0x1000, region_unused1.length);
 
         // PA overflow
-        let mut region_physical0 =
-            VirtualRegion::new_with_pa(0x8000_0000_0000_0000, 0x4000_0000, 0x8000_0000_0000_0000);
-        let region_physical1 =
-            VirtualRegion::new_with_pa(0x9000_0000, 0x8000_0000_4000_0000, 0x1000);
+        let mut region_physical0 = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x8000_0000_0000_0000),
+            VirtualAddress(0x4000_0000),
+            0x8000_0000_0000_0000,
+        );
+        let region_physical1 = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x9000_0000),
+            VirtualAddress(0x8000_0000_4000_0000),
+            0x1000,
+        );
         assert!(!region_physical0.try_append(&region_physical1));
-        assert_eq!(0x4000_0000, region_physical0.va);
+        assert_eq!(VirtualAddress(0x4000_0000), region_physical0.va);
         assert_eq!(0x8000_0000_0000_0000, region_physical0.length);
-        assert_eq!(0x8000_0000_4000_0000, region_physical1.va);
+        assert_eq!(VirtualAddress(0x8000_0000_4000_0000), region_physical1.va);
         assert_eq!(0x1000, region_physical1.length);
     }
 
     #[test]
     fn test_virtual_region_create_split_by_used() {
-        let region_unused = VirtualRegion::new(0x4000_0000, 0x4000);
+        let region_unused = VirtualRegion::new(VirtualAddress(0x4000_0000), 0x4000);
 
         // New region at the start
         let (new_region, splitted_regions) = region_unused.create_split(
-            0x4000_0000,
+            VirtualAddress(0x4000_0000),
             0x1000,
-            Some(PhysicalRegion::PhysicalAddress(0x8000_0000)),
+            Some(PhysicalRegion::PhysicalAddress(PhysicalAddress(
+                0x8000_0000,
+            ))),
         );
 
-        assert_eq!(0x4000_0000, new_region.va);
+        assert_eq!(VirtualAddress(0x4000_0000), new_region.va);
         assert_eq!(0x1000, new_region.length);
-        assert_eq!(0x8000_0000, new_region.get_pa());
+        assert_eq!(PhysicalAddress(0x8000_0000), new_region.get_pa());
         assert!(matches!(
             new_region.physical_region,
             PhysicalRegion::PhysicalAddress(_)
         ));
 
-        assert_eq!(0x4000_0000, splitted_regions[0].va);
+        assert_eq!(VirtualAddress(0x4000_0000), splitted_regions[0].va);
         assert_eq!(0x1000, splitted_regions[0].length);
-        assert_eq!(0x8000_0000, splitted_regions[0].get_pa());
+        assert_eq!(PhysicalAddress(0x8000_0000), splitted_regions[0].get_pa());
         assert!(matches!(
             splitted_regions[0].physical_region,
             PhysicalRegion::PhysicalAddress(_)
         ));
 
-        assert_eq!(0x4000_1000, splitted_regions[1].va);
+        assert_eq!(VirtualAddress(0x4000_1000), splitted_regions[1].va);
         assert_eq!(0x3000, splitted_regions[1].length);
         assert!(matches!(
             splitted_regions[1].physical_region,
@@ -459,35 +502,37 @@
 
         // New region in the middle
         let (new_region, splitted_regions) = region_unused.create_split(
-            0x4000_1000,
+            VirtualAddress(0x4000_1000),
             0x1000,
-            Some(PhysicalRegion::PhysicalAddress(0x8000_0000)),
+            Some(PhysicalRegion::PhysicalAddress(PhysicalAddress(
+                0x8000_0000,
+            ))),
         );
 
-        assert_eq!(0x4000_1000, new_region.va);
+        assert_eq!(VirtualAddress(0x4000_1000), new_region.va);
         assert_eq!(0x1000, new_region.length);
-        assert_eq!(0x8000_0000, new_region.get_pa());
+        assert_eq!(PhysicalAddress(0x8000_0000), new_region.get_pa());
         assert!(matches!(
             new_region.physical_region,
             PhysicalRegion::PhysicalAddress(_)
         ));
 
-        assert_eq!(0x4000_0000, splitted_regions[0].va);
+        assert_eq!(VirtualAddress(0x4000_0000), splitted_regions[0].va);
         assert_eq!(0x1000, splitted_regions[0].length);
         assert!(matches!(
             splitted_regions[0].physical_region,
             PhysicalRegion::Unused
         ));
 
-        assert_eq!(0x4000_1000, splitted_regions[1].va);
+        assert_eq!(VirtualAddress(0x4000_1000), splitted_regions[1].va);
         assert_eq!(0x1000, splitted_regions[1].length);
-        assert_eq!(0x8000_0000, splitted_regions[1].get_pa());
+        assert_eq!(PhysicalAddress(0x8000_0000), splitted_regions[1].get_pa());
         assert!(matches!(
             splitted_regions[1].physical_region,
             PhysicalRegion::PhysicalAddress(_)
         ));
 
-        assert_eq!(0x4000_2000, splitted_regions[2].va);
+        assert_eq!(VirtualAddress(0x4000_2000), splitted_regions[2].va);
         assert_eq!(0x2000, splitted_regions[2].length);
         assert!(matches!(
             splitted_regions[2].physical_region,
@@ -496,29 +541,31 @@
 
         // New region at the end
         let (new_region, splitted_regions) = region_unused.create_split(
-            0x4000_3000,
+            VirtualAddress(0x4000_3000),
             0x1000,
-            Some(PhysicalRegion::PhysicalAddress(0x8000_0000)),
+            Some(PhysicalRegion::PhysicalAddress(PhysicalAddress(
+                0x8000_0000,
+            ))),
         );
 
-        assert_eq!(0x4000_3000, new_region.va);
+        assert_eq!(VirtualAddress(0x4000_3000), new_region.va);
         assert_eq!(0x1000, new_region.length);
-        assert_eq!(0x8000_0000, new_region.get_pa());
+        assert_eq!(PhysicalAddress(0x8000_0000), new_region.get_pa());
         assert!(matches!(
             new_region.physical_region,
             PhysicalRegion::PhysicalAddress(_)
         ));
 
-        assert_eq!(0x4000_0000, splitted_regions[0].va);
+        assert_eq!(VirtualAddress(0x4000_0000), splitted_regions[0].va);
         assert_eq!(0x3000, splitted_regions[0].length);
         assert!(matches!(
             splitted_regions[0].physical_region,
             PhysicalRegion::Unused
         ));
 
-        assert_eq!(0x4000_3000, splitted_regions[1].va);
+        assert_eq!(VirtualAddress(0x4000_3000), splitted_regions[1].va);
         assert_eq!(0x1000, splitted_regions[1].length);
-        assert_eq!(0x8000_0000, splitted_regions[1].get_pa());
+        assert_eq!(PhysicalAddress(0x8000_0000), splitted_regions[1].get_pa());
         assert!(matches!(
             splitted_regions[1].physical_region,
             PhysicalRegion::PhysicalAddress(_)
@@ -527,76 +574,83 @@
 
     #[test]
     fn test_virtual_region_create_split_by_unused() {
-        let region_unused = VirtualRegion::new_with_pa(0x8000_0000, 0x4000_0000, 0x4000);
+        let region_unused = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x8000_0000),
+            VirtualAddress(0x4000_0000),
+            0x4000,
+        );
 
         // New region at the start
-        let (new_region, splitted_regions) = region_unused.create_split(0x4000_0000, 0x1000, None);
+        let (new_region, splitted_regions) =
+            region_unused.create_split(VirtualAddress(0x4000_0000), 0x1000, None);
 
-        assert_eq!(0x4000_0000, new_region.va);
+        assert_eq!(VirtualAddress(0x4000_0000), new_region.va); // TODO: why do we need explicit type here?
         assert_eq!(0x1000, new_region.length);
         assert!(matches!(new_region.physical_region, PhysicalRegion::Unused));
 
-        assert_eq!(0x4000_0000, splitted_regions[0].va);
+        assert_eq!(VirtualAddress(0x4000_0000), splitted_regions[0].va);
         assert_eq!(0x1000, splitted_regions[0].length);
         assert!(matches!(
             splitted_regions[0].physical_region,
             PhysicalRegion::Unused
         ));
 
-        assert_eq!(0x4000_1000, splitted_regions[1].va);
+        assert_eq!(VirtualAddress(0x4000_1000), splitted_regions[1].va);
         assert_eq!(0x3000, splitted_regions[1].length);
-        assert_eq!(0x8000_1000, splitted_regions[1].get_pa());
+        assert_eq!(PhysicalAddress(0x8000_1000), splitted_regions[1].get_pa());
         assert!(matches!(
             splitted_regions[1].physical_region,
             PhysicalRegion::PhysicalAddress(_)
         ));
 
         // New region in the middle
-        let (new_region, splitted_regions) = region_unused.create_split(0x4000_1000, 0x1000, None);
+        let (new_region, splitted_regions) =
+            region_unused.create_split(VirtualAddress(0x4000_1000), 0x1000, None);
 
-        assert_eq!(0x4000_1000, new_region.va);
+        assert_eq!(VirtualAddress(0x4000_1000), new_region.va);
         assert_eq!(0x1000, new_region.length);
         assert!(matches!(new_region.physical_region, PhysicalRegion::Unused));
 
-        assert_eq!(0x4000_0000, splitted_regions[0].va);
+        assert_eq!(VirtualAddress(0x4000_0000), splitted_regions[0].va);
         assert_eq!(0x1000, splitted_regions[0].length);
-        assert_eq!(0x8000_0000, splitted_regions[0].get_pa());
+        assert_eq!(PhysicalAddress(0x8000_0000), splitted_regions[0].get_pa());
         assert!(matches!(
             splitted_regions[0].physical_region,
             PhysicalRegion::PhysicalAddress(_)
         ));
 
-        assert_eq!(0x4000_1000, splitted_regions[1].va);
+        assert_eq!(VirtualAddress(0x4000_1000), splitted_regions[1].va);
         assert_eq!(0x1000, splitted_regions[1].length);
         assert!(matches!(
             splitted_regions[1].physical_region,
             PhysicalRegion::Unused
         ));
 
-        assert_eq!(0x4000_2000, splitted_regions[2].va);
+        assert_eq!(VirtualAddress(0x4000_2000), splitted_regions[2].va);
         assert_eq!(0x2000, splitted_regions[2].length);
-        assert_eq!(0x8000_2000, splitted_regions[2].get_pa());
+        assert_eq!(PhysicalAddress(0x8000_2000), splitted_regions[2].get_pa());
         assert!(matches!(
             splitted_regions[2].physical_region,
             PhysicalRegion::PhysicalAddress(_)
         ));
 
         // New region at the end
-        let (new_region, splitted_regions) = region_unused.create_split(0x4000_3000, 0x1000, None);
+        let (new_region, splitted_regions) =
+            region_unused.create_split(VirtualAddress(0x4000_3000), 0x1000, None);
 
-        assert_eq!(0x4000_3000, new_region.va);
+        assert_eq!(VirtualAddress(0x4000_3000), new_region.va);
         assert_eq!(0x1000, new_region.length);
         assert!(matches!(new_region.physical_region, PhysicalRegion::Unused));
 
-        assert_eq!(0x4000_0000, splitted_regions[0].va);
+        assert_eq!(VirtualAddress(0x4000_0000), splitted_regions[0].va);
         assert_eq!(0x3000, splitted_regions[0].length);
-        assert_eq!(0x8000_0000, splitted_regions[0].get_pa());
+        assert_eq!(PhysicalAddress(0x8000_0000), splitted_regions[0].get_pa());
         assert!(matches!(
             splitted_regions[0].physical_region,
             PhysicalRegion::PhysicalAddress(_)
         ));
 
-        assert_eq!(0x4000_3000, splitted_regions[1].va);
+        assert_eq!(VirtualAddress(0x4000_3000), splitted_regions[1].va);
         assert_eq!(0x1000, splitted_regions[1].length);
 
         assert!(matches!(
@@ -608,26 +662,29 @@
     #[test]
     #[should_panic]
     fn test_virtual_region_does_not_contain() {
-        let region = VirtualRegion::new(0x4000_0000, 0x1000);
+        let region = VirtualRegion::new(VirtualAddress(0x4000_0000), 0x1000);
         region.create_split(
-            0x8000_0000,
+            VirtualAddress(0x8000_0000),
             0x1000,
-            Some(PhysicalRegion::PhysicalAddress(0xc000_0000)),
+            Some(PhysicalRegion::PhysicalAddress(PhysicalAddress(
+                0xc000_0000,
+            ))),
         );
     }
 
     #[test]
     #[should_panic]
     fn test_virtual_region_create_split_same_used() {
-        let region = VirtualRegion::new(0x4000_0000, 0x1000);
-        region.create_split(0x4000_0000, 0x1000, Some(PhysicalRegion::Unused));
+        let region = VirtualRegion::new(VirtualAddress(0x4000_0000), 0x1000);
+        region.create_split(
+            VirtualAddress(0x4000_0000),
+            0x1000,
+            Some(PhysicalRegion::Unused),
+        );
     }
 
     #[test]
     fn test_virtual_region_drop() {
-        const PA: usize = 0x0123_4567_89ab_cdef;
-        const LENGTH: usize = 0x8000_0000_0000;
-
         static PAGE_POOL_AREA: PagePoolArea<8192> = PagePoolArea::new();
         let page_pool = PagePool::new(&PAGE_POOL_AREA);
         let page = page_pool.allocate_pages(4096).unwrap();
@@ -635,7 +692,8 @@
         let physical_region = PhysicalRegion::Allocated(page_pool, page);
 
         // Testing physical region drop through virtualregion
-        let virtual_region = VirtualRegion::new_from_fields(0x4000_0000, 1000, physical_region);
+        let virtual_region =
+            VirtualRegion::new_from_fields(VirtualAddress(0x4000_0000), 1000, physical_region);
         drop(virtual_region);
     }
 }