Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
polazarus committed Jun 6, 2024
1 parent 999696c commit 573150c
Show file tree
Hide file tree
Showing 5 changed files with 312 additions and 56 deletions.
78 changes: 78 additions & 0 deletions src/bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,9 @@
use core::borrow::Borrow;
use core::hash::Hash;
use core::mem::MaybeUninit;
use core::ops::{Bound, Deref, DerefMut, Range, RangeBounds};
use core::ptr;

use super::raw::Raw;
use crate::alloc::fmt;
Expand Down Expand Up @@ -88,6 +90,10 @@ where
Self(Raw::empty())
}

pub(crate) fn with_capacity(new_len: usize) -> Self {
Self(Raw::with_capacity(new_len))
}

/// Creates a new `HipByt` from a borrowed slice without copying the slice.
///
/// # Examples
Expand Down Expand Up @@ -677,6 +683,78 @@ where
pub fn repeat(&self, n: usize) -> Self {
Self(self.0.repeat(n))
}

/// Returns the remaining spare capacity of the vector as a slice of
/// `MaybeUninit<T>`.
///
/// The returned slice can be used to fill the vector with data (e.g. by
/// reading from a file) before marking the data as initialized using the
/// [`set_len`] method.
///
/// [`set_len`]: Raw::set_len
pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<u8>] {
self.0.spare_capacity_mut()
}

/// Forces the length of the vector to `new_len`.
///
/// # Safety
///
/// * `new_len` should be must be less than or equal to `INLINE_CAPACITY`.
/// * The elements at `old_len..new_len` must be initialized.
/// * The vector should not be shared.
pub unsafe fn set_len(&mut self, new_len: usize) {
// SAFETY: precondition
unsafe { self.0.set_len(new_len) }
}

pub fn concat_slices(slices: &[&[u8]]) -> Self {
let new_len = slices.iter().map(|e| e.len()).sum();

let mut raw = Raw::with_capacity(new_len);
let dst = raw.spare_capacity_mut();
let dst_ptr = dst.as_mut_ptr().cast();
let final_ptr = slices.iter().fold(dst_ptr, |dst, slice| {
let len = slice.len();
unsafe {
ptr::copy_nonoverlapping(slice.as_ptr(), dst, len);
dst.add(len)
}
});

unsafe { raw.set_len(new_len) };

// check end pointer
debug_assert_eq!(final_ptr.cast_const(), raw.as_slice().as_ptr_range().end);
Self(raw)
}

pub fn concat(slices: impl IntoIterator<Item = impl Borrow<[u8]>> + Clone) -> Self {
let new_len = slices.clone().into_iter().map(|e| e.borrow().len()).sum();

let mut raw = Raw::with_capacity(new_len);
let dst = raw.spare_capacity_mut();
let dst_ptr: *mut u8 = dst.as_mut_ptr().cast();

// compute the final pointer
let final_ptr = unsafe { dst_ptr.add(new_len) };

let _ = slices.into_iter().fold(dst_ptr, |dst, slice| {
let slice = slice.borrow();
let len = slice.len();
let end_ptr = unsafe { dst_ptr.add(len) };
assert!(end_ptr <= final_ptr, "slices changed during concat");
unsafe {
ptr::copy_nonoverlapping(slice.as_ptr(), dst, len);
end_ptr
}
});

unsafe { raw.set_len(new_len) };
debug_assert_eq!(final_ptr.cast_const(), raw.as_slice().as_ptr_range().end);

Self(raw)
}
}

impl<B> HipByt<'static, B>
Expand Down
180 changes: 125 additions & 55 deletions src/raw.rs
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,16 @@ enum RawSplit<'a, 'borrow, B: Backend> {
Borrowed(&'a Borrowed<'borrow>),
}

/// Helper enum to split this raw byte string into its possible representation mutably.
enum RawSplitMut<'a, 'borrow, B: Backend> {
/// Inline representation
Inline(&'a mut Inline),
/// Allocated and shared representation
Allocated(&'a mut Allocated<B>),
/// Borrowed slice representation
Borrowed(&'a mut Borrowed<'borrow>),
}

impl<'borrow, B: Backend> Raw<'borrow, B> {
/// Retrieves a reference on the union.
#[inline]
Expand Down Expand Up @@ -166,6 +176,59 @@ impl<'borrow, B: Backend> Raw<'borrow, B> {
Union { borrowed }.into_raw()
}

/// Retrieves the tag.
const fn tag(&self) -> Tag {
match self.tag_byte.get() & MASK {
TAG_INLINE => Tag::Inline,
TAG_BORROWED => Tag::Borrowed,
TAG_ALLOCATED => Tag::Allocated,
// SAFETY: type invariant
_ => unsafe { unreachable_unchecked() },
}
}

/// Splits this raw into its possible representation.
#[inline]
const fn split(&self) -> RawSplit<'_, 'borrow, B> {
let tag = self.tag();
let union = self.union();
match tag {
Tag::Inline => {
// SAFETY: representation checked
RawSplit::Inline(unsafe { &union.inline })
}
Tag::Borrowed => {
// SAFETY: representation checked
RawSplit::Borrowed(unsafe { &union.borrowed })
}
Tag::Allocated => {
// SAFETY: representation checked
RawSplit::Allocated(unsafe { &union.allocated })
}
}
}

/// Splits this raw into its possible representation.
#[inline]
fn split_mut(&mut self) -> RawSplitMut<'_, 'borrow, B> {
let tag = self.tag();
let union = self.union_mut();
match tag {
Tag::Inline => {
// SAFETY: representation checked
RawSplitMut::Inline(unsafe { &mut union.inline })
}
Tag::Borrowed => {
// SAFETY: representation checked
RawSplitMut::Borrowed(unsafe { &mut union.borrowed })
}
Tag::Allocated => {
// SAFETY: representation checked
RawSplitMut::Allocated(unsafe { &mut union.allocated })
}
}
}

/// Creates a new `Raw` from a vector.
///
/// The vector's length should be strictly greater than `INLINE_CAPACITY`.
Expand Down Expand Up @@ -196,12 +259,6 @@ impl<'borrow, B: Backend> Raw<'borrow, B> {
Self::from_inline(inline)
}

/// Creates a new empty `Raw`.
#[inline]
pub const fn empty() -> Self {
Self::from_inline(Inline::empty())
}

/// Creates a new `Raw` from a static slice.
///
/// # Representation
Expand All @@ -218,6 +275,12 @@ impl<'borrow, B: Backend> Raw<'borrow, B> {

// derived constructors

/// Creates a new empty `Raw`.
#[inline]
pub const fn empty() -> Self {
Self::from_inline(Inline::empty())
}

/// Creates a new `Raw` from a vector.
///
/// Will normalize the representation depending on the size of the vector.
Expand Down Expand Up @@ -245,24 +308,14 @@ impl<'borrow, B: Backend> Raw<'borrow, B> {
}
}

/// Splits this raw into its possible representation.
#[inline]
const fn split(&self) -> RawSplit<'_, 'borrow, B> {
let tag = self.tag();
let union = self.union();
match tag {
Tag::Inline => {
// SAFETY: representation checked
RawSplit::Inline(unsafe { &union.inline })
}
Tag::Borrowed => {
// SAFETY: representation checked
RawSplit::Borrowed(unsafe { &union.borrowed })
}
Tag::Allocated => {
// SAFETY: representation checked
RawSplit::Allocated(unsafe { &union.allocated })
}
/// Creates a new `Raw` with the given capacity.
///
/// **This representation may not be normalized.**
pub fn with_capacity(capacity: usize) -> Self {
if capacity <= INLINE_CAPACITY {
Self::from_inline(Inline::empty())
} else {
Self::from_vec(Vec::with_capacity(capacity))
}
}

Expand Down Expand Up @@ -417,14 +470,10 @@ impl<'borrow, B: Backend> Raw<'borrow, B> {

#[inline]
pub fn as_mut_slice(&mut self) -> Option<&mut [u8]> {
let tag = self.tag();
let union = self.union_mut();

// SAFETY: representation is checked
match tag {
Tag::Inline => Some(unsafe { union.inline.as_mut_slice() }),
Tag::Borrowed => None,
Tag::Allocated => unsafe { union.allocated.as_mut_slice() },
match self.split_mut() {
RawSplitMut::Inline(inline) => Some(inline.as_mut_slice()),
RawSplitMut::Allocated(allocated) => allocated.as_mut_slice(),
RawSplitMut::Borrowed(_) => None,
}
}

Expand All @@ -433,22 +482,6 @@ impl<'borrow, B: Backend> Raw<'borrow, B> {
pub fn push_slice(&mut self, addition: &[u8]) {
let new_len = self.len() + addition.len();

if new_len <= INLINE_CAPACITY {
// can be inlined

if !self.is_inline() {
// make it inline first
// SAFETY: `new_len` is checked before, so current len <= INLINE_CAPACITY
*self = unsafe { Self::inline_unchecked(self.as_slice()) };
}

// SAFETY: `new_len` is checked above
unsafe {
self.union_mut().inline.push_slice_unchecked(addition);
}
return;
}

if self.is_allocated() {
// current allocation may be pushed into it directly?

Expand All @@ -464,6 +497,20 @@ impl<'borrow, B: Backend> Raw<'borrow, B> {
}
}

if new_len <= INLINE_CAPACITY {
if !self.is_inline() {
// make it inline first
// SAFETY: `new_len` is checked before, so current len <= INLINE_CAPACITY
*self = unsafe { Self::inline_unchecked(self.as_slice()) };
}

// SAFETY: `new_len` is checked above
unsafe {
self.union_mut().inline.push_slice_unchecked(addition);
}
return;
}

// requires a new vector
let mut vec = Vec::with_capacity(new_len);
vec.extend_from_slice(self.as_slice());
Expand Down Expand Up @@ -685,13 +732,36 @@ impl<'borrow, B: Backend> Raw<'borrow, B> {
}
}

const fn tag(&self) -> Tag {
match self.tag_byte.get() & MASK {
TAG_INLINE => Tag::Inline,
TAG_BORROWED => Tag::Borrowed,
TAG_ALLOCATED => Tag::Allocated,
// SAFETY: type invariant
_ => unsafe { unreachable_unchecked() },
/// Returns the remaining spare capacity of the vector as a slice of
/// `MaybeUninit<T>`.
///
/// The returned slice can be used to fill the vector with data (e.g. by
/// reading from a file) before marking the data as initialized using the
/// [`set_len`] method.
///
/// [`set_len`]: Raw::set_len
pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<u8>] {
match self.split_mut() {
RawSplitMut::Borrowed(_) => &mut [],
RawSplitMut::Inline(inline) => inline.spare_capacity_mut(),
RawSplitMut::Allocated(allocated) => allocated.spare_capacity_mut(),
}
}

/// Forces the length of the vector to `new_len`.
///
/// # Safety
///
/// * `new_len` should be must be less than or equal to `INLINE_CAPACITY`.
/// * The elements at `old_len..new_len` must be initialized.
/// * The vector should not be shared (if `new_len != old_len`).
pub unsafe fn set_len(&mut self, new_len: usize) {
match self.split_mut() {
RawSplitMut::Borrowed(borrowed) => {
debug_assert!(borrowed.len() == new_len, "set_len on borrowed")
}
RawSplitMut::Inline(inline) => unsafe { inline.set_len(new_len) },
RawSplitMut::Allocated(allocated) => unsafe { allocated.set_len(new_len) },
}
}
}
Expand Down
Loading

0 comments on commit 573150c

Please sign in to comment.