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/descriptor.rs b/src/descriptor.rs
index 24ed8df..0731856 100644
--- a/src/descriptor.rs
+++ b/src/descriptor.rs
@@ -12,6 +12,7 @@
 use crate::MemoryAttributesIndex;
 
 use super::address::PhysicalAddress;
+use super::TranslationGranule;
 
 /// Memory shareability
 #[derive(PrimitiveEnum_u8, Clone, Copy, Debug, PartialEq, Eq, Default)]
@@ -112,17 +113,14 @@
 impl Descriptor {
     const ATTR_MASK: u64 = 0xfff8_0000_0000_0ffc;
     const DESCRIPTOR_TYPE_MASK: u64 = 0b11;
-    pub const GRANULE_SIZES: [usize; 4] = [0, 0x4000_0000, 0x0020_0000, 0x0000_1000];
     const INVALID_DESCRIPTOR_VALUE: u64 = 0x0;
     const NEXT_ATTR_MASK: u64 = 0xf800_0000_0000_0000;
-    const OA_MASK: u64 = 0x0000_ffff_ffff_f000;
     const TABLE_BIT: u64 = 0b10;
-    const TABLE_ENTRY_COUNT: usize = 512;
     const TA_MASK: u64 = 0x0000_ffff_ffff_f000;
     const VALID_BIT: u64 = 0b01;
 
     /// Query descriptor type
-    pub fn get_descriptor_type(&self, level: usize) -> DescriptorType {
+    pub fn get_descriptor_type(&self, level: isize) -> DescriptorType {
         assert!(level <= 3);
 
         let desc_type_bits = self.get() & Self::DESCRIPTOR_TYPE_MASK;
@@ -150,9 +148,10 @@
     // Block descriptor functions
 
     /// Set block descriptor
-    pub fn set_block_descriptor(
+    pub fn set_block_descriptor<const VA_BITS: usize>(
         &mut self,
-        level: usize,
+        granule: TranslationGranule<VA_BITS>,
+        level: isize,
         output_address: PhysicalAddress,
         attributes: Attributes,
     ) {
@@ -160,24 +159,107 @@
 
         assert!(level <= 3);
         assert!(self.get_descriptor_type(level) != DescriptorType::Table);
-        assert_eq!(0, output_address.0 & !Self::get_oa_mask(level));
         assert_eq!(0, attr & !Self::ATTR_MASK);
 
+        let oa_granule_mask = !(granule.block_size_at_level(level) - 1);
+
+        // Figure D8-14 VMSAv8-64 Block descriptor formats
+        let oa_bits = match granule {
+            // 4KB and 16KB granules, 52-bit OA
+            #[cfg(feature = "feat_lpa2")]
+            TranslationGranule::Granule4k | TranslationGranule::Granule16k => {
+                let oa_mask = oa_granule_mask & 0x000f_ffff_ffff_f000;
+                assert_eq!(0, output_address.0 & !oa_mask);
+
+                let address = output_address.0 & oa_mask;
+
+                // OA[49:n] remains in place, OA[51:50] is mapped to [9:8]
+                let lsbs = address & 0x0003_ffff_ffff_f000;
+                let msbs = ((address >> 50) & 0x3) << 8;
+
+                lsbs | msbs
+            }
+
+            // 64KB granule, 52-bit OA
+            #[cfg(feature = "feat_lpa")]
+            TranslationGranule::Granule64k => {
+                let oa_mask = oa_granule_mask & 0x000f_ffff_ffff_0000;
+                assert_eq!(0, output_address.0 & !oa_mask);
+
+                let address = output_address.0 & oa_mask;
+
+                // OA[47:n] remains in place, OA[51:48] is mapped to [15:12]
+                let lsbs = address & 0x0000_ffff_ffff_0000;
+                let msbs = ((address >> 48) & 0xf) << 12;
+
+                lsbs | msbs
+            }
+
+            // 4KB, 16KB, and 64KB granules, 48-bit OA
+            #[cfg(not(all(feature = "feat_lpa", feature = "feat_lpa2")))]
+            _ => {
+                let oa_mask = oa_granule_mask & 0x0000_ffff_ffff_f000;
+                assert_eq!(0, output_address.0 & !oa_mask);
+                output_address.0 & oa_mask
+            }
+        };
+
         let table_bit = if level < 3 { 0 } else { Self::TABLE_BIT };
 
-        self.set(Self::VALID_BIT | table_bit | output_address.0 as u64 | attr);
+        self.set(Self::VALID_BIT | table_bit | oa_bits as u64 | attr);
     }
 
     /// Get output address from the block descriptor
-    pub fn get_block_output_address(&self, level: usize) -> PhysicalAddress {
+    pub fn get_block_output_address<const VA_BITS: usize>(
+        &self,
+        granule: TranslationGranule<VA_BITS>,
+        level: isize,
+    ) -> PhysicalAddress {
         assert!(level <= 3);
         assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));
 
-        PhysicalAddress((self.get() & Self::OA_MASK) as usize)
+        let oa_granule_mask = !(granule.block_size_at_level(level) - 1);
+        let descriptor_value = self.get();
+
+        // Figure D8-14 VMSAv8-64 Block descriptor formats
+        let pa = match granule {
+            // 4KB and 16KB granules, 52-bit OA
+            #[cfg(feature = "feat_lpa2")]
+            TranslationGranule::Granule4k | TranslationGranule::Granule16k => {
+                let oa_mask = oa_granule_mask & 0x000f_ffff_ffff_f000;
+
+                // OA[49:n] remains in place, OA[51:50] is mapped from [9:8]
+                let lsbs = descriptor_value & 0x0003_ffff_ffff_f000;
+                let msbs = ((descriptor_value >> 8) & 0x3) << 50;
+
+                (lsbs | msbs) as usize & oa_mask
+            }
+
+            // 64KB granule, 52-bit OA
+            #[cfg(feature = "feat_lpa")]
+            TranslationGranule::Granule64k => {
+                let oa_mask = oa_granule_mask & 0x000f_ffff_ffff_0000;
+
+                // OA[47:n] remains in place, OA[51:48] is mapped from [15:12]
+                let lsbs = descriptor_value & 0x0003_ffff_ffff_0000;
+                let msbs = ((descriptor_value >> 12) & 0xf) << 48;
+
+                (lsbs | msbs) as usize & oa_mask
+            }
+
+            // 4KB, 16KB, and 64KB granules, 48-bit OA
+            #[cfg(not(all(feature = "feat_lpa", feature = "feat_lpa2")))]
+            _ => {
+                let oa_mask = oa_granule_mask & 0x0000_ffff_ffff_f000;
+                descriptor_value as usize & oa_mask
+            }
+        };
+
+        PhysicalAddress(pa)
     }
 
     /// Set the attributes of the block descriptor
-    pub fn set_block_attributes(&mut self, level: usize, attributes: Attributes) {
+    pub fn set_block_attributes(&mut self, level: isize, attributes: Attributes) {
         assert!(level <= 3);
         let attr: u64 = attributes.into();
         assert_eq!(0, attr & !Self::ATTR_MASK);
@@ -187,7 +269,7 @@
     }
 
     /// Get the attributes of the block descriptor
-    pub fn get_block_attributes(&self, level: usize) -> Attributes {
+    pub fn get_block_attributes(&self, level: isize) -> Attributes {
         assert!(level <= 3);
         assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));
 
@@ -195,7 +277,7 @@
     }
 
     /// Set block descriptor to invalid
-    pub fn set_block_descriptor_to_invalid(&mut self, level: usize) {
+    pub fn set_block_descriptor_to_invalid(&mut self, level: isize) {
         assert!(level <= 3);
         assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));
 
@@ -203,7 +285,7 @@
     }
 
     /// Set block or invalid descriptor to invalid
-    pub fn set_block_or_invalid_descriptor_to_invalid(&mut self, level: usize) {
+    pub fn set_block_or_invalid_descriptor_to_invalid(&mut self, level: isize) {
         assert!(level <= 3);
         assert!(DescriptorType::Table != self.get_descriptor_type(level));
 
@@ -216,12 +298,11 @@
     /// the descriptor.
     pub unsafe fn set_table_descriptor(
         &mut self,
-        level: usize,
+        level: isize,
         next_level_table: &mut [Descriptor],
         next_level_attributes: Option<NextLevelAttributes>,
     ) {
         assert!(level <= 2);
-        assert_eq!(Self::TABLE_ENTRY_COUNT, next_level_table.len());
         assert!(self.get_descriptor_type(level) != DescriptorType::Table);
 
         let table_addr = KernelSpace::kernel_to_pa(next_level_table.as_ptr() as u64);
@@ -242,30 +323,38 @@
     ///
     /// Safety: The returned next level table is based on the address read from the descriptor.
     /// The caller has to ensure that no other references are being used of the table.
-    pub unsafe fn get_next_level_table(&self, level: usize) -> &[Descriptor] {
+    pub unsafe fn get_next_level_table<const VA_BITS: usize>(
+        &self,
+        granule: TranslationGranule<VA_BITS>,
+        level: isize,
+    ) -> &[Descriptor] {
         assert!(level <= 2);
         assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));
 
         let table_address =
             KernelSpace::pa_to_kernel(self.get() & Self::TA_MASK) as *const Descriptor;
-        core::slice::from_raw_parts(table_address, Self::TABLE_ENTRY_COUNT)
+        core::slice::from_raw_parts(table_address, granule.entry_count_at_level(level + 1))
     }
 
     /// Get mutable next level table
     ///
     /// Safety: The returned next level table is based on the address read from the descriptor.
     /// The caller has to ensure that no other references are being used of the table.
-    pub unsafe fn get_next_level_table_mut(&mut self, level: usize) -> &mut [Descriptor] {
+    pub unsafe fn get_next_level_table_mut<const VA_BITS: usize>(
+        &mut self,
+        granule: TranslationGranule<VA_BITS>,
+        level: isize,
+    ) -> &mut [Descriptor] {
         assert!(level <= 2);
         assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));
 
         let table_address =
             KernelSpace::pa_to_kernel(self.get() & Self::TA_MASK) as *mut Descriptor;
-        core::slice::from_raw_parts_mut(table_address, Self::TABLE_ENTRY_COUNT)
+        core::slice::from_raw_parts_mut(table_address, granule.entry_count_at_level(level + 1))
     }
 
     /// Get next level attributes
-    pub fn get_next_level_attributes(&self, level: usize) -> NextLevelAttributes {
+    pub fn get_next_level_attributes(&self, level: isize) -> NextLevelAttributes {
         assert!(level <= 2);
         assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));
 
@@ -276,14 +365,18 @@
     ///
     /// **Unsafe:** The returned descriptor reference must be released by the caller, i.e. release
     /// to `PagePool`
-    pub unsafe fn set_table_descriptor_to_invalid(&mut self, level: usize) -> &mut [Descriptor] {
+    pub unsafe fn set_table_descriptor_to_invalid<const VA_BITS: usize>(
+        &mut self,
+        granule: TranslationGranule<VA_BITS>,
+        level: isize,
+    ) -> &mut [Descriptor] {
         assert!(level <= 2);
         assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));
 
         let table_address =
             KernelSpace::pa_to_kernel(self.get() & Self::TA_MASK) as *mut Descriptor;
         self.set(Self::INVALID_DESCRIPTOR_VALUE);
-        core::slice::from_raw_parts_mut(table_address, Self::TABLE_ENTRY_COUNT)
+        core::slice::from_raw_parts_mut(table_address, granule.entry_count_at_level(level + 1))
     }
 
     /// Get raw descriptor value
@@ -309,11 +402,6 @@
     {
         self.set(f(self.get()))
     }
-
-    /// Get output address mask
-    fn get_oa_mask(level: usize) -> usize {
-        Self::OA_MASK as usize & !(Self::GRANULE_SIZES[level] - 1)
-    }
 }
 
 #[cfg(test)]
@@ -477,7 +565,12 @@
             cell: UnsafeCell::new(1),
         };
 
-        descriptor.set_block_descriptor(1, PhysicalAddress(0), Attributes::default());
+        descriptor.set_block_descriptor::<48>(
+            TranslationGranule::Granule4k,
+            1,
+            PhysicalAddress(0),
+            Attributes::default(),
+        );
         assert_eq!(0x1, descriptor.get());
     }
 
@@ -488,16 +581,23 @@
             cell: UnsafeCell::new(0),
         };
 
-        descriptor.set_block_descriptor(1, PhysicalAddress(1 << 63), Attributes::default());
+        descriptor.set_block_descriptor::<48>(
+            TranslationGranule::Granule4k,
+            1,
+            PhysicalAddress(1 << 63),
+            Attributes::default(),
+        );
     }
 
     #[test]
     fn test_descriptor_block() {
+        let granule: TranslationGranule<48> = TranslationGranule::Granule4k;
         let mut descriptor = Descriptor {
             cell: UnsafeCell::new(0),
         };
 
         descriptor.set_block_descriptor(
+            granule,
             1,
             PhysicalAddress(0x0000000f_c0000000),
             Attributes {
@@ -512,6 +612,7 @@
         };
 
         descriptor.set_block_descriptor(
+            granule,
             3,
             PhysicalAddress(0x0000000f_fffff000),
             Attributes {
@@ -523,7 +624,7 @@
 
         assert_eq!(
             PhysicalAddress(0x0000000f_fffff000),
-            descriptor.get_block_output_address(3)
+            descriptor.get_block_output_address(granule, 3)
         );
         assert_eq!(
             Attributes {
@@ -687,7 +788,9 @@
             cell: UnsafeCell::new(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003),
         };
         assert_eq!(KernelSpace::pa_to_kernel(NEXT_LEVEL_ADDR), unsafe {
-            descriptor.get_next_level_table(0).as_ptr() as u64
+            descriptor
+                .get_next_level_table::<36>(TranslationGranule::Granule4k, 0)
+                .as_ptr() as u64
         });
     }
 
@@ -698,7 +801,9 @@
             cell: UnsafeCell::new(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003),
         };
         assert_eq!(KernelSpace::pa_to_kernel(NEXT_LEVEL_ADDR), unsafe {
-            descriptor.get_next_level_table_mut(0).as_ptr() as *mut Descriptor as u64
+            descriptor
+                .get_next_level_table_mut::<36>(TranslationGranule::Granule4k, 0)
+                .as_ptr() as *mut Descriptor as u64
         });
     }
 
@@ -724,7 +829,9 @@
             cell: UnsafeCell::new(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003),
         };
         assert_eq!(KernelSpace::pa_to_kernel(NEXT_LEVEL_ADDR), unsafe {
-            descriptor.set_table_descriptor_to_invalid(0).as_ptr() as *mut Descriptor as u64
+            descriptor
+                .set_table_descriptor_to_invalid::<36>(TranslationGranule::Granule4k, 0)
+                .as_ptr() as *mut Descriptor as u64
         });
         assert_eq!(0, descriptor.get());
     }