Use NonNull<u8> for pc instead of *mut u8 (#9274)

Copyright (c) 2024, Arm Limited.

Signed-off-by: Karl Meakin <karl.meakin@arm.com>

pulley/fuzz/src/interp.rs

@@ -32,7 +32,7 @@ pub fn interp(ops: Vec<Op>) {
         match vm.call(NonNull::from(&encoded[0]), args, rets.into_iter().copied()) {
             Ok(rets) => assert_eq!(rets.count(), 0),
             Err(pc) => {
-                let pc = pc as usize;
+                let pc = pc.as_ptr() as usize;
 
                 let start = &encoded[0] as *const u8 as usize;
                 let end = encoded.last().unwrap() as *const u8 as usize;

pulley/src/decode.rs

@@ -1,5 +1,7 @@
 //! Decoding support for pulley bytecode.
 
+use core::ptr::NonNull;
+
 use alloc::vec::Vec;
 use cranelift_bitset::scalar::ScalarBitSetStorage;
 use cranelift_bitset::ScalarBitSet;
@@ -194,7 +196,7 @@ pub enum Uninhabited {}
 ///
 /// This is a wrapper over a raw pointer to bytecode somewhere in memory.
 #[derive(Clone, Copy, Debug)]
-pub struct UnsafeBytecodeStream(*mut u8);
+pub struct UnsafeBytecodeStream(NonNull<u8>);
 
 impl UnsafeBytecodeStream {
     /// Construct a new `UnsafeBytecodeStream` pointing at the given PC.
@@ -208,8 +210,7 @@ impl UnsafeBytecodeStream {
     /// new PC, this stream must not be used to read just after the
     /// unconditional jump instruction because there is no guarantee that that
     /// memory is part of the bytecode stream or not.
-    pub unsafe fn new(pc: *mut u8) -> Self {
-        assert!(!pc.is_null());
+    pub unsafe fn new(pc: NonNull<u8>) -> Self {
         UnsafeBytecodeStream(pc)
     }
 
@@ -222,20 +223,19 @@ impl UnsafeBytecodeStream {
     /// that the address at `self._as_ptr() + offset` contains valid Pulley
     /// bytecode.
     pub unsafe fn offset(&self, offset: isize) -> Self {
-        UnsafeBytecodeStream(self.0.offset(offset))
+        UnsafeBytecodeStream(NonNull::new_unchecked(self.0.as_ptr().offset(offset)))
     }
 
     /// Get this stream's underlying raw pointer.
-    pub fn as_ptr(&self) -> *mut u8 {
+    pub fn as_ptr(&self) -> NonNull<u8> {
         self.0
     }
 }
 
 impl BytecodeStream for UnsafeBytecodeStream {
     fn read<const N: usize>(&mut self) -> Result<[u8; N], Self::Error> {
-        debug_assert!(!self.0.is_null());
-        let bytes = unsafe { self.0.cast::<[u8; N]>().read() };
-        self.0 = unsafe { self.0.add(N) };
+        let bytes = unsafe { self.0.cast::<[u8; N]>().as_ptr().read() };
+        self.0 = unsafe { NonNull::new_unchecked(self.0.as_ptr().add(N)) };
         Ok(bytes)
     }
 

pulley/src/interp.rs

@@ -72,7 +72,7 @@ impl Vm {
         func: NonNull<u8>,
         args: &[Val],
         rets: impl IntoIterator<Item = RegType> + 'a,
-    ) -> Result<impl Iterator<Item = Val> + 'a, *mut u8> {
+    ) -> Result<impl Iterator<Item = Val> + 'a, NonNull<u8>> {
         // NB: make sure this method stays in sync with
         // `PulleyMachineDeps::compute_arg_locs`!
 
@@ -97,7 +97,7 @@ impl Vm {
             }
         }
 
-        self.run(func.as_ptr())?;
+        self.run(func)?;
 
         let mut x_rets = (0..16).map(|x| XReg::new_unchecked(x));
         let mut f_rets = (0..16).map(|f| FReg::new_unchecked(f));
@@ -119,7 +119,7 @@ impl Vm {
         }))
     }
 
-    unsafe fn run(&mut self, pc: *mut u8) -> Result<(), *mut u8> {
+    unsafe fn run(&mut self, pc: NonNull<u8>) -> Result<(), NonNull<u8>> {
         let mut visitor = InterpreterVisitor {
             state: &mut self.state,
             pc: UnsafeBytecodeStream::new(pc),
@@ -148,23 +148,25 @@ impl Vm {
 
     #[cold]
     #[inline(never)]
-    fn return_to_host(&self) -> Result<(), *mut u8> {
+    fn return_to_host(&self) -> Result<(), NonNull<u8>> {
         Ok(())
     }
 
     #[cold]
     #[inline(never)]
-    fn trap(&self, pc: *mut u8) -> Result<(), *mut u8> {
+    fn trap(&self, pc: NonNull<u8>) -> Result<(), NonNull<u8>> {
         // We are given the VM's PC upon having executed a trap instruction,
         // which is actually pointing to the next instruction after the
         // trap. Back the PC up to point exactly at the trap.
-        let trap_pc = unsafe { pc.byte_sub(ExtendedOpcode::ENCODED_SIZE_OF_TRAP) };
+        let trap_pc = unsafe {
+            NonNull::new_unchecked(pc.as_ptr().byte_sub(ExtendedOpcode::ENCODED_SIZE_OF_TRAP))
+        };
         Err(trap_pc)
     }
 
     #[cold]
     #[inline(never)]
-    fn host_call(&self) -> Result<(), *mut u8> {
+    fn host_call(&self) -> Result<(), NonNull<u8>> {
         todo!()
     }
 }
@@ -624,15 +626,15 @@ impl OpVisitor for InterpreterVisitor<'_> {
             Continuation::ReturnToHost
         } else {
             let return_addr = self.state[XReg::lr].get_ptr();
-            self.pc = unsafe { UnsafeBytecodeStream::new(return_addr) };
+            self.pc = unsafe { UnsafeBytecodeStream::new(NonNull::new_unchecked(return_addr)) };
             // log::trace!("returning to {return_addr:#p}");
             Continuation::Continue
         }
     }
 
     fn call(&mut self, offset: PcRelOffset) -> Self::Return {
-        let return_addr = u64::try_from(self.pc.as_ptr() as usize).unwrap();
-        self.state[XReg::lr].set_u64(return_addr);
+        let return_addr = self.pc.as_ptr();
+        self.state[XReg::lr].set_ptr(return_addr.as_ptr());
         self.pc_rel_jump(offset, 5)
     }
 

pulley/tests/all/interp.rs

@@ -13,7 +13,7 @@ fn encoded(ops: &[Op]) -> Vec<u8> {
     encoded
 }
 
-unsafe fn run(vm: &mut Vm, ops: &[Op]) -> Result<(), *mut u8> {
+unsafe fn run(vm: &mut Vm, ops: &[Op]) -> Result<(), NonNull<u8>> {
     let _ = env_logger::try_init();
     let ops = encoded(ops);
     let _ = vm.call(NonNull::from(&ops[..]).cast(), &[], [])?;