use super::*;
use crate::pointer::invariant::{Aligned, Exclusive, Invariants, Shared, Valid};
#[cfg_attr(feature = "derive", doc = "[derive]: zerocopy_derive::SplitAt")]
#[cfg_attr(
not(feature = "derive"),
doc = concat!("[derive]: https://docs.rs/zerocopy/", env!("CARGO_PKG_VERSION"), "/zerocopy/derive.SplitAt.html"),
)]
#[cfg_attr(
zerocopy_diagnostic_on_unimplemented_1_78_0,
diagnostic::on_unimplemented(note = "Consider adding `#[derive(SplitAt)]` to `{Self}`")
)]
pub unsafe trait SplitAt: KnownLayout<PointerMetadata = usize> {
type Elem;
#[doc(hidden)]
fn only_derive_is_allowed_to_implement_this_trait()
where
Self: Sized;
#[inline]
#[must_use]
unsafe fn split_at_unchecked(&self, l_len: usize) -> Split<&Self> {
unsafe { Split::<&Self>::new(self, l_len) }
}
#[inline]
#[must_use = "has no side effects"]
fn split_at(&self, l_len: usize) -> Option<Split<&Self>> {
MetadataOf::new_in_bounds(self, l_len).map(
#[inline(always)]
|l_len| {
unsafe { Split::new(self, l_len.get()) }
},
)
}
#[inline]
#[must_use]
unsafe fn split_at_mut_unchecked(&mut self, l_len: usize) -> Split<&mut Self> {
unsafe { Split::<&mut Self>::new(self, l_len) }
}
#[inline]
fn split_at_mut(&mut self, l_len: usize) -> Option<Split<&mut Self>> {
MetadataOf::new_in_bounds(self, l_len).map(
#[inline(always)]
|l_len| {
unsafe { Split::new(self, l_len.get()) }
},
)
}
}
unsafe impl<T> SplitAt for [T] {
type Elem = T;
#[inline]
#[allow(dead_code)]
fn only_derive_is_allowed_to_implement_this_trait()
where
Self: Sized,
{
}
}
#[derive(Debug)]
pub struct Split<T> {
source: T,
l_len: usize,
}
impl<T> Split<T> {
#[inline(always)]
unsafe fn new(source: T, l_len: usize) -> Self {
Self { source, l_len }
}
}
impl<'a, T> Split<&'a T>
where
T: ?Sized + SplitAt,
{
#[inline(always)]
fn into_ptr(self) -> Split<Ptr<'a, T, (Shared, Aligned, Valid)>> {
let source = Ptr::from_ref(self.source);
unsafe { Split::new(source, self.l_len) }
}
#[must_use = "has no side effects"]
#[inline(always)]
pub fn via_immutable(self) -> (&'a T, &'a [T::Elem])
where
T: Immutable,
{
let (l, r) = self.into_ptr().via_immutable();
(l.as_ref(), r.as_ref())
}
#[must_use = "has no side effects"]
#[inline(always)]
pub fn via_into_bytes(self) -> (&'a T, &'a [T::Elem])
where
T: IntoBytes,
{
let (l, r) = self.into_ptr().via_into_bytes();
(l.as_ref(), r.as_ref())
}
#[must_use = "has no side effects"]
#[inline(always)]
pub fn via_unaligned(self) -> (&'a T, &'a [T::Elem])
where
T: Unaligned,
{
let (l, r) = self.into_ptr().via_unaligned();
(l.as_ref(), r.as_ref())
}
#[must_use = "has no side effects"]
#[inline(always)]
pub fn via_runtime_check(self) -> Result<(&'a T, &'a [T::Elem]), Self> {
match self.into_ptr().via_runtime_check() {
Ok((l, r)) => Ok((l.as_ref(), r.as_ref())),
Err(s) => Err(s.into_ref()),
}
}
#[must_use = "has no side effects"]
#[inline(always)]
pub unsafe fn via_unchecked(self) -> (&'a T, &'a [T::Elem]) {
let (l, r) = unsafe { self.into_ptr().via_unchecked() };
(l.as_ref(), r.as_ref())
}
}
impl<'a, T> Split<&'a mut T>
where
T: ?Sized + SplitAt,
{
#[inline(always)]
fn into_ptr(self) -> Split<Ptr<'a, T, (Exclusive, Aligned, Valid)>> {
let source = Ptr::from_mut(self.source);
unsafe { Split::new(source, self.l_len) }
}
#[must_use = "has no side effects"]
#[inline(always)]
pub fn via_into_bytes(self) -> (&'a mut T, &'a mut [T::Elem])
where
T: IntoBytes,
{
let (l, r) = self.into_ptr().via_into_bytes();
(l.as_mut(), r.as_mut())
}
#[must_use = "has no side effects"]
#[inline(always)]
pub fn via_unaligned(self) -> (&'a mut T, &'a mut [T::Elem])
where
T: Unaligned,
{
let (l, r) = self.into_ptr().via_unaligned();
(l.as_mut(), r.as_mut())
}
#[must_use = "has no side effects"]
#[inline(always)]
pub fn via_runtime_check(self) -> Result<(&'a mut T, &'a mut [T::Elem]), Self> {
match self.into_ptr().via_runtime_check() {
Ok((l, r)) => Ok((l.as_mut(), r.as_mut())),
Err(s) => Err(s.into_mut()),
}
}
#[must_use = "has no side effects"]
#[inline(always)]
pub unsafe fn via_unchecked(self) -> (&'a mut T, &'a mut [T::Elem]) {
let (l, r) = unsafe { self.into_ptr().via_unchecked() };
(l.as_mut(), r.as_mut())
}
}
impl<'a, T, I> Split<Ptr<'a, T, I>>
where
T: ?Sized + SplitAt,
I: Invariants<Alignment = Aligned, Validity = Valid>,
{
fn into_ref(self) -> Split<&'a T>
where
I: Invariants<Aliasing = Shared>,
{
unsafe { Split::new(self.source.as_ref(), self.l_len) }
}
fn into_mut(self) -> Split<&'a mut T>
where
I: Invariants<Aliasing = Exclusive>,
{
unsafe { Split::new(self.source.unify_invariants().as_mut(), self.l_len) }
}
#[inline(always)]
fn l_len(&self) -> MetadataOf<T> {
unsafe { MetadataOf::<T>::new_unchecked(self.l_len) }
}
#[inline(always)]
fn via_immutable(self) -> (Ptr<'a, T, I>, Ptr<'a, [T::Elem], I>)
where
T: Immutable,
I: Invariants<Aliasing = Shared>,
{
unsafe { self.via_unchecked() }
}
#[inline(always)]
fn via_into_bytes(self) -> (Ptr<'a, T, I>, Ptr<'a, [T::Elem], I>)
where
T: IntoBytes,
{
unsafe { self.via_unchecked() }
}
#[inline(always)]
fn via_unaligned(self) -> (Ptr<'a, T, I>, Ptr<'a, [T::Elem], I>)
where
T: Unaligned,
{
unsafe { self.via_unchecked() }
}
#[inline(always)]
fn via_runtime_check(self) -> Result<(Ptr<'a, T, I>, Ptr<'a, [T::Elem], I>), Self> {
let l_len = self.l_len();
if l_len.padding_needed_for() == 0 {
Ok(unsafe { self.via_unchecked() })
} else {
Err(self)
}
}
#[inline(always)]
unsafe fn via_unchecked(self) -> (Ptr<'a, T, I>, Ptr<'a, [T::Elem], I>) {
let l_len = self.l_len();
let inner = self.source.as_inner();
let (left, right) = unsafe { inner.split_at_unchecked(l_len) };
let left = unsafe { Ptr::from_inner(left) };
let right = unsafe { Ptr::from_inner(right) };
(left, right)
}
}
#[cfg(test)]
mod tests {
#[cfg(feature = "derive")]
#[test]
fn test_split_at() {
use crate::{FromBytes, Immutable, IntoBytes, KnownLayout, SplitAt};
#[derive(FromBytes, KnownLayout, SplitAt, IntoBytes, Immutable, Debug)]
#[repr(C)]
struct SliceDst<const OFFSET: usize> {
prefix: [u8; OFFSET],
trailing: [u8],
}
#[allow(clippy::as_conversions)]
fn test_split_at<const OFFSET: usize, const BUFFER_SIZE: usize>() {
let n: usize = BUFFER_SIZE - OFFSET;
let arr = [1; BUFFER_SIZE];
let dst = SliceDst::<OFFSET>::ref_from_bytes(&arr[..]).unwrap();
for i in 0..=n {
let (l, r) = dst.split_at(i).unwrap().via_runtime_check().unwrap();
let l_sum: u8 = l.trailing.iter().sum();
let r_sum: u8 = r.iter().sum();
assert_eq!(l_sum, i as u8);
assert_eq!(r_sum, (n - i) as u8);
assert_eq!(l_sum + r_sum, n as u8);
}
let n: usize = BUFFER_SIZE - OFFSET;
let mut arr = [1; BUFFER_SIZE];
let dst = SliceDst::<OFFSET>::mut_from_bytes(&mut arr[..]).unwrap();
for i in 0..=n {
let (l, r) = dst.split_at_mut(i).unwrap().via_runtime_check().unwrap();
let l_sum: u8 = l.trailing.iter().sum();
let r_sum: u8 = r.iter().sum();
assert_eq!(l_sum, i as u8);
assert_eq!(r_sum, (n - i) as u8);
assert_eq!(l_sum + r_sum, n as u8);
}
}
test_split_at::<0, 16>();
test_split_at::<1, 17>();
test_split_at::<2, 18>();
}
#[cfg(feature = "derive")]
#[test]
#[allow(clippy::as_conversions)]
fn test_split_at_overlapping() {
use crate::{FromBytes, Immutable, IntoBytes, KnownLayout, SplitAt};
#[derive(FromBytes, KnownLayout, SplitAt, Immutable)]
#[repr(C, align(2))]
struct SliceDst {
prefix: u8,
trailing: [u8],
}
const N: usize = 16;
let arr = [1u16; N];
let dst = SliceDst::ref_from_bytes(arr.as_bytes()).unwrap();
for i in 0..N {
let split = dst.split_at(i).unwrap().via_runtime_check();
if i % 2 == 1 {
assert!(split.is_ok());
} else {
assert!(split.is_err());
}
}
}
}