Add wasmi_collections::Stable{Arena,Vec}::get_mut_ptr method (#1992)

* add Stable{Arena,Vec}::get_mut_ptr method

* add unit tests for new method

Author: Robbepop

crates/collections/src/arena/stable_arena.rs

@@ -3,6 +3,7 @@ use core::{
     iter::{Enumerate, FusedIterator, repeat_with},
     marker::PhantomData,
     ops::{Index, IndexMut, Range},
+    ptr::NonNull,
 };
 
 /// An append-only [`Arena`] whose entities have stable addresses.
@@ -155,6 +156,21 @@ where
             .ok_or(ArenaError::KeyOutOfBounds)
     }
 
+    /// Returns a raw pointer to the entity at the given key.
+    ///
+    /// Unlike [`get_mut`](StableArena::get_mut), this never forms an intermediate `&mut T`, so the
+    /// pointer carries the entity's underlying allocation provenance.
+    ///
+    /// # Errors
+    ///
+    /// If the `key` is out of bounds.
+    #[inline]
+    pub fn get_mut_ptr(&mut self, key: Key) -> Result<NonNull<T>, ArenaError> {
+        self.items
+            .get_mut_ptr(key.into_usize())
+            .ok_or(ArenaError::KeyOutOfBounds)
+    }
+
     /// Returns exclusive references to the pair of entities at the given keys if any.
     ///
     /// # Errors
@@ -406,6 +422,19 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn get_mut_ptr_writes_through() {
+        let mut arena: Arena = (0..10).collect();
+        let ptr = arena.get_mut_ptr(3).unwrap();
+        // Safety: key `3` is in bounds and uniquely owned here.
+        unsafe { *ptr.as_ptr() = 30 };
+        assert_eq!(arena.get(3).unwrap(), &30);
+        assert!(matches!(
+            arena.get_mut_ptr(10),
+            Err(ArenaError::KeyOutOfBounds)
+        ));
+    }
+
     #[test]
     fn iter_yields_keys_and_is_double_ended() {
         let arena: Arena = (0..5).collect();

crates/collections/src/arena/stable_vec.rs

@@ -121,6 +121,21 @@ impl<T> StableVec<T> {
         Some(unsafe { &mut *ptr.as_ptr().add(slot) })
     }
 
+    /// Returns a raw pointer to the item at `index`, or `None` if out of bounds.
+    ///
+    /// Unlike [`get_mut`](StableVec::get_mut), this never forms an intermediate `&mut T`, so the
+    /// pointer carries the bucket allocation's own provenance.
+    #[inline]
+    pub fn get_mut_ptr(&mut self, index: usize) -> Option<NonNull<T>> {
+        if index >= self.len {
+            return None;
+        }
+        let (bucket_index, slot) = Self::locate(index);
+        // Safety: `index < len` implies the bucket is allocated and `slot` is initialized.
+        let ptr = unsafe { self.buckets[bucket_index].unwrap_unchecked() };
+        Some(unsafe { NonNull::new_unchecked(ptr.as_ptr().add(slot)) })
+    }
+
     /// Returns exclusive references to the items at `a` and `b`.
     ///
     /// # Errors
@@ -482,6 +497,23 @@ mod tests {
         assert_eq!(vector.get_mut(100), None);
     }
 
+    #[test]
+    fn get_mut_ptr_writes_through() {
+        let mut vector = StableVec::new();
+        for i in 0..100 {
+            vector.push(i);
+        }
+        for i in 0..100 {
+            let ptr = vector.get_mut_ptr(i).unwrap();
+            // Safety: `i < len`, so the slot is initialized and uniquely owned here.
+            unsafe { *ptr.as_ptr() *= 2 };
+        }
+        let collected: Vec<usize> = vector.iter().copied().collect();
+        let expected: Vec<usize> = (0..100).map(|i| i * 2).collect();
+        assert_eq!(collected, expected);
+        assert!(vector.get_mut_ptr(100).is_none());
+    }
+
     #[test]
     fn iter_double_ended_and_exact_size() {
         let mut vector = StableVec::new();