Skip to content

Commit

Permalink
Add SplitByteSlice::{try_split_at, split_at_unchecked}
Browse files Browse the repository at this point in the history 
Given the importance of avoiding panic paths to some of our
customers, we should, as a rule, avoid only offering (or relying
upon) APIs with panic paths.

Related to #202
  • Loading branch information
@jswrenn
jswrenn committed Mar 25, 2024
1 parent 83c68a0 commit 748ec27
Show file tree
Hide file tree
Showing 2 changed files with 136 additions and 47 deletions.
145 changes: 98 additions & 47 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -325,14 +325,14 @@ use core::alloc::Layout;
#[doc(hidden)]
pub use crate::pointer::{Maybe, MaybeAligned, Ptr};

// For each polyfill, as soon as the corresponding feature is stable, the
// For each trait polyfill, as soon as the corresponding feature is stable, the
// polyfill import will be unused because method/function resolution will prefer
// the inherent method/function over a trait method/function. Thus, we suppress
// the `unused_imports` warning.
//
// See the documentation on `util::polyfills` for more information.
#[allow(unused_imports)]
use crate::util::polyfills::NonNullExt as _;
use crate::util::polyfills::{self, NonNullExt as _};

#[rustversion::nightly]
#[cfg(all(test, not(__INTERNAL_USE_ONLY_NIGHTLY_FEATURES_IN_TESTS)))]
Expand Down Expand Up @@ -4883,10 +4883,10 @@ where
#[must_use = "has no side effects"]
fn bikeshed_new_from_prefix_known_layout(bytes: B) -> Option<(Ref<B, T>, B)> {
let (_, split_at) = Ptr::from_ref(bytes.deref()).try_cast_into::<T>(CastType::Prefix)?;
let (bytes, suffix) = bytes.split_at(split_at);
let (bytes, suffix) = bytes.try_split_at(split_at)?;
// INVARIANTS: `try_cast_into` validates size and alignment, and returns
// a `split_at` that indicates how many bytes of `bytes` correspond to a
// valid `T`. By safety invariant on `SplitByteSlice`, `split_at` is
// valid `T`. By safety invariant on `SplitByteSlice`, `try_split_at` is
// implemented correctly, and so we can rely on it to produce the
// correct `bytes` and `suffix`.
Some((Ref(bytes, PhantomData), suffix))
Expand All @@ -4895,10 +4895,10 @@ where
#[must_use = "has no side effects"]
fn bikeshed_new_from_suffix_known_layout(bytes: B) -> Option<(B, Ref<B, T>)> {
let (_, split_at) = Ptr::from_ref(bytes.deref()).try_cast_into::<T>(CastType::Suffix)?;
let (prefix, bytes) = bytes.split_at(split_at);
let (prefix, bytes) = bytes.try_split_at(split_at)?;
// INVARIANTS: `try_cast_into` validates size and alignment, and returns
// a `split_at` that indicates how many bytes of `bytes` correspond to a
// valid `T`. By safety invariant on `SplitByteSlice`, `split_at` is
// valid `T`. By safety invariant on `SplitByteSlice`, `try_split_at` is
// implemented correctly, and so we can rely on it to produce the
// correct `prefix` and `bytes`.
Some((prefix, Ref(bytes, PhantomData)))
Expand Down Expand Up @@ -4928,30 +4928,30 @@ where
if bytes.len() < mem::size_of::<T>() || !util::aligned_to::<_, T>(bytes.deref()) {
return None;
}
let (bytes, suffix) = bytes.split_at(mem::size_of::<T>());
let (bytes, suffix) = bytes.try_split_at(mem::size_of::<T>())?;
// INVARIANTS: We just validated alignment and that `bytes` is at least
// as large as `T`. `bytes.split_at(mem::size_of::<T>())` ensures that
// the new `bytes` is exactly the size of `T`. By safety invariant on
// `SplitByteSlice`, `split_at` is implemented correctly, and so we can
// rely on it to produce the correct `bytes` and `suffix`.
// as large as `T`. `bytes.try_split_at(mem::size_of::<T>())?` ensures
// that the new `bytes` is exactly the size of `T`. By safety invariant
// on `SplitByteSlice`, `try_split_at` is implemented correctly, and so
// we can rely on it to produce the correct `bytes` and `suffix`.
Some((Ref(bytes, PhantomData), suffix))
}

#[must_use = "has no side effects"]
fn new_sized_from_suffix(bytes: B) -> Option<(B, Ref<B, T>)> {
let bytes_len = bytes.len();
let split_at = bytes_len.checked_sub(mem::size_of::<T>())?;
let (prefix, bytes) = bytes.split_at(split_at);
let (prefix, bytes) = bytes.try_split_at(split_at)?;
if !util::aligned_to::<_, T>(bytes.deref()) {
return None;
}
// INVARIANTS: Since `split_at` is defined as `bytes_len -
// size_of::<T>()`, the `bytes` which results from `let (prefix, bytes)
// = bytes.split_at(split_at)` has length `size_of::<T>()`. After
// = bytes.try_split_at(split_at)?` has length `size_of::<T>()`. After
// constructing `bytes`, we validate that it has the proper alignment.
// By safety invariant on `SplitByteSlice`, `split_at` is implemented
// correctly, and so we can rely on it to produce the correct `prefix`
// and `bytes`.
// By safety invariant on `SplitByteSlice`, `try_split_at` is
// implemented correctly, and so we can rely on it to produce the
// correct `prefix` and `bytes`.
Some((prefix, Ref(bytes, PhantomData)))
}
}
Expand Down Expand Up @@ -4991,10 +4991,10 @@ where
#[inline]
pub fn new_from_prefix(bytes: B) -> Option<(Ref<B, T>, B)> {
let (_, split_at) = Ptr::from_ref(bytes.deref()).try_cast_into::<T>(CastType::Prefix)?;
let (bytes, suffix) = bytes.split_at(split_at);
let (bytes, suffix) = bytes.try_split_at(split_at)?;
// INVARIANTS: `try_cast_into` validates size and alignment, and returns
// a `split_at` that indicates how many bytes of `bytes` correspond to a
// valid `T`. By safety invariant on `SplitByteSlice`, `split_at` is
// valid `T`. By safety invariant on `SplitByteSlice`, `try_split_at` is
// implemented correctly, and so we can rely on it to produce the
// correct `bytes` and `suffix`.
Some((Ref(bytes, PhantomData), suffix))
Expand All @@ -5012,12 +5012,12 @@ where
#[inline]
pub fn new_from_suffix(bytes: B) -> Option<(B, Ref<B, T>)> {
let (_, split_at) = Ptr::from_ref(bytes.deref()).try_cast_into::<T>(CastType::Suffix)?;
let (prefix, bytes) = bytes.split_at(split_at);
let (prefix, bytes) = bytes.try_split_at(split_at)?;
// INVARIANTS: `try_cast_into` validates size and alignment, and returns
// a `split_at` that indicates how many bytes of `bytes` correspond to a
// valid `T`. By safety invariant on `SplitByteSlice`, `split_at` is
// implemented correctly, and so we can rely on it to produce the
// correct `prefix` and `bytes`.
// a `try_split_at` that indicates how many bytes of `bytes` correspond
// to a valid `T`. By safety invariant on `SplitByteSlice`,
// `try_split_at` is implemented correctly, and so we can rely on it to
// produce the correct `prefix` and `bytes`.
Some((prefix, Ref(bytes, PhantomData)))
}
}
Expand Down Expand Up @@ -5049,7 +5049,7 @@ where
if bytes.len() < expected_len {
return None;
}
let (prefix, bytes) = bytes.split_at(expected_len);
let (prefix, bytes) = bytes.try_split_at(expected_len)?;
Self::new(prefix).map(move |l| (l, bytes))
}

Expand All @@ -5073,7 +5073,7 @@ where
None => return None,
};
let split_at = bytes.len().checked_sub(expected_len)?;
let (bytes, suffix) = bytes.split_at(split_at);
let (bytes, suffix) = bytes.try_split_at(split_at)?;
Self::new(suffix).map(move |l| (bytes, l))
}
}
Expand Down Expand Up @@ -5671,11 +5671,16 @@ pub unsafe trait ByteSliceMut: ByteSlice + DerefMut {}
///
/// # Safety
///
/// Unsafe code may depend for its soundness on the assumption that `split_at`
/// is implemented correctly. In particular, given `B: SplitByteSlice` and `b:
/// B`, if `b.deref()` returns a byte slice with address `addr` and length
/// `len`, then if `split <= len`, `b.split_at(split)` will return `(first,
/// second)` such that:
/// Unsafe code may depend for its soundness on the assumption that `split_at`,
/// `try_split_at` and `split_at_unchecked` are implemented correctly. In
/// particular, given `B: SplitByteSlice` and `b: B`, if `b.deref()` returns a
/// byte slice with address `addr` and length `len`, then if `split <= len`, all
/// three of these invocations:
/// - `b.split_at(split)`
/// - `b.split_at_unchecked(split)`
/// - `b.try_split_at(split)`
///
/// ...will return `(first, second)` such that:
/// - `first`'s address is `addr` and its length is `split`
/// - `second`'s address is `addr + split` and its length is `len - split`
pub unsafe trait SplitByteSlice: ByteSlice {
Expand All @@ -5687,7 +5692,39 @@ pub unsafe trait SplitByteSlice: ByteSlice {
///
/// `x.split_at(mid)` panics if `mid > x.deref().len()`.
#[must_use]
fn split_at(self, mid: usize) -> (Self, Self);
#[inline]
fn split_at(self, mid: usize) -> (Self, Self) {
if let Some(splits) = self.try_split_at(mid) {
splits
} else {
panic!("mid > len")
}
}

/// Splits the slice at the midpoint, with minimal runtime checks.
///
/// `x.split_at_unchecked(mid)` returns `x[..mid]` and `x[mid..]`.
///
/// # Safety
///
/// `mid` must not be greater than `x.deref().len()`
#[must_use]
unsafe fn split_at_unchecked(self, mid: usize) -> (Self, Self);

/// Attempts to split the slice at the midpoint.
///
/// `x.try_split_at(mid)` returns `Some((x[..mid], x[mid..]))`, if `mid <=
/// x.deref().len()`; otherwise `None`.
#[must_use]
#[inline]
fn try_split_at(self, mid: usize) -> Option<(Self, Self)> {
if mid <= self.deref().len() {
// SAFETY: Above, we ensure that `mid <= self.deref().len()`.
unsafe { Some(self.split_at_unchecked(mid)) }
} else {
None
}
}
}

/// A shorthand for [`SplitByteSlice`] and [`ByteSliceMut`].
Expand Down Expand Up @@ -5722,12 +5759,14 @@ pub trait IntoByteSliceMut<'a>: ByteSliceMut + Into<&'a mut [u8]> {}
#[allow(clippy::undocumented_unsafe_blocks)]
unsafe impl<'a> ByteSlice for &'a [u8] {}

// SAFETY: This delegates to the stdlib implementation, which is assumed to be
// correct.
// SAFETY: This delegates to `polyfills:split_at_unchecked`, which is assumed to
// be correct.
unsafe impl<'a> SplitByteSlice for &'a [u8] {
#[inline]
fn split_at(self, mid: usize) -> (Self, Self) {
<[u8]>::split_at(self, mid)
unsafe fn split_at_unchecked(self, mid: usize) -> (Self, Self) {
// SAFETY: By precondition on caller, `mid` is not greater than
// `self.len()`.
unsafe { polyfills::split_at_unchecked(self, mid) }
}
}

Expand All @@ -5737,12 +5776,14 @@ impl<'a> IntoByteSlice<'a> for &'a [u8] {}
#[allow(clippy::undocumented_unsafe_blocks)]
unsafe impl<'a> ByteSlice for &'a mut [u8] {}

// SAFETY: This delegates to the stdlib implementation, which is assumed to be
// correct.
// SAFETY: This delegates to the stdlib implementation of
// `polyfills::split_at_unchecked_mut`, which is assumed to be correct.
unsafe impl<'a> SplitByteSlice for &'a mut [u8] {
#[inline]
fn split_at(self, mid: usize) -> (Self, Self) {
<[u8]>::split_at_mut(self, mid)
unsafe fn split_at_unchecked(self, mid: usize) -> (Self, Self) {
// SAFETY: By precondition on caller, `mid` is not greater than
// `self.len()`.
unsafe { polyfills::split_at_mut_unchecked(self, mid) }
}
}

Expand All @@ -5752,25 +5793,35 @@ impl<'a> IntoByteSliceMut<'a> for &'a mut [u8] {}
#[allow(clippy::undocumented_unsafe_blocks)]
unsafe impl<'a> ByteSlice for cell::Ref<'a, [u8]> {}

// SAFETY: This delegates to stdlib implementations, which are assumed to be
// correct.
// SAFETY: This delegates to stdlib implementation of `Ref::map_split` and
// `polyfills::split_at_unchecked`, which are assumed to be correct.
unsafe impl<'a> SplitByteSlice for cell::Ref<'a, [u8]> {
#[inline]
fn split_at(self, mid: usize) -> (Self, Self) {
cell::Ref::map_split(self, |slice| <[u8]>::split_at(slice, mid))
unsafe fn split_at_unchecked(self, mid: usize) -> (Self, Self) {
cell::Ref::map_split(self, |slice|
// SAFETY: By precondition on caller, `mid` is not greater than
// `slice.len()`.
unsafe {
polyfills::split_at_unchecked(slice, mid)
})
}
}

// TODO(#429): Add a "SAFETY" comment and remove this `allow`.
#[allow(clippy::undocumented_unsafe_blocks)]
unsafe impl<'a> ByteSlice for RefMut<'a, [u8]> {}

// SAFETY: This delegates to stdlib implementations, which are assumed to be
// correct.
// SAFETY: This delegates to stdlib implementation of `Ref::map_split` and
// `polyfills::split_at_mut_unchecked`, which are assumed to be correct.
unsafe impl<'a> SplitByteSlice for RefMut<'a, [u8]> {
#[inline]
fn split_at(self, mid: usize) -> (Self, Self) {
RefMut::map_split(self, |slice| <[u8]>::split_at_mut(slice, mid))
unsafe fn split_at_unchecked(self, mid: usize) -> (Self, Self) {
RefMut::map_split(self, |slice|
// SAFETY: By precondition on caller, `mid` is not greater than
// `slice.len()`
unsafe {
polyfills::split_at_mut_unchecked(slice, mid)
})
}
}

Expand Down
38 changes: 38 additions & 0 deletions src/util.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -623,6 +623,44 @@ pub(crate) mod polyfills {
unsafe { NonNull::new_unchecked(ptr) }
}
}

/// Splits `bytes` at `mid` with minimal runtime checks.
///
/// # Safety
///
/// `mid` must not be greater than `bytes.len()`.
// TODO(#67): Once `slice::split_at_unchecked` is stabilized and available
// on our MSRV, remove this.
pub(crate) unsafe fn split_at_unchecked(bytes: &[u8], mid: usize) -> (&[u8], &[u8]) {
// SAFETY: By contract on caller, `mid` is not greater than
// `bytes.len()`.
unsafe { (<[u8]>::get_unchecked(bytes, ..mid), <[u8]>::get_unchecked(bytes, mid..)) }
}

/// Splits `bytes` at `mid` with minimal runtime checks.
///
/// # Safety
///
/// `mid` must not be greater than `bytes.len()`.
// TODO(#67): Once `slice::split_at_mut_unchecked` is stabilized and
// available on our MSRV, remove this.
pub(crate) unsafe fn split_at_mut_unchecked(
bytes: &mut [u8],
mid: usize,
) -> (&mut [u8], &mut [u8]) {
use core::slice::from_raw_parts_mut;

let l_ptr = bytes.as_mut_ptr();
// SAFETY: By contract on caller, `mid` is not greater than
// `bytes.len()`.
let r_ptr = unsafe { l_ptr.add(mid) };

let l_len = mid;
let r_len = bytes.len() - mid;

// SAFETY: TODO.
unsafe { (from_raw_parts_mut(l_ptr, l_len), from_raw_parts_mut(r_ptr, r_len)) }
}
}

#[cfg(test)]
Expand Down

0 comments on commit 748ec27

Please sign in to comment.