Implement TLB maintenance

Invalidate TLB entries after changing descriptor to invalid and
implement break-before-make sequence when changing translation table
entries.

Signed-off-by: Imre Kis <imre.kis@arm.com>
Change-Id: I8e384de306c185315a1d189c01ff275ac646d539
diff --git a/src/lib.rs b/src/lib.rs
index dd7de75..62cac5c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -560,6 +560,7 @@
             1,
             self.base_table.descriptors.as_mut_slice(),
             &self.page_pool,
+            &self.regime,
         );
     }
 
@@ -573,6 +574,7 @@
     /// * level: Translation table level
     /// * table: Translation table on the given level
     /// * page_pool: Page pool where the function can allocate pages for the translation tables
+    #[allow(clippy::too_many_arguments)]
     fn set_block_descriptor_recursively(
         attributes: Attributes,
         pa: PhysicalAddress,
@@ -581,12 +583,16 @@
         level: usize,
         table: &mut [Descriptor],
         page_pool: &PagePool,
+        regime: &TranslationRegime,
     ) {
         // Get descriptor of the current level
         let descriptor = &mut table[va.get_level_index(level)];
 
         // We reached the required granule level
         if Self::GRANULE_SIZES[level] == granule {
+            // Follow break-before-make sequence
+            descriptor.set_block_or_invalid_descriptor_to_invalid(level);
+            Self::invalidate(regime, Some(va));
             descriptor.set_block_descriptor(level, pa, attributes);
             return;
         }
@@ -607,6 +613,7 @@
                     level + 1,
                     unsafe { descriptor.get_next_level_table_mut(level) },
                     page_pool,
+                    regime,
                 )
             }
             DescriptorType::Block => {
@@ -638,12 +645,13 @@
                         level + 1,
                         unsafe { descriptor.get_next_level_table_mut(level) },
                         page_pool,
+                        regime,
                     )
                 }
 
                 // Invoke self to continue recursion on the newly created level
                 Self::set_block_descriptor_recursively(
-                    attributes, pa, va, granule, level, table, page_pool,
+                    attributes, pa, va, granule, level, table, page_pool, regime,
                 );
             }
             DescriptorType::Table => Self::set_block_descriptor_recursively(
@@ -654,6 +662,7 @@
                 level + 1,
                 unsafe { descriptor.get_next_level_table_mut(level) },
                 page_pool,
+                regime,
             ),
         }
     }
@@ -668,6 +677,7 @@
             1,
             self.base_table.descriptors.as_mut_slice(),
             &self.page_pool,
+            &self.regime,
         );
     }
 
@@ -685,6 +695,7 @@
         level: usize,
         table: &mut [Descriptor],
         page_pool: &PagePool,
+        regime: &TranslationRegime,
     ) {
         // Get descriptor of the current level
         let descriptor = &mut table[va.get_level_index(level)];
@@ -692,6 +703,7 @@
         // We reached the required granule level
         if Self::GRANULE_SIZES[level] == granule {
             descriptor.set_block_descriptor_to_invalid(level);
+            Self::invalidate(regime, Some(va));
             return;
         }
 
@@ -711,6 +723,7 @@
                     level + 1,
                     next_level_table,
                     page_pool,
+                    regime,
                 );
 
                 if next_level_table.iter().all(|d| !d.is_valid()) {
@@ -757,6 +770,69 @@
             }
         }
     }
+
+    fn invalidate(regime: &TranslationRegime, va: Option<VirtualAddress>) {
+        // SAFETY: The assembly code invalidates the translation table entry of
+        // the VA or all entries of the translation regime.
+        #[cfg(target_arch = "aarch64")]
+        unsafe {
+            if let Some(VirtualAddress(va)) = va {
+                match regime {
+                    TranslationRegime::EL1_0(_, _) => {
+                        core::arch::asm!(
+                        "tlbi vaae1is, {0}
+                        dsb nsh
+                        isb",
+                        in(reg) va)
+                    }
+                    #[cfg(target_feature = "vh")]
+                    TranslationRegime::EL2_0(_, _) => {
+                        core::arch::asm!(
+                        "tlbi vaae1is, {0}
+                        dsb nsh
+                        isb",
+                        in(reg) va)
+                    }
+                    TranslationRegime::EL2 => core::arch::asm!(
+                        "tlbi vae2is, {0}
+                        dsb nsh
+                        isb",
+                        in(reg) va),
+                    TranslationRegime::EL3 => core::arch::asm!(
+                        "tlbi vae3is, {0}
+                        dsb nsh
+                        isb",
+                        in(reg) va),
+                }
+            } else {
+                match regime {
+                    TranslationRegime::EL1_0(_, asid) => core::arch::asm!(
+                        "tlbi aside1, {0}
+                        dsb nsh
+                        isb",
+                        in(reg) (*asid as u64) << 48
+                    ),
+                    #[cfg(target_feature = "vh")]
+                    TranslationRegime::EL2_0(_, asid) => core::arch::asm!(
+                        "tlbi aside1, {0}
+                        dsb nsh
+                        isb",
+                        in(reg) (*asid as u64) << 48
+                    ),
+                    TranslationRegime::EL2 => core::arch::asm!(
+                        "tlbi alle2
+                        dsb nsh
+                        isb"
+                    ),
+                    TranslationRegime::EL3 => core::arch::asm!(
+                        "tlbi alle3
+                        dsb nsh
+                        isb"
+                    ),
+                }
+            }
+        }
+    }
 }
 
 #[cfg(test)]