diff --git a/pvmfw/src/memory.rs b/pvmfw/src/memory.rs
index 604aa80..7eecb97 100644
--- a/pvmfw/src/memory.rs
+++ b/pvmfw/src/memory.rs
@@ -321,8 +321,7 @@
         handle_alloc_error(layout);
     };
 
-    let vaddr = buffer.as_ptr() as usize;
-    let paddr = virt_to_phys(vaddr);
+    let paddr = virt_to_phys(buffer);
     // If share_range fails then we will leak the allocation, but that seems better than having it
     // be reused while maybe still partially shared with the host.
     share_range(&(paddr..paddr + layout.size()), granule)?;
@@ -338,7 +337,7 @@
 ///
 /// The memory must have been allocated by `alloc_shared` with the same size, and not yet
 /// deallocated.
-pub unsafe fn dealloc_shared(vaddr: usize, size: usize) -> smccc::Result<()> {
+pub unsafe fn dealloc_shared(vaddr: NonNull<u8>, size: usize) -> smccc::Result<()> {
     let layout = shared_buffer_layout(size)?;
     let granule = layout.align();
 
@@ -346,7 +345,7 @@
     unshare_range(&(paddr..paddr + layout.size()), granule)?;
     // Safe because the memory was allocated by `alloc_shared` above using the same allocator, and
     // the layout is the same as was used then.
-    unsafe { dealloc(vaddr as *mut u8, layout) };
+    unsafe { dealloc(vaddr.as_ptr(), layout) };
 
     Ok(())
 }
@@ -372,8 +371,16 @@
 
 /// Returns the intermediate physical address corresponding to the given virtual address.
 ///
-/// As we use identity mapping for everything, this is just the identity function, but it's useful
-/// to use it to be explicit about where we are converting from virtual to physical address.
-pub fn virt_to_phys(vaddr: usize) -> usize {
-    vaddr
+/// As we use identity mapping for everything, this is just a cast, but it's useful to use it to be
+/// explicit about where we are converting from virtual to physical address.
+pub fn virt_to_phys(vaddr: NonNull<u8>) -> usize {
+    vaddr.as_ptr() as _
+}
+
+/// Returns a pointer for the virtual address corresponding to the given non-zero intermediate
+/// physical address.
+///
+/// Panics if `paddr` is 0.
+pub fn phys_to_virt(paddr: usize) -> NonNull<u8> {
+    NonNull::new(paddr as _).unwrap()
 }
diff --git a/pvmfw/src/virtio/hal.rs b/pvmfw/src/virtio/hal.rs
index c1c8ae6..c6c7a99 100644
--- a/pvmfw/src/virtio/hal.rs
+++ b/pvmfw/src/virtio/hal.rs
@@ -1,23 +1,22 @@
-use crate::memory::{alloc_shared, dealloc_shared, virt_to_phys};
+use crate::memory::{alloc_shared, dealloc_shared, phys_to_virt, virt_to_phys};
 use core::ptr::{copy_nonoverlapping, NonNull};
 use log::debug;
-use virtio_drivers::{BufferDirection, Hal, PhysAddr, VirtAddr, PAGE_SIZE};
+use virtio_drivers::{BufferDirection, Hal, PhysAddr, PAGE_SIZE};
 
 pub struct HalImpl;
 
 impl Hal for HalImpl {
-    fn dma_alloc(pages: usize) -> PhysAddr {
+    fn dma_alloc(pages: usize, _direction: BufferDirection) -> (PhysAddr, NonNull<u8>) {
         debug!("dma_alloc: pages={}", pages);
         let size = pages * PAGE_SIZE;
-        let vaddr = alloc_shared(size)
-            .expect("Failed to allocate and share VirtIO DMA range with host")
-            .as_ptr() as VirtAddr;
-        virt_to_phys(vaddr)
+        let vaddr =
+            alloc_shared(size).expect("Failed to allocate and share VirtIO DMA range with host");
+        let paddr = virt_to_phys(vaddr);
+        (paddr, vaddr)
     }
 
-    fn dma_dealloc(paddr: PhysAddr, pages: usize) -> i32 {
+    fn dma_dealloc(paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
         debug!("dma_dealloc: paddr={:#x}, pages={}", paddr, pages);
-        let vaddr = Self::phys_to_virt(paddr);
         let size = pages * PAGE_SIZE;
         // Safe because the memory was allocated by `dma_alloc` above using the same allocator, and
         // the layout is the same as was used then.
@@ -27,8 +26,8 @@
         0
     }
 
-    fn phys_to_virt(paddr: PhysAddr) -> VirtAddr {
-        paddr
+    fn mmio_phys_to_virt(paddr: PhysAddr, _size: usize) -> NonNull<u8> {
+        phys_to_virt(paddr)
     }
 
     fn share(buffer: NonNull<[u8]>, direction: BufferDirection) -> PhysAddr {
@@ -43,11 +42,11 @@
                 copy_nonoverlapping(buffer.as_ptr() as *mut u8, copy.as_ptr(), size);
             }
         }
-        virt_to_phys(copy.as_ptr() as VirtAddr)
+        virt_to_phys(copy)
     }
 
     fn unshare(paddr: PhysAddr, buffer: NonNull<[u8]>, direction: BufferDirection) {
-        let vaddr = Self::phys_to_virt(paddr);
+        let vaddr = phys_to_virt(paddr);
         let size = buffer.len();
         if direction == BufferDirection::DeviceToDriver {
             debug!(
@@ -56,7 +55,7 @@
                 buffer.as_ptr() as *mut u8 as usize
             );
             unsafe {
-                copy_nonoverlapping(vaddr as *const u8, buffer.as_ptr() as *mut u8, size);
+                copy_nonoverlapping(vaddr.as_ptr(), buffer.as_ptr() as *mut u8, size);
             }
         }
 
diff --git a/vmbase/example/src/pci.rs b/vmbase/example/src/pci.rs
index 438ff9e..c0a2d2b 100644
--- a/vmbase/example/src/pci.rs
+++ b/vmbase/example/src/pci.rs
@@ -15,7 +15,7 @@
 //! Functions to scan the PCI bus for VirtIO device.
 
 use aarch64_paging::paging::MemoryRegion;
-use alloc::alloc::{alloc, dealloc, Layout};
+use alloc::alloc::{alloc, dealloc, handle_alloc_error, Layout};
 use core::{mem::size_of, ptr::NonNull};
 use fdtpci::PciInfo;
 use log::{debug, info};
@@ -25,7 +25,7 @@
         pci::{bus::PciRoot, virtio_device_type, PciTransport},
         DeviceType, Transport,
     },
-    BufferDirection, Hal, PhysAddr, VirtAddr, PAGE_SIZE,
+    BufferDirection, Hal, PhysAddr, PAGE_SIZE,
 };
 
 /// The standard sector size of a VirtIO block device, in bytes.
@@ -87,32 +87,34 @@
 struct HalImpl;
 
 impl Hal for HalImpl {
-    fn dma_alloc(pages: usize) -> PhysAddr {
+    fn dma_alloc(pages: usize, _direction: BufferDirection) -> (PhysAddr, NonNull<u8>) {
         debug!("dma_alloc: pages={}", pages);
         let layout = Layout::from_size_align(pages * PAGE_SIZE, PAGE_SIZE).unwrap();
         // Safe because the layout has a non-zero size.
-        let vaddr = unsafe { alloc(layout) } as VirtAddr;
-        virt_to_phys(vaddr)
+        let vaddr = unsafe { alloc(layout) };
+        let vaddr =
+            if let Some(vaddr) = NonNull::new(vaddr) { vaddr } else { handle_alloc_error(layout) };
+        let paddr = virt_to_phys(vaddr);
+        (paddr, vaddr)
     }
 
-    fn dma_dealloc(paddr: PhysAddr, pages: usize) -> i32 {
+    fn dma_dealloc(paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
         debug!("dma_dealloc: paddr={:#x}, pages={}", paddr, pages);
-        let vaddr = Self::phys_to_virt(paddr);
         let layout = Layout::from_size_align(pages * PAGE_SIZE, PAGE_SIZE).unwrap();
         // Safe because the memory was allocated by `dma_alloc` above using the same allocator, and
         // the layout is the same as was used then.
         unsafe {
-            dealloc(vaddr as *mut u8, layout);
+            dealloc(vaddr.as_ptr(), layout);
         }
         0
     }
 
-    fn phys_to_virt(paddr: PhysAddr) -> VirtAddr {
-        paddr
+    fn mmio_phys_to_virt(paddr: PhysAddr, _size: usize) -> NonNull<u8> {
+        NonNull::new(paddr as _).unwrap()
     }
 
     fn share(buffer: NonNull<[u8]>, _direction: BufferDirection) -> PhysAddr {
-        let vaddr = buffer.as_ptr() as *mut u8 as usize;
+        let vaddr = buffer.cast();
         // Nothing to do, as the host already has access to all memory.
         virt_to_phys(vaddr)
     }
@@ -123,6 +125,6 @@
     }
 }
 
-fn virt_to_phys(vaddr: VirtAddr) -> PhysAddr {
-    vaddr
+fn virt_to_phys(vaddr: NonNull<u8>) -> PhysAddr {
+    vaddr.as_ptr() as _
 }