Use safe-mmio crate for 8250 UART driver.

This takes care of working around the Rust volatile write issue on
aarch64 and also lets us remove some unsafe code here.

Test: atest vmbase_example.integration_test
Change-Id: I73e2d16ed55b134d5851f94eb226a1e9104da495

diff --git a/libs/libvmbase/Android.bp b/libs/libvmbase/Android.bp
index 2d0294e..c02466e 100644
--- a/libs/libvmbase/Android.bp
+++ b/libs/libvmbase/Android.bp
@@ -83,6 +83,7 @@
         "liblibfdt_nostd",
         "liblog_rust_nostd",
         "libonce_cell_nostd",
+        "libsafe_mmio",
         "libspin_nostd",
         "libstatic_assertions",
         "libthiserror_nostd",
diff --git a/libs/libvmbase/src/arch.rs b/libs/libvmbase/src/arch.rs
index cc42939..29d3a32 100644
--- a/libs/libvmbase/src/arch.rs
+++ b/libs/libvmbase/src/arch.rs
@@ -38,25 +38,6 @@
 #[cfg(target_arch = "aarch64")]
 pub use aarch64_paging::paging::VirtualAddress;
 
-/// Write with well-defined compiled behavior.
-///
-/// See https://github.com/rust-lang/rust/issues/131894
-///
-/// # Safety
-///
-/// `dst` must be valid for writes.
-#[inline]
-pub unsafe fn write_volatile_u8(dst: *mut u8, src: u8) {
-    cfg_if::cfg_if! {
-        if #[cfg(target_arch = "aarch64")] {
-            // SAFETY: `dst` is valid for writes.
-            unsafe { aarch64::strb(dst, src) }
-        } else {
-            compile_error!("Unsupported target_arch")
-        }
-    }
-}
-
 /// Flush `size` bytes of data cache by virtual address.
 #[inline]
 pub(crate) fn flush_region(start: usize, size: usize) {
diff --git a/libs/libvmbase/src/arch/aarch64.rs b/libs/libvmbase/src/arch/aarch64.rs
index 9c4d256..de3bfa8 100644
--- a/libs/libvmbase/src/arch/aarch64.rs
+++ b/libs/libvmbase/src/arch/aarch64.rs
@@ -119,26 +119,6 @@
     }};
 }
 
-/// STRB intrinsics.
-///
-/// See https://github.com/rust-lang/rust/issues/131894
-///
-/// # Safety
-///
-/// `dst` must be valid for writes.
-#[inline]
-pub unsafe fn strb(dst: *mut u8, src: u8) {
-    // SAFETY: strb only modifies *dst, which must be valid for writes.
-    unsafe {
-        core::arch::asm!(
-            "strb {value:w}, [{ptr}]",
-            value = in(reg) src,
-            ptr = in(reg) dst,
-            options(preserves_flags),
-        );
-    }
-}
-
 /// Reads the number of words in the smallest cache line of all the data caches and unified caches.
 #[inline]
 pub fn min_dcache_line_size() -> usize {
diff --git a/libs/libvmbase/src/arch/aarch64/uart.rs b/libs/libvmbase/src/arch/aarch64/uart.rs
index 2ef7d7e..0e94fc6 100644
--- a/libs/libvmbase/src/arch/aarch64/uart.rs
+++ b/libs/libvmbase/src/arch/aarch64/uart.rs
@@ -14,36 +14,29 @@
 
 //! Uart driver with backend for aarch64 using MMIO
 
-use crate::arch::write_volatile_u8;
 use crate::uart::UartBackend;
+use core::ptr::NonNull;
+use safe_mmio::{fields::ReadWrite, UniqueMmioPointer};
 
 /// Alias for default Uart for aarch64 backend with [`MmioBackend`]
-pub type Uart = crate::uart::Uart<MmioBackend>;
+pub type Uart = crate::uart::Uart<MmioBackend<'static>>;
 
-/// Backend for [`crate::uart::Uart`] that uses [`crate::arch::write_volatile_u8`] for writing to
-/// hardware registers.
-pub struct MmioBackend {
-    base_address: *mut u8,
+/// Backend for [`crate::uart::Uart`] that uses `safe-mmio` for writing to hardware registers.
+pub struct MmioBackend<'a> {
+    registers: UniqueMmioPointer<'a, [ReadWrite<u8>; 8]>,
 }
 
-impl MmioBackend {
-    /// Constructs a new instance of the UART driver backend for a device at the given base address.
-    ///
-    /// # Safety
-    ///
-    /// The given base address must point to the 8 MMIO control registers of an appropriate UART
-    /// device, which must be mapped into the address space of the process as device memory and not
-    /// have any other aliases.
-    pub unsafe fn new(base_address: usize) -> Self {
-        Self { base_address: base_address as *mut u8 }
+impl<'a> MmioBackend<'a> {
+    /// Constructs a new instance of the UART driver backend for a device with the given data
+    /// register.
+    pub fn new(registers: UniqueMmioPointer<'a, [ReadWrite<u8>; 8]>) -> Self {
+        Self { registers }
     }
 }
 
-impl UartBackend for MmioBackend {
-    fn write_register_u8(&self, offset: usize, byte: u8) {
-        // SAFETY: We know that the base address points to the control registers of a UART device
-        // which is appropriately mapped.
-        unsafe { write_volatile_u8(self.base_address.add(offset), byte) }
+impl UartBackend for MmioBackend<'_> {
+    fn write_register_u8(&mut self, offset: usize, byte: u8) {
+        self.registers.get(offset).expect("Register offset out of bounds").write(byte);
     }
 }
 
@@ -52,15 +45,15 @@
     ///
     /// # Safety
     ///
-    /// The given base address must point to the 8 MMIO control registers of an appropriate UART
-    /// device, which must be mapped into the address space of the process as device memory and not
-    /// have any other aliases.
+    /// The given base address must point to the 8 MMIO control registers of an appropriate 8250
+    /// UART device, which must be mapped into the address space of the process as device memory and
+    /// not have any other aliases.
     pub unsafe fn new(base_address: usize) -> Self {
-        // SAFETY: Delegated to caller
-        unsafe { Self::create(MmioBackend::new(base_address)) }
+        // SAFETY: The caller promises that base_address points to a mapped 8250 UART's registers
+        // with no aliases. That implies that there are 8 single-byte registers we can safely
+        // access, as required by `UniqueMmioPointer`.
+        let data_register =
+            unsafe { UniqueMmioPointer::new(NonNull::new(base_address as _).unwrap()) };
+        Self::create(MmioBackend::new(data_register))
     }
 }
-
-// SAFETY: `MmioBackend` just contains a pointer to device memory, which can be accessed from any
-// context.
-unsafe impl Send for MmioBackend {}
diff --git a/libs/libvmbase/src/uart.rs b/libs/libvmbase/src/uart.rs
index 857a22e..6a256ea 100644
--- a/libs/libvmbase/src/uart.rs
+++ b/libs/libvmbase/src/uart.rs
@@ -20,7 +20,7 @@
 /// The backend for [`Uart`] that abstracts the access to 8250 register map
 pub trait UartBackend {
     /// Writes a byte value on the offset to the hardware registers.
-    fn write_register_u8(&self, offset: usize, byte: u8);
+    fn write_register_u8(&mut self, offset: usize, byte: u8);
 }
 
 /// Minimal driver for an 8250 UART. This only implements enough to work with the emulated 8250
@@ -36,7 +36,7 @@
     }
 
     /// Writes a single byte to the UART.
-    pub fn write_byte(&self, byte: u8) {
+    pub fn write_byte(&mut self, byte: u8) {
         self.backend.write_register_u8(0, byte)
     }
 }