Handle 16k and 64k translation granules

Enable Xlat to handle 16k and 64k translation granules along different
VA bit counts.

Signed-off-by: Imre Kis <imre.kis@arm.com>
Change-Id: Iab4fe066e813d5b75a5a6d45ba8498867cc5c541
diff --git a/src/page_pool.rs b/src/page_pool.rs
index 2d966d9..e0b6dc4 100644
--- a/src/page_pool.rs
+++ b/src/page_pool.rs
@@ -58,7 +58,7 @@
     /// Zero init pages
     pub fn zero_init(&mut self) {
         unsafe {
-            self.get_as_slice::<u8>().fill(0);
+            self.get_as_mut_slice::<u8>().fill(0);
         }
     }
 
@@ -67,11 +67,24 @@
         PhysicalAddress(self.pa)
     }
 
+    /// Get as slice
+    ///
+    /// **Unsafe**: The returned slice is created from its address and length which is stored in the
+    /// object. The caller has to ensure that no other references are being used of the pages.
+    pub unsafe fn get_as_slice<T>(&self) -> &[T] {
+        assert!((core::mem::align_of::<T>() - 1) & self.pa == 0);
+
+        core::slice::from_raw_parts(
+            KernelSpace::pa_to_kernel(self.pa as u64) as *const T,
+            self.length / core::mem::size_of::<T>(),
+        )
+    }
+
     /// Get as mutable slice
     ///
     /// **Unsafe**: The returned slice is created from its address and length which is stored in the
     /// object. The caller has to ensure that no other references are being used of the pages.
-    pub unsafe fn get_as_slice<T>(&mut self) -> &mut [T] {
+    pub unsafe fn get_as_mut_slice<T>(&mut self) -> &mut [T] {
         assert!((core::mem::align_of::<T>() - 1) & self.pa == 0);
 
         core::slice::from_raw_parts_mut(
@@ -200,10 +213,18 @@
     }
 
     /// Allocate pages for given length
-    pub fn allocate_pages(&self, length: usize) -> Result<Pages, PagePoolError> {
-        self.pages
-            .lock()
-            .allocate(Self::round_up_to_page_size(length), (), None)
+    pub fn allocate_pages(
+        &self,
+        length: usize,
+        alignment: Option<usize>,
+    ) -> Result<Pages, PagePoolError> {
+        let aligned_length = if let Some(alignment) = alignment {
+            length.next_multiple_of(alignment)
+        } else {
+            length
+        };
+
+        self.pages.lock().allocate(aligned_length, (), alignment)
     }
 
     /// Release pages
@@ -239,7 +260,7 @@
         pages.zero_init();
         assert_eq!([0, 0, 0, 0, 0, 0, 0, 0], area[0..8]);
 
-        let s = unsafe { pages.get_as_slice() };
+        let s = unsafe { pages.get_as_mut_slice() };
         for (i, e) in s.iter_mut().enumerate().take(8) {
             *e = i as u8;
         }