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/descriptor.rs b/src/descriptor.rs
index bf10380..24ed8df 100644
--- a/src/descriptor.rs
+++ b/src/descriptor.rs
@@ -202,9 +202,17 @@
         self.set(Self::INVALID_DESCRIPTOR_VALUE)
     }
 
+    /// Set block or invalid descriptor to invalid
+    pub fn set_block_or_invalid_descriptor_to_invalid(&mut self, level: usize) {
+        assert!(level <= 3);
+        assert!(DescriptorType::Table != self.get_descriptor_type(level));
+
+        self.set(Self::INVALID_DESCRIPTOR_VALUE)
+    }
+
     /// Set table descriptor
     ///
-    /// **Unsafe**: The caller has to ensure that the passed next level table has the same life as
+    /// Safety: The caller has to ensure that the passed next level table has the same life as
     /// the descriptor.
     pub unsafe fn set_table_descriptor(
         &mut self,
@@ -232,7 +240,7 @@
 
     /// Get next level table
     ///
-    /// **Unsafe**: The returned next level table is based on the address read from the descriptor.
+    /// 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] {
         assert!(level <= 2);
@@ -245,7 +253,7 @@
 
     /// Get mutable next level table
     ///
-    /// **Unsafe**: The returned next level table is based on the address read from the descriptor.
+    /// 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] {
         assert!(level <= 2);
@@ -285,7 +293,13 @@
 
     /// Set raw descriptor value
     fn set(&mut self, value: u64) {
-        unsafe { ptr::write_volatile(self.cell.get(), value) }
+        // Safety: The cell should point to a valid address and the assembly code is just a data
+        // barrier.
+        unsafe {
+            ptr::write_volatile(self.cell.get(), value);
+            #[cfg(target_arch = "aarch64")]
+            core::arch::asm!("dsb nsh");
+        }
     }
 
     /// Modify raw descriptor value
@@ -556,6 +570,23 @@
     }
 
     #[test]
+    fn test_descriptor_block_or_invalid_block_to_invalid() {
+        let mut descriptor = Descriptor {
+            cell: UnsafeCell::new(0),
+        };
+
+        descriptor.set_block_or_invalid_descriptor_to_invalid(0);
+        assert_eq!(0, descriptor.get());
+
+        let mut descriptor = Descriptor {
+            cell: UnsafeCell::new(3),
+        };
+
+        descriptor.set_block_or_invalid_descriptor_to_invalid(3);
+        assert_eq!(0, descriptor.get());
+    }
+
+    #[test]
     #[should_panic]
     fn test_descriptor_level3_to_table() {
         let mut next_level_table = [Descriptor {