Add TranslationGranule module

Collect all translation granule related calculations into the
TranslationGranule module.

Signed-off-by: Imre Kis <imre.kis@arm.com>
Change-Id: I96286fe3f62d87541615189ce879b944507d493b
diff --git a/src/granule.rs b/src/granule.rs
new file mode 100644
index 0000000..0263c90
--- /dev/null
+++ b/src/granule.rs
@@ -0,0 +1,281 @@
+// SPDX-FileCopyrightText: Copyright 2024 Arm Limited and/or its affiliates <open-source-office@arm.com>
+// SPDX-License-Identifier: MIT OR Apache-2.0
+
+#[derive(Clone, Copy)]
+pub enum TranslationGranule<const VA_BITS: usize> {
+    Granule4k = 4 * 1024,
+    Granule16k = 16 * 1024,
+    Granule64k = 64 * 1024,
+}
+
+impl<const VA_BITS: usize> TranslationGranule<VA_BITS> {
+    pub const fn initial_lookup_level(&self) -> isize {
+        match self {
+            TranslationGranule::Granule4k => match VA_BITS {
+                #[cfg(feature = "feat_ttst")]
+                16..=21 => 3,
+                #[cfg(feature = "feat_ttst")]
+                22..=24 => 2,
+                25..=30 => 2,
+                31..=39 => 1,
+                40..=48 => 0,
+                #[cfg(feature = "feat_lpa2")]
+                49..=52 => -1,
+                _ => panic!("Invalid VA_BITS"),
+            },
+            TranslationGranule::Granule16k => match VA_BITS {
+                #[cfg(feature = "feat_ttst")]
+                16..=24 => 3,
+                25 => 3,
+                26..=36 => 2,
+                37..=47 => 1,
+                48 => 0,
+                #[cfg(feature = "feat_lpa2")]
+                49..=52 => 0,
+                _ => panic!("Invalid VA_BITS"),
+            },
+            TranslationGranule::Granule64k => match VA_BITS {
+                #[cfg(feature = "feat_ttst")]
+                17..=24 => 3,
+                25..=29 => 3,
+                30..=42 => 2,
+                43..=48 => 1,
+                #[cfg(feature = "feat_lva")]
+                49..=52 => 1,
+                _ => panic!("Invalid VA_BITS"),
+            },
+        }
+    }
+
+    pub const fn entry_count_at_level(&self, level: isize) -> usize {
+        let total_bits_at_level = self.total_bits_at_level(level);
+        let bits_per_level = match self {
+            TranslationGranule::Granule4k => 9,
+            TranslationGranule::Granule16k => 11,
+            TranslationGranule::Granule64k => 13,
+        };
+
+        let bits = if total_bits_at_level + bits_per_level < VA_BITS {
+            total_bits_at_level + bits_per_level
+        } else {
+            VA_BITS
+        } - total_bits_at_level;
+
+        1 << bits
+    }
+
+    pub const fn table_size<D>(&self, level: isize) -> usize {
+        self.entry_count_at_level(level) * core::mem::size_of::<D>()
+    }
+
+    pub const fn table_alignment<D>(&self, level: isize) -> usize {
+        // D8.2.5.2
+        // For the VMSAv8-64 translation system, if the translation table has fewer than eight
+        // entries and an OA size greater than 48 bits is used, then the table is aligned to 64
+        // bytes. Otherwise, the translation table is aligned to the size of that translation table.
+
+        let alignment = self.table_size::<D>(level);
+
+        if VA_BITS > 48 && alignment < 64 {
+            64
+        } else {
+            alignment
+        }
+    }
+
+    pub const fn block_size_at_level(&self, level: isize) -> usize {
+        1 << self.total_bits_at_level(level)
+    }
+
+    pub const fn total_bits_at_level(&self, level: isize) -> usize {
+        assert!(self.initial_lookup_level() <= level && level <= 3);
+
+        match self {
+            TranslationGranule::Granule4k => match level {
+                -1 => 48,
+                0 => 39,
+                1 => 30,
+                2 => 21,
+                3 => 12,
+                _ => panic!("Invalid translation level"),
+            },
+            TranslationGranule::Granule16k => match level {
+                0 => 47,
+                1 => 36,
+                2 => 25,
+                3 => 14,
+                _ => panic!("Invalid translation level"),
+            },
+            TranslationGranule::Granule64k => match level {
+                1 => 42,
+                2 => 29,
+                3 => 16,
+                _ => panic!("Invalid translation level"),
+            },
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    macro_rules! validate_initial_lookup_level {
+        ( $granule:expr, $tnsz_min:literal, $tnsz_max:literal, $level:literal ) => {
+            assert_eq!(
+                $level,
+                TranslationGranule::<{ 64 - $tnsz_min }>::initial_lookup_level(&$granule)
+            );
+            assert_eq!(
+                $level,
+                TranslationGranule::<{ 64 - $tnsz_max }>::initial_lookup_level(&$granule)
+            );
+        };
+    }
+
+    #[test]
+    fn test_initial_lookup_level() {
+        // Table D8-16 4KB granule, determining stage 1 initial lookup level
+        #[cfg(feature = "feat_lpa2")]
+        validate_initial_lookup_level!(TranslationGranule::Granule4k, 12, 15, -1);
+        validate_initial_lookup_level!(TranslationGranule::Granule4k, 16, 24, 0);
+        validate_initial_lookup_level!(TranslationGranule::Granule4k, 25, 33, 1);
+        validate_initial_lookup_level!(TranslationGranule::Granule4k, 34, 39, 2);
+        #[cfg(feature = "feat_ttst")]
+        validate_initial_lookup_level!(TranslationGranule::Granule4k, 40, 42, 2);
+        #[cfg(feature = "feat_ttst")]
+        validate_initial_lookup_level!(TranslationGranule::Granule4k, 43, 48, 3);
+
+        //// Table D8-26 16KB granule, determining stage 1 initial lookup level
+        #[cfg(feature = "feat_lpa2")]
+        validate_initial_lookup_level!(TranslationGranule::Granule16k, 12, 15, 0);
+        validate_initial_lookup_level!(TranslationGranule::Granule16k, 16, 16, 0);
+        validate_initial_lookup_level!(TranslationGranule::Granule16k, 17, 27, 1);
+        validate_initial_lookup_level!(TranslationGranule::Granule16k, 28, 38, 2);
+        validate_initial_lookup_level!(TranslationGranule::Granule16k, 39, 39, 3);
+        #[cfg(feature = "feat_ttst")]
+        validate_initial_lookup_level!(TranslationGranule::Granule16k, 40, 48, 3);
+        //
+        //// D8.2.10.1 VMSAv8-64 Stage 1 address translation using the 64KB translation granule
+        #[cfg(feature = "feat_lva")]
+        validate_initial_lookup_level!(TranslationGranule::Granule64k, 12, 15, 1);
+        validate_initial_lookup_level!(TranslationGranule::Granule64k, 16, 21, 1);
+        validate_initial_lookup_level!(TranslationGranule::Granule64k, 22, 34, 2);
+        validate_initial_lookup_level!(TranslationGranule::Granule64k, 35, 39, 3);
+        #[cfg(feature = "feat_ttst")]
+        validate_initial_lookup_level!(TranslationGranule::Granule64k, 40, 47, 3);
+    }
+
+    #[test]
+    fn test_entry_count_at_level() {
+        // Table D8-14 4KB granule translation table properties at each lookup level
+        #[cfg(feature = "feat_lpa2")]
+        {
+            let granule = TranslationGranule::<52>::Granule4k;
+            assert_eq!(16, granule.entry_count_at_level(-1));
+            assert_eq!(512, granule.entry_count_at_level(0));
+            assert_eq!(512, granule.entry_count_at_level(1));
+            assert_eq!(512, granule.entry_count_at_level(2));
+            assert_eq!(512, granule.entry_count_at_level(3));
+        }
+        #[cfg(not(feature = "feat_lpa2"))]
+        {
+            let granule = TranslationGranule::<48>::Granule4k;
+            assert_eq!(512, granule.entry_count_at_level(0));
+            assert_eq!(512, granule.entry_count_at_level(1));
+            assert_eq!(512, granule.entry_count_at_level(2));
+            assert_eq!(512, granule.entry_count_at_level(3));
+        }
+
+        let granule = TranslationGranule::<42>::Granule4k;
+        assert_eq!(8, granule.entry_count_at_level(0));
+        assert_eq!(512, granule.entry_count_at_level(1));
+        assert_eq!(512, granule.entry_count_at_level(2));
+        assert_eq!(512, granule.entry_count_at_level(3));
+
+        let granule = TranslationGranule::<32>::Granule4k;
+        assert_eq!(4, granule.entry_count_at_level(1));
+        assert_eq!(512, granule.entry_count_at_level(2));
+        assert_eq!(512, granule.entry_count_at_level(3));
+
+        let granule = TranslationGranule::<26>::Granule4k;
+        assert_eq!(32, granule.entry_count_at_level(2));
+        assert_eq!(512, granule.entry_count_at_level(3));
+
+        // Table D8-24 16KB granule translation table properties at each lookup level
+        #[cfg(feature = "feat_lpa2")]
+        {
+            let granule = TranslationGranule::<52>::Granule16k;
+            assert_eq!(32, granule.entry_count_at_level(0));
+            assert_eq!(2048, granule.entry_count_at_level(1));
+            assert_eq!(2048, granule.entry_count_at_level(2));
+            assert_eq!(2048, granule.entry_count_at_level(3));
+        }
+        #[cfg(not(feature = "feat_lpa2"))]
+        {
+            let granule = TranslationGranule::<48>::Granule16k;
+            assert_eq!(2, granule.entry_count_at_level(0));
+            assert_eq!(2048, granule.entry_count_at_level(1));
+            assert_eq!(2048, granule.entry_count_at_level(2));
+            assert_eq!(2048, granule.entry_count_at_level(3));
+        }
+
+        let granule = TranslationGranule::<40>::Granule16k;
+        assert_eq!(16, granule.entry_count_at_level(1));
+        assert_eq!(2048, granule.entry_count_at_level(2));
+        assert_eq!(2048, granule.entry_count_at_level(3));
+
+        let granule = TranslationGranule::<30>::Granule16k;
+        assert_eq!(32, granule.entry_count_at_level(2));
+        assert_eq!(2048, granule.entry_count_at_level(3));
+
+        // Table D8-33 64KB granule translation table properties at each lookup level
+        #[cfg(feature = "feat_lva")]
+        {
+            let granule = TranslationGranule::<52>::Granule64k;
+            assert_eq!(1024, granule.entry_count_at_level(1));
+        }
+
+        let granule = TranslationGranule::<48>::Granule64k;
+        #[cfg(not(feature = "feat_lva"))]
+        {
+            assert_eq!(64, granule.entry_count_at_level(1));
+        }
+
+        assert_eq!(8192, granule.entry_count_at_level(2));
+        assert_eq!(8192, granule.entry_count_at_level(3));
+
+        let granule = TranslationGranule::<40>::Granule64k;
+        assert_eq!(2048, granule.entry_count_at_level(2));
+        assert_eq!(8192, granule.entry_count_at_level(3));
+    }
+
+    #[test]
+    fn test_granule_size_at_level() {
+        // Table D8-15 4KB granule block descriptor properties at each lookup level
+        #[cfg(feature = "feat_lpa2")]
+        {
+            let granule = TranslationGranule::<52>::Granule4k;
+            assert_eq!(256 << 40, granule.block_size_at_level(-1));
+        }
+
+        let granule = TranslationGranule::<48>::Granule4k;
+        assert_eq!(512 << 30, granule.block_size_at_level(0));
+        assert_eq!(1 << 30, granule.block_size_at_level(1));
+        assert_eq!(2 << 20, granule.block_size_at_level(2));
+        assert_eq!(4 << 10, granule.block_size_at_level(3));
+
+        // Table D8-25 16KB granule block descriptor properties at each lookup level
+        let granule = TranslationGranule::<48>::Granule16k;
+        assert_eq!(128 << 40, granule.block_size_at_level(0));
+        assert_eq!(64 << 30, granule.block_size_at_level(1));
+        assert_eq!(32 << 20, granule.block_size_at_level(2));
+        assert_eq!(16 << 10, granule.block_size_at_level(3));
+
+        // Table D8-34 64KB granule block descriptor properties at each lookup level
+        let granule = TranslationGranule::<48>::Granule64k;
+        assert_eq!(4 << 40, granule.block_size_at_level(1));
+        assert_eq!(512 << 20, granule.block_size_at_level(2));
+        assert_eq!(64 << 10, granule.block_size_at_level(3));
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 74d4f88..f5b1962 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -30,6 +30,7 @@
 
 pub mod address;
 mod descriptor;
+mod granule;
 pub mod kernel_space;
 pub mod page_pool;
 mod region;
@@ -161,6 +162,8 @@
     EL3,                      // EL3, TTBR0_EL3
 }
 
+pub type TranslationGranule<const VA_BITS: usize> = granule::TranslationGranule<VA_BITS>;
+
 pub struct Xlat {
     base_table: Box<BaseTable>,
     page_pool: PagePool,