review: add UnsafeBorrowResult

crates/wiggle/src/lib.rs

@@ -528,7 +528,11 @@ impl<'a, T> GuestPtr<'a, [T]> {
     where
         T: GuestTypeTransparent<'a>,
     {
-        self.as_unsafe_slice_mut()?.shared_borrow()
+        match self.as_unsafe_slice_mut()?.shared_borrow() {
+            UnsafeBorrowResult::Ok(slice) => Ok(Some(slice)),
+            UnsafeBorrowResult::Shared(_) => Ok(None),
+            UnsafeBorrowResult::Err(e) => Err(e),
+        }
     }
 
     /// Attempts to create a [`GuestSliceMut<'_, T>`] from this pointer,
@@ -548,7 +552,11 @@ impl<'a, T> GuestPtr<'a, [T]> {
     where
         T: GuestTypeTransparent<'a>,
     {
-        self.as_unsafe_slice_mut()?.mut_borrow()
+        match self.as_unsafe_slice_mut()?.mut_borrow() {
+            UnsafeBorrowResult::Ok(slice) => Ok(Some(slice)),
+            UnsafeBorrowResult::Shared(_) => Ok(None),
+            UnsafeBorrowResult::Err(e) => Err(e),
+        }
     }
 
     /// Similar to `as_slice_mut`, this function will attempt to create a smart
@@ -630,12 +638,18 @@ impl<'a, T> GuestPtr<'a, [T]> {
         if guest_slice.len != slice.len() {
             return Err(GuestError::SliceLengthsDiffer);
         }
-        // ... and copy the bytes. SAFETY:  we copy the same regardless of
-        // whether the memory is shared or non-shared and accept that the guest
-        // data may be concurrently modified. TODO: audit that the underlying
-        // use of `std::ptr::copy` is safe with shared memory
-        // (https://github.com/bytecodealliance/wasmtime/issues/4203)
-        unsafe { std::ptr::copy(slice.as_ptr(), guest_slice.ptr, guest_slice.len) };
+        // ... and copy the bytes.
+        match guest_slice.mut_borrow() {
+            UnsafeBorrowResult::Ok(mut dst) => dst.copy_from_slice(slice),
+            UnsafeBorrowResult::Shared(guest_slice) => {
+                // SAFETY:  in the shared memory case, we copy and accept that
+                // the guest data may be concurrently modified. 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(slice.as_ptr(), guest_slice.ptr, guest_slice.len) };
+            }
+            UnsafeBorrowResult::Err(e) => return Err(e),
+        }
         Ok(())
     }
 
@@ -893,17 +907,21 @@ impl<'a, T> UnsafeGuestSlice<'a, T> {
     /// This function is safe if and only if:
     /// - the memory is not shared (it will return `None` in this case) and
     /// - there are no overlapping mutable borrows for this region.
-    fn shared_borrow(self) -> Result<Option<GuestSlice<'a, T>>, GuestError> {
+    fn shared_borrow(self) -> UnsafeBorrowResult<GuestSlice<'a, T>, Self> {
         if self.mem.is_shared_memory() {
-            Ok(None)
+            UnsafeBorrowResult::Shared(self)
         } else {
-            let borrow = self.mem.shared_borrow(self.region)?;
-            let ptr = unsafe { slice::from_raw_parts(self.ptr, self.len) };
-            Ok(Some(GuestSlice {
-                ptr,
-                mem: self.mem,
-                borrow,
-            }))
+            match self.mem.shared_borrow(self.region) {
+                Ok(borrow) => {
+                    let ptr = unsafe { slice::from_raw_parts(self.ptr, self.len) };
+                    UnsafeBorrowResult::Ok(GuestSlice {
+                        ptr,
+                        mem: self.mem,
+                        borrow,
+                    })
+                }
+                Err(e) => UnsafeBorrowResult::Err(e),
+            }
         }
     }
 
@@ -915,21 +933,45 @@ impl<'a, T> UnsafeGuestSlice<'a, T> {
     /// - the memory is not shared (it will return `None` in this case) and
     /// - there are no overlapping borrows of any kind (shared or mutable) for
     ///   this region.
-    fn mut_borrow(self) -> Result<Option<GuestSliceMut<'a, T>>, GuestError> {
+    fn mut_borrow(self) -> UnsafeBorrowResult<GuestSliceMut<'a, T>, Self> {
         if self.mem.is_shared_memory() {
-            Ok(None)
+            UnsafeBorrowResult::Shared(self)
         } else {
-            let borrow = self.mem.mut_borrow(self.region)?;
-            let ptr = unsafe { slice::from_raw_parts_mut(self.ptr, self.len) };
-            Ok(Some(GuestSliceMut {
-                ptr,
-                mem: self.mem,
-                borrow,
-            }))
+            match self.mem.mut_borrow(self.region) {
+                Ok(borrow) => {
+                    let ptr = unsafe { slice::from_raw_parts_mut(self.ptr, self.len) };
+                    UnsafeBorrowResult::Ok(GuestSliceMut {
+                        ptr,
+                        mem: self.mem,
+                        borrow,
+                    })
+                }
+                Err(e) => UnsafeBorrowResult::Err(e),
+            }
         }
     }
 }
 
+/// A three-way result type for expressing that borrowing from an
+/// [`UnsafeGuestSlice`] could fail in multiple ways. Retaining the
+/// [`UnsafeGuestSlice`] in the `Shared` case allows us to reuse it.
+enum UnsafeBorrowResult<T, S> {
+    /// The borrow succeeded.
+    Ok(T),
+    /// The borrow failed because the underlying memory was shared--we cannot
+    /// safely borrow in this case and return the original unsafe slice.
+    Shared(S),
+    /// The borrow failed for some other reason, e.g., the region was already
+    /// borrowed.
+    Err(GuestError),
+}
+
+impl<T, S> From<GuestError> for UnsafeBorrowResult<T, S> {
+    fn from(e: GuestError) -> Self {
+        UnsafeBorrowResult::Err(e)
+    }
+}
+
 /// A smart pointer to an shareable `str` in guest memory.
 /// Usable as a `&'a str` via [`std::ops::Deref`].
 pub struct GuestStr<'a> {