More table safety improvements

cranelift/entity/src/primary.rs

@@ -7,7 +7,6 @@ use alloc::boxed::Box;
 use alloc::vec::Vec;
 use core::fmt;
 use core::marker::PhantomData;
-use core::mem;
 use core::ops::{Index, IndexMut};
 use core::slice;
 #[cfg(feature = "enable-serde")]
@@ -163,35 +162,11 @@ where
     ///
     /// Returns an error if an element does not exist, or if the same key was passed more than
     /// once.
-    // This implementation is taken from the unstable `get_many_mut`.
-    //
-    // Once it has been stabilised we can call that method directly.
-    pub fn get_many_mut<const N: usize>(
+    pub fn get_disjoint_mut<const N: usize>(
         &mut self,
         indices: [K; N],
-    ) -> Result<[&mut V; N], GetManyMutError<K>> {
-        for (i, &idx) in indices.iter().enumerate() {
-            if idx.index() >= self.len() {
-                return Err(GetManyMutError::DoesNotExist(idx));
-            }
-            for &idx2 in &indices[..i] {
-                if idx == idx2 {
-                    return Err(GetManyMutError::MultipleOf(idx));
-                }
-            }
-        }
-
-        let slice: *mut V = self.elems.as_mut_ptr();
-        let mut arr: mem::MaybeUninit<[&mut V; N]> = mem::MaybeUninit::uninit();
-        let arr_ptr = arr.as_mut_ptr();
-
-        unsafe {
-            for i in 0..N {
-                let idx = *indices.get_unchecked(i);
-                *(*arr_ptr).get_unchecked_mut(i) = &mut *slice.add(idx.index());
-            }
-            Ok(arr.assume_init())
-        }
+    ) -> Result<[&mut V; N], slice::GetDisjointMutError> {
+        self.elems.get_disjoint_mut(indices.map(|k| k.index()))
     }
 
     /// Performs a binary search on the values with a key extraction function.
@@ -236,12 +211,6 @@ where
     }
 }
 
-#[derive(Debug, PartialEq, Eq, Clone, Copy)]
-pub enum GetManyMutError<K> {
-    DoesNotExist(K),
-    MultipleOf(K),
-}
-
 impl<K, V> Default for PrimaryMap<K, V>
 where
     K: EntityRef,
@@ -557,14 +526,14 @@ mod tests {
         let _1 = m.push(1);
         let _2 = m.push(2);
 
-        assert_eq!([&mut 0, &mut 2], m.get_many_mut([_0, _2]).unwrap());
+        assert_eq!([&mut 0, &mut 2], m.get_disjoint_mut([_0, _2]).unwrap());
         assert_eq!(
-            m.get_many_mut([_0, _0]),
-            Err(GetManyMutError::MultipleOf(_0))
+            m.get_disjoint_mut([_0, _0]),
+            Err(slice::GetDisjointMutError::OverlappingIndices)
         );
         assert_eq!(
-            m.get_many_mut([E(4)]),
-            Err(GetManyMutError::DoesNotExist(E(4)))
+            m.get_disjoint_mut([E(4)]),
+            Err(slice::GetDisjointMutError::IndexOutOfBounds)
         );
     }
 }

crates/wasmtime/src/runtime/externals/table.rs

@@ -2,7 +2,7 @@ use crate::prelude::*;
 use crate::runtime::vm::{self as runtime, GcStore};
 use crate::store::{AutoAssertNoGc, StoreInstanceId, StoreOpaque};
 use crate::trampoline::generate_table_export;
-use crate::{AnyRef, AsContext, AsContextMut, ExternRef, Func, HeapType, Ref, TableType};
+use crate::{AnyRef, AsContext, AsContextMut, ExternRef, Func, HeapType, Ref, TableType, Trap};
 use core::iter;
 use wasmtime_environ::DefinedTableIndex;
 
@@ -344,24 +344,74 @@ impl Table {
                  destination table's element type",
             )?;
 
-        let (dst_table, _) = dst_table.wasmtime_table(store, iter::empty());
-        // FIXME(#11179) shouldn't need to subvert the borrow checker
-        let dst_table: *mut _ = dst_table;
-        let src_range = src_index..(src_index.checked_add(len).unwrap_or(u64::MAX));
-        let (src_table, _) = src_table.wasmtime_table(store, src_range);
-        // FIXME(#11179) shouldn't need to subvert the borrow checker
-        let src_table: *mut _ = src_table;
+        // SAFETY: the the two tables have the same type, as type-checked above.
         unsafe {
-            runtime::Table::copy(
-                store.optional_gc_store_mut(),
-                dst_table,
-                src_table,
+            Self::copy_raw(store, dst_table, dst_index, src_table, src_index, len)?;
+        }
+        Ok(())
+    }
+
+    /// Copies the elements of `src_table` to `dst_table`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the either table doesn't belong to `store`.
+    ///
+    /// # Safety
+    ///
+    /// Requires that the two tables have previously been type-checked to have
+    /// the same type.
+    pub(crate) unsafe fn copy_raw(
+        store: &mut StoreOpaque,
+        dst_table: &Table,
+        dst_index: u64,
+        src_table: &Table,
+        src_index: u64,
+        len: u64,
+    ) -> Result<(), Trap> {
+        // Handle lazy initialization of the source table first before doing
+        // anything else.
+        let src_range = src_index..(src_index.checked_add(len).unwrap_or(u64::MAX));
+        src_table.wasmtime_table(store, src_range);
+
+        // validate `dst_table` belongs to `store`.
+        dst_table.wasmtime_table(store, iter::empty());
+
+        // Figure out which of the three cases we're in:
+        //
+        // 1. Cross-instance table copy.
+        // 2. Intra-instance table copy.
+        // 3. Intra-table copy.
+        let src_instance = src_table.instance.instance();
+        let dst_instance = dst_table.instance.instance();
+        if src_instance != dst_instance {
+            // SAFETY: accessing two instances requires only accessing defined items
+            // on each instance which is done below with `get_defined_*` methods.
+            let (gc_store, [src_instance, dst_instance]) =
+                unsafe { store.optional_gc_store_and_instances_mut([src_instance, dst_instance]) };
+            src_instance.get_defined_table(src_table.index).copy_to(
+                dst_instance.get_defined_table(dst_table.index),
+                gc_store,
                 dst_index,
                 src_index,
                 len,
-            )?;
+            )
+        } else if src_table.index != dst_table.index {
+            assert_eq!(src_instance, dst_instance);
+            let (gc_store, instance) = store.optional_gc_store_and_instance_mut(src_instance);
+            let [(_, src_table), (_, dst_table)] = instance
+                .tables_mut()
+                .get_disjoint_mut([src_table.index, dst_table.index])
+                .unwrap();
+            src_table.copy_to(dst_table, gc_store, dst_index, src_index, len)
+        } else {
+            assert_eq!(src_instance, dst_instance);
+            assert_eq!(src_table.index, dst_table.index);
+            let (gc_store, instance) = store.optional_gc_store_and_instance_mut(src_instance);
+            instance
+                .get_defined_table(src_table.index)
+                .copy_within(gc_store, dst_index, src_index, len)
         }
-        Ok(())
     }
 
     /// Fill `table[dst..(dst + len)]` with the given value.

crates/wasmtime/src/runtime/store.rs

@@ -1305,6 +1305,32 @@ impl StoreOpaque {
         self.instances[id].handle.get_mut()
     }
 
+    /// Access multiple instances specified via `ids`.
+    ///
+    /// # Panics
+    ///
+    /// This method will panic if any indices in `ids` overlap.
+    ///
+    /// # Safety
+    ///
+    /// This method is not safe if the returned instances are used to traverse
+    /// "laterally" between other instances. For example accessing imported
+    /// items in an instance may traverse laterally to a sibling instance thus
+    /// aliasing a returned value here. The caller must ensure that only defined
+    /// items within the instances themselves are accessed.
+    #[inline]
+    pub unsafe fn optional_gc_store_and_instances_mut<const N: usize>(
+        &mut self,
+        ids: [InstanceId; N],
+    ) -> (Option<&mut GcStore>, [Pin<&mut vm::Instance>; N]) {
+        let instances = self
+            .instances
+            .get_disjoint_mut(ids)
+            .unwrap()
+            .map(|h| h.handle.get_mut());
+        (self.gc_store.as_mut(), instances)
+    }
+
     /// Pair of `Self::optional_gc_store_mut` and `Self::instance_mut`
     pub fn optional_gc_store_and_instance_mut(
         &mut self,

crates/wasmtime/src/runtime/vm/debug_builtins.rs

@@ -1,6 +1,6 @@
 #![doc(hidden)]
 
-use crate::runtime::vm::instance::Instance;
+use crate::runtime::vm::instance::InstanceAndStore;
 use crate::runtime::vm::vmcontext::VMContext;
 use core::ptr::NonNull;
 use wasmtime_environ::{EntityRef, MemoryIndex};
@@ -18,7 +18,8 @@ pub unsafe extern "C" fn resolve_vmctx_memory_ptr(p: *const u32) -> *const u8 {
         VMCTX_AND_MEMORY.0 != NonNull::dangling(),
         "must call `__vmctx->set()` before resolving Wasm pointers"
     );
-    Instance::from_vmctx(VMCTX_AND_MEMORY.0, |handle| {
+    InstanceAndStore::from_vmctx(VMCTX_AND_MEMORY.0, |handle| {
+        let (handle, _) = handle.unpack_mut();
         assert!(
             VMCTX_AND_MEMORY.1 < handle.env_module().memories.len(),
             "memory index for debugger is out of bounds"