Implement aligned allocation in region pool

Allow regions to be allocated from region pools with an optional
alignment parameter. Region trait now requires the implementation of the
try_alloc_aligned function.

Signed-off-by: Imre Kis <imre.kis@arm.com>
Change-Id: I35d5546e1612020fa5eef22c0a2b9c75cba36834
diff --git a/src/address.rs b/src/address.rs
index 12bbd8e..a8f2c39 100644
--- a/src/address.rs
+++ b/src/address.rs
@@ -77,6 +77,10 @@
     pub const fn diff(self, rhs: Self) -> Option<usize> {
         self.0.checked_sub(rhs.0)
     }
+
+    pub const fn align_up(self, alignment: usize) -> Self {
+        Self(self.0.next_multiple_of(alignment))
+    }
 }
 
 impl From<VirtualAddress> for usize {
diff --git a/src/lib.rs b/src/lib.rs
index 62cac5c..74d4f88 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -243,7 +243,7 @@
             self.regions
                 .acquire(required_va, pages_length, physical_region)
         } else {
-            self.regions.allocate(pages_length, physical_region)
+            self.regions.allocate(pages_length, physical_region, None)
         }
         .map_err(XlatError::RegionPoolError)?;
 
@@ -275,7 +275,7 @@
             self.regions
                 .acquire(required_va, pages_length, physical_region)
         } else {
-            self.regions.allocate(pages_length, physical_region)
+            self.regions.allocate(pages_length, physical_region, None)
         }
         .map_err(XlatError::RegionPoolError)?;
 
@@ -301,7 +301,7 @@
         let region = if let Some(required_va) = va {
             self.regions.acquire(required_va, length, resource)
         } else {
-            self.regions.allocate(length, resource)
+            self.regions.allocate(length, resource, None)
         }
         .map_err(XlatError::RegionPoolError)?;
 
diff --git a/src/page_pool.rs b/src/page_pool.rs
index 8636af7..2d966d9 100644
--- a/src/page_pool.rs
+++ b/src/page_pool.rs
@@ -96,6 +96,7 @@
     type Resource = ();
     type Base = usize;
     type Length = usize;
+    type Alignment = usize;
 
     fn base(&self) -> usize {
         self.pa
@@ -119,6 +120,22 @@
         }
     }
 
+    fn try_alloc_aligned(
+        &self,
+        length: Self::Length,
+        alignment: Self::Alignment,
+    ) -> Option<Self::Base> {
+        let aligned_base = self.pa.next_multiple_of(alignment);
+        let base_offset = aligned_base.checked_sub(self.pa)?;
+
+        let required_length = base_offset.checked_add(length)?;
+        if required_length <= self.length {
+            Some(aligned_base)
+        } else {
+            None
+        }
+    }
+
     fn try_append(&mut self, other: &Self) -> bool {
         if let (Some(self_end), Some(new_length)) = (
             self.pa.checked_add(self.length),
@@ -186,7 +203,7 @@
     pub fn allocate_pages(&self, length: usize) -> Result<Pages, PagePoolError> {
         self.pages
             .lock()
-            .allocate(Self::round_up_to_page_size(length), ())
+            .allocate(Self::round_up_to_page_size(length), (), None)
     }
 
     /// Release pages
@@ -245,4 +262,13 @@
 
         // Overflow tests
     }
+
+    #[test]
+    fn test_pages_try_alloc() {
+        let pages = Pages::new(0x4000_1000, 0x10000, false);
+
+        assert_eq!(Some(0x4000_1000), pages.try_alloc(0x1000, None));
+        assert_eq!(Some(0x4000_2000), pages.try_alloc(0x1000, Some(0x2000)));
+        assert_eq!(None, pages.try_alloc(0x1000, Some(0x10_0000)));
+    }
 }
diff --git a/src/region.rs b/src/region.rs
index 83c54ab..09958e8 100644
--- a/src/region.rs
+++ b/src/region.rs
@@ -84,6 +84,7 @@
     type Resource = PhysicalRegion;
     type Base = VirtualAddress;
     type Length = usize;
+    type Alignment = usize;
 
     fn base(&self) -> Self::Base {
         self.va
@@ -107,6 +108,22 @@
         }
     }
 
+    fn try_alloc_aligned(
+        &self,
+        length: Self::Length,
+        alignment: Self::Alignment,
+    ) -> Option<Self::Base> {
+        let aligned_base = self.va.align_up(alignment);
+        let base_offset = aligned_base.diff(self.va)?;
+
+        let required_length = base_offset.checked_add(length)?;
+        if required_length <= self.length {
+            Some(aligned_base)
+        } else {
+            None
+        }
+    }
+
     fn try_append(&mut self, other: &Self) -> bool {
         if let (Some(self_end), Some(new_length)) = (
             self.va.add_offset(self.length),
@@ -345,6 +362,25 @@
     }
 
     #[test]
+    fn test_virtual_region_try_alloc() {
+        let region = VirtualRegion::new_with_pa(
+            PhysicalAddress(0x8000_1000),
+            VirtualAddress(0x4000_1000),
+            0x10000,
+        );
+
+        assert_eq!(
+            Some(VirtualAddress(0x4000_1000)),
+            region.try_alloc(0x1000, None)
+        );
+        assert_eq!(
+            Some(VirtualAddress(0x4000_2000)),
+            region.try_alloc(0x1000, Some(0x2000))
+        );
+        assert_eq!(None, region.try_alloc(0x1000, Some(0x10_0000)));
+    }
+
+    #[test]
     fn test_virtual_region_contains() {
         const VA: usize = 0x8000_0000_0000_0000;
         const LENGTH: usize = 0x8000_0000_0000;
diff --git a/src/region_pool.rs b/src/region_pool.rs
index 4c81160..116e523 100644
--- a/src/region_pool.rs
+++ b/src/region_pool.rs
@@ -16,6 +16,7 @@
     type Resource;
     type Base: Ord + Copy;
     type Length: Ord + Copy;
+    type Alignment: Copy;
 
     /// Get base address
     fn base(&self) -> Self::Base;
@@ -29,6 +30,34 @@
     /// Check if an area defined by its base address and length is inside the region
     fn contains(&self, base: Self::Base, length: Self::Length) -> bool;
 
+    /// Try to allocate a base address from the region. If there's no alignment specified, the
+    /// function will check if the requested length fits into the buffer and if yes, it will
+    /// return the base of the region.
+    /// If an alignment is specified, it will call try_alloc_aligned which has to be implemeted
+    /// by the trait users.
+    /// This function is part of the trait so the trait users can override it if they want to
+    /// achieve a better implementation specific behavior.
+    fn try_alloc(
+        &self,
+        length: Self::Length,
+        alignment: Option<Self::Alignment>,
+    ) -> Option<Self::Base> {
+        if let Some(alignment) = alignment {
+            self.try_alloc_aligned(length, alignment)
+        } else if length <= self.length() {
+            Some(self.base())
+        } else {
+            None
+        }
+    }
+
+    /// Try to allocate aligned address from the region.
+    fn try_alloc_aligned(
+        &self,
+        length: Self::Length,
+        alignment: Self::Alignment,
+    ) -> Option<Self::Base>;
+
     /// Try append the parameter region. Return true on success.
     fn try_append(&mut self, other: &Self) -> bool;
 
@@ -92,17 +121,21 @@
         &mut self,
         length: T::Length,
         resource: T::Resource,
+        alignment: Option<T::Alignment>,
     ) -> Result<T, RegionPoolError> {
         let region_to_allocate_from = self
             .regions
             .iter()
             .enumerate()
-            .filter(|(_index, region)| region.length() >= length && !region.used())
+            .filter(|(_index, region)| {
+                !region.used() && region.try_alloc(length, alignment).is_some()
+            })
             .min_by_key(|(_index_a, region)| region.length());
 
         if let Some((index, region)) = region_to_allocate_from {
-            let (new_region, split_regions) =
-                region.create_split(region.base(), length, Some(resource));
+            // The previous call to try_alloc ensures that the following unwrap will not fail.
+            let base = region.try_alloc(length, alignment).unwrap();
+            let (new_region, split_regions) = region.create_split(base, length, Some(resource));
             self.replace(index, split_regions);
             Ok(new_region)
         } else {
@@ -211,6 +244,7 @@
         type Resource = ();
         type Base = usize;
         type Length = usize;
+        type Alignment = usize;
 
         fn base(&self) -> usize {
             self.base
@@ -228,6 +262,22 @@
             self.base <= base && base + length <= self.base + self.length
         }
 
+        fn try_alloc_aligned(
+            &self,
+            length: Self::Length,
+            alignment: Self::Alignment,
+        ) -> Option<Self::Base> {
+            let aligned_base = self.base.next_multiple_of(alignment);
+            let base_offset = aligned_base.checked_sub(self.base)?;
+
+            let required_length = base_offset.checked_add(length)?;
+            if required_length <= self.length {
+                Some(aligned_base)
+            } else {
+                None
+            }
+        }
+
         fn try_append(&mut self, other: &Self) -> bool {
             if self.used == other.used && self.base + self.length == other.base {
                 self.length += other.length;
@@ -271,6 +321,15 @@
     }
 
     #[test]
+    fn test_region_try_alloc() {
+        let region = RegionExample::new(0x4000_1000, 0x1_0000, false);
+
+        assert_eq!(Some(0x4000_1000), region.try_alloc(0x1000, None));
+        assert_eq!(Some(0x4000_2000), region.try_alloc(0x1000, Some(0x2000)));
+        assert_eq!(None, region.try_alloc(0x1000, Some(0x10_0000)));
+    }
+
+    #[test]
     fn test_region_try_append() {
         // Normal append
         let mut region = RegionExample::new(0x4000_0000, 0x1000_0000, false);
@@ -342,15 +401,30 @@
         pool.add(RegionExample::new(0x4000_0000, 0x1000_0000, false))
             .unwrap();
 
-        let res = pool.allocate(0x1000, ());
+        let res = pool.allocate(0x1000, (), None);
         assert_eq!(Ok(RegionExample::new(0x4000_0000, 0x1000, true)), res);
-        let res = pool.allocate(0x1_0000, ());
+        let res = pool.allocate(0x1_0000, (), None);
         assert_eq!(Ok(RegionExample::new(0x4000_1000, 0x1_0000, true)), res);
-        let res = pool.allocate(0x2000_0000, ());
+        let res = pool.allocate(0x2000_0000, (), None);
         assert_eq!(Err(RegionPoolError::NoSpace), res);
     }
 
     #[test]
+    fn test_pool_allocate_aligned() {
+        let mut pool = RegionPool::new();
+        pool.add(RegionExample::new(0x4000_0000, 0x1000_0000, false))
+            .unwrap();
+
+        let res = pool.allocate(0x1000, (), Some(0x2000));
+        assert_eq!(Ok(RegionExample::new(0x4000_0000, 0x1000, true)), res);
+        let res = pool.allocate(0x2000, (), Some(0x1_0000));
+        assert_eq!(Ok(RegionExample::new(0x4001_0000, 0x2000, true)), res);
+        let res = pool.allocate(0x1000, (), None);
+        // This region is allocated from the empty space between the first two regions
+        assert_eq!(Ok(RegionExample::new(0x4000_1000, 0x1000, true)), res);
+    }
+
+    #[test]
     fn test_region_pool_acquire() {
         let mut pool = RegionPool::new();
         pool.add(RegionExample::new(0x4000_0000, 0x1000_0000, false))
@@ -358,7 +432,7 @@
 
         let res = pool.acquire(0x4000_1000, 0x1000, ());
         assert_eq!(Ok(RegionExample::new(0x4000_1000, 0x1000, true)), res);
-        let res = pool.allocate(0x1_0000, ());
+        let res = pool.allocate(0x1_0000, (), None);
         assert_eq!(Ok(RegionExample::new(0x4000_2000, 0x10000, true)), res);
         let res = pool.acquire(0x4000_1000, 0x1000, ());
         assert_eq!(Err(RegionPoolError::AlreadyUsed), res);
@@ -370,11 +444,11 @@
         pool.add(RegionExample::new(0x4000_0000, 0x1000_0000, false))
             .unwrap();
 
-        let res = pool.allocate(0x1000, ());
+        let res = pool.allocate(0x1000, (), None);
         assert_eq!(Ok(RegionExample::new(0x4000_0000, 0x1000, true)), res);
         assert_eq!(Ok(()), pool.release(res.unwrap()));
 
-        let res = pool.allocate(0x1_0000, ());
+        let res = pool.allocate(0x1_0000, (), None);
         assert_eq!(Ok(RegionExample::new(0x4000_0000, 0x10000, true)), res);
         assert_eq!(Ok(()), pool.release(res.unwrap()));