Optimize some functions in the wiggle crate (#5566)

* wiggle: Inline some trivial functions

This commit marks a number of functions in wiggle as `#[inline]` as
they're otherwise trivial, mostly returning constants. This comes out of
some work I looked at recently with Andrew where some of these functions
showed up in profiles when they shouldn't.

* wiggle: Optimize the `GuestMemory` for shared memory

This commit implements a minor optimization to the `GuestMemory`
implementation for Wasmtime to skip most methods if a shared memory is
in play. Shared memories never get borrowed and this can be used to
internally skip some borrow-checker methods.

* wiggle: Optimize `GuestPtr::to_vec`

This commit replaces the safe implementation of `GuestPtr::to_vec` with
an unsafe implementation. The purpose of this is to speed up the
function when used with shared memory which otherwise performs a bunch
of atomic reads for types like `u8` which does validation-per-element
and isn't vectorizable. On a benchmark I was helping Andrew with this
sped up the host code enough to the point that guest code dwarfed the
execution time.

* Fix build

crates/wiggle/generate/src/types/flags.rs

@@ -42,6 +42,7 @@ pub(super) fn define_flags(
 
         impl TryFrom<#repr> for #ident {
             type Error = wiggle::GuestError;
+            #[inline]
             fn try_from(value: #repr) -> Result<Self, wiggle::GuestError> {
                 if #repr::from(!#ident::all()) & value != 0 {
                     Err(wiggle::GuestError::InvalidFlagValue(stringify!(#ident)))
@@ -53,22 +54,26 @@ pub(super) fn define_flags(
 
         impl TryFrom<#abi_repr> for #ident {
             type Error = wiggle::GuestError;
+            #[inline]
             fn try_from(value: #abi_repr) -> Result<Self, wiggle::GuestError> {
                 #ident::try_from(#repr::try_from(value)?)
             }
         }
 
         impl From<#ident> for #repr {
+            #[inline]
             fn from(e: #ident) -> #repr {
                 e.bits
             }
         }
 
         impl<'a> wiggle::GuestType<'a> for #ident {
+            #[inline]
             fn guest_size() -> u32 {
                 #repr::guest_size()
             }
 
+            #[inline]
             fn guest_align() -> usize {
                 #repr::guest_align()
             }

crates/wiggle/generate/src/types/handle.rs

@@ -14,29 +14,34 @@ pub(super) fn define_handle(name: &witx::Id, h: &witx::HandleDatatype) -> TokenS
         pub struct #ident(u32);
 
         impl #ident {
+            #[inline]
             pub unsafe fn inner(&self) -> u32 {
                 self.0
             }
         }
 
         impl From<#ident> for u32 {
+            #[inline]
             fn from(e: #ident) -> u32 {
                 e.0
             }
         }
 
         impl From<#ident> for i32 {
+            #[inline]
             fn from(e: #ident) -> i32 {
                 e.0 as i32
             }
         }
 
         impl From<u32> for #ident {
+            #[inline]
             fn from(e: u32) -> #ident {
                 #ident(e)
             }
         }
         impl From<i32> for #ident {
+            #[inline]
             fn from(e: i32) -> #ident {
                 #ident(e as u32)
             }
@@ -49,18 +54,22 @@ pub(super) fn define_handle(name: &witx::Id, h: &witx::HandleDatatype) -> TokenS
         }
 
         impl<'a> wiggle::GuestType<'a> for #ident {
+            #[inline]
             fn guest_size() -> u32 {
                 #size
             }
 
+            #[inline]
             fn guest_align() -> usize {
                 #align
             }
 
+            #[inline]
             fn read(location: &wiggle::GuestPtr<'a, #ident>) -> Result<#ident, wiggle::GuestError> {
                 Ok(#ident(u32::read(&location.cast())?))
             }
 
+            #[inline]
             fn write(location: &wiggle::GuestPtr<'_, Self>, val: Self) -> Result<(), wiggle::GuestError> {
                 u32::write(&location.cast(), val.0)
             }

crates/wiggle/generate/src/types/record.rs

@@ -87,10 +87,12 @@ pub(super) fn define_struct(name: &witx::Id, s: &witx::RecordDatatype) -> TokenS
         }
 
         impl<'a> wiggle::GuestType<'a> for #ident #struct_lifetime {
+            #[inline]
             fn guest_size() -> u32 {
                 #size
             }
 
+            #[inline]
             fn guest_align() -> usize {
                 #align
             }

crates/wiggle/generate/src/types/variant.rs

@@ -80,6 +80,7 @@ pub(super) fn define_variant(
         quote! {
             impl TryFrom<#tag_ty> for #ident {
                 type Error = wiggle::GuestError;
+                #[inline]
                 fn try_from(value: #tag_ty) -> Result<#ident, wiggle::GuestError> {
                     match value {
                         #(#tryfrom_repr_cases),*,
@@ -90,6 +91,7 @@ pub(super) fn define_variant(
 
             impl TryFrom<#abi_ty> for #ident {
                 type Error = wiggle::GuestError;
+                #[inline]
                 fn try_from(value: #abi_ty) -> Result<#ident, wiggle::GuestError> {
                     #ident::try_from(#tag_ty::try_from(value)?)
                 }
@@ -107,6 +109,7 @@ pub(super) fn define_variant(
         });
         quote! {
             impl From<#ident> for #tag_ty {
+                #[inline]
                 fn from(v: #ident) -> #tag_ty {
                     match v {
                         #(#from_repr_cases),*,
@@ -148,10 +151,12 @@ pub(super) fn define_variant(
         #enum_from
 
         impl<'a> wiggle::GuestType<'a> for #ident #enum_lifetime {
+            #[inline]
             fn guest_size() -> u32 {
                 #size
             }
 
+            #[inline]
             fn guest_align() -> usize {
                 #align
             }

crates/wiggle/src/guest_type.rs

@@ -61,7 +61,9 @@ pub unsafe trait GuestTypeTransparent<'a>: GuestType<'a> {}
 macro_rules! integer_primitives {
     ($([$ty:ident, $ty_atomic:ident],)*) => ($(
         impl<'a> GuestType<'a> for $ty {
+            #[inline]
             fn guest_size() -> u32 { mem::size_of::<Self>() as u32 }
+            #[inline]
             fn guest_align() -> usize { mem::align_of::<Self>() }
 
             #[inline]
@@ -122,7 +124,9 @@ macro_rules! integer_primitives {
 macro_rules! float_primitives {
     ($([$ty:ident, $ty_unsigned:ident, $ty_atomic:ident],)*) => ($(
         impl<'a> GuestType<'a> for $ty {
+            #[inline]
             fn guest_size() -> u32 { mem::size_of::<Self>() as u32 }
+            #[inline]
             fn guest_align() -> usize { mem::align_of::<Self>() }
 
             #[inline]
@@ -183,10 +187,12 @@ float_primitives! {
 
 // Support pointers-to-pointers where pointers are always 32-bits in wasm land
 impl<'a, T> GuestType<'a> for GuestPtr<'a, T> {
+    #[inline]
     fn guest_size() -> u32 {
         u32::guest_size()
     }
 
+    #[inline]
     fn guest_align() -> usize {
         u32::guest_align()
     }
@@ -206,10 +212,12 @@ impl<'a, T> GuestType<'a> for GuestPtr<'a, [T]>
 where
     T: GuestType<'a>,
 {
+    #[inline]
     fn guest_size() -> u32 {
         u32::guest_size() * 2
     }
 
+    #[inline]
     fn guest_align() -> usize {
         u32::guest_align()
     }

crates/wiggle/src/lib.rs

@@ -610,10 +610,22 @@ impl<'a, T> GuestPtr<'a, [T]> {
         T: GuestTypeTransparent<'a> + Copy + 'a,
     {
         let guest_slice = self.as_unsafe_slice_mut()?;
-        let mut vec = Vec::with_capacity(guest_slice.ptr.len());
+        let len = guest_slice.ptr.len();
-        for offs in 0..guest_slice.ptr.len() {
+        let mut vec = Vec::with_capacity(len);
-            let elem = self.get(offs as u32).expect("already validated the size");
+
-            vec.push(elem.read()?);
+        // SAFETY: The `guest_slice` variable is already a valid pointer into
+        // the guest's memory, and it may or may not be a pointer into shared
+        // memory. We can't naively use `.to_vec(..)` which could introduce data
+        // races but all that needs to happen is to copy data into our local
+        // `vec` as all the data is `Copy` and transparent anyway. For this
+        // purpose the `ptr::copy` function should be sufficient for copying
+        // over all the data.
+        //
+        // TODO: audit that this use of `std::ptr::copy` is safe with shared
+        // memory (https://github.com/bytecodealliance/wasmtime/issues/4203)
+        unsafe {
+            std::ptr::copy(guest_slice.ptr.as_ptr().cast::<T>(), vec.as_mut_ptr(), len);
+            vec.set_len(len);
         }
         Ok(vec)
     }

crates/wiggle/src/wasmtime.rs

@@ -51,30 +51,50 @@ impl<'a> WasmtimeGuestMemory<'a> {
 }
 
 unsafe impl GuestMemory for WasmtimeGuestMemory<'_> {
+    #[inline]
     fn base(&self) -> &[UnsafeCell<u8>] {
         self.mem
     }
+
+    // Note that this implementation has special cases for shared memory
+    // specifically because no regions of a shared memory can ever be borrowed.
+    // In the shared memory cases `shared_borrow` and `mut_borrow` are never
+    // called so that can be used to optimize the other methods by quickly
+    // checking a flag before calling the more expensive borrow-checker methods.
+
+    #[inline]
     fn has_outstanding_borrows(&self) -> bool {
-        self.bc.has_outstanding_borrows()
+        !self.shared && self.bc.has_outstanding_borrows()
     }
+    #[inline]
     fn is_shared_borrowed(&self, r: Region) -> bool {
-        self.bc.is_shared_borrowed(r)
+        !self.shared && self.bc.is_shared_borrowed(r)
     }
+    #[inline]
     fn is_mut_borrowed(&self, r: Region) -> bool {
-        self.bc.is_mut_borrowed(r)
+        !self.shared && self.bc.is_mut_borrowed(r)
     }
+    #[inline]
     fn shared_borrow(&self, r: Region) -> Result<BorrowHandle, GuestError> {
+        debug_assert!(!self.shared);
         self.bc.shared_borrow(r)
     }
+    #[inline]
     fn mut_borrow(&self, r: Region) -> Result<BorrowHandle, GuestError> {
+        debug_assert!(!self.shared);
         self.bc.mut_borrow(r)
     }
+    #[inline]
     fn shared_unborrow(&self, h: BorrowHandle) {
+        debug_assert!(!self.shared);
         self.bc.shared_unborrow(h)
     }
+    #[inline]
     fn mut_unborrow(&self, h: BorrowHandle) {
+        debug_assert!(!self.shared);
         self.bc.mut_unborrow(h)
     }
+    #[inline]
     fn is_shared_memory(&self) -> bool {
         self.shared
     }