Skip to content

Commit

Permalink
Replace get_split with split_first_chunk builtin
Browse files Browse the repository at this point in the history
[`split_first_chunk`][0] is a nightly only builtin. It looks like it
will be [stabillized soon][1], so I decided to use it and copy over the
implementation in the meantime.

Amusingly, benchmark show a consistent 3% improvement to throughput
compared to using `get_split`.

[0]: https://doc.rust-lang.org/std/primitive.slice.html#method.split_first_chunk
[1]: rust-lang/rust#117561
  • Loading branch information
nickbabcock committed Nov 24, 2023
1 parent 65bd6f1 commit 63cccd9
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 17 deletions.
34 changes: 17 additions & 17 deletions src/binary/de.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use super::{tokens::*, Rgb};
use crate::{
binary::{BinaryFlavor, FailedResolveStrategy, TokenResolver},
de::ColorSequence,
util::get_split,
util::split_first_chunk,
BinaryTape, BinaryToken, DeserializeError, DeserializeErrorKind, Error, ErrorKind,
};
use serde::de::{self, Deserialize, DeserializeSeed, MapAccess, SeqAccess, Visitor};
Expand All @@ -24,8 +24,8 @@ impl<'data> OndemandParser<'data> {

#[inline]
pub fn next(&mut self) -> Option<u16> {
let (data, token) =
get_split::<2>(self.data).map(|(head, rest)| (rest, u16::from_le_bytes(head)))?;
let (data, token) = split_first_chunk::<2>(self.data)
.map(|(head, rest)| (rest, u16::from_le_bytes(*head)))?;
self.data = data;
Some(token)
}
Expand All @@ -37,8 +37,8 @@ impl<'data> OndemandParser<'data> {

#[inline]
pub fn read_string(&mut self) -> Result<&'data [u8], Error> {
let (head, rest) = get_split::<2>(self.data).ok_or_else(Error::eof)?;
let text_len = usize::from(u16::from_le_bytes(head));
let (head, rest) = split_first_chunk::<2>(self.data).ok_or_else(Error::eof)?;
let text_len = usize::from(u16::from_le_bytes(*head));
if text_len <= rest.len() {
let (text, rest) = rest.split_at(text_len);
self.data = rest;
Expand All @@ -57,44 +57,44 @@ impl<'data> OndemandParser<'data> {

#[inline]
fn read_u32(&mut self) -> Result<u32, Error> {
let (head, rest) = get_split::<4>(self.data).ok_or_else(Error::eof)?;
let (head, rest) = split_first_chunk::<4>(self.data).ok_or_else(Error::eof)?;
self.data = rest;
Ok(u32::from_le_bytes(head))
Ok(u32::from_le_bytes(*head))
}

#[inline]
fn read_u64(&mut self) -> Result<u64, Error> {
let (head, rest) = get_split::<8>(self.data).ok_or_else(Error::eof)?;
let (head, rest) = split_first_chunk::<8>(self.data).ok_or_else(Error::eof)?;
self.data = rest;
Ok(u64::from_le_bytes(head))
Ok(u64::from_le_bytes(*head))
}

#[inline]
fn read_i64(&mut self) -> Result<i64, Error> {
let (head, rest) = get_split::<8>(self.data).ok_or_else(Error::eof)?;
let (head, rest) = split_first_chunk::<8>(self.data).ok_or_else(Error::eof)?;
self.data = rest;
Ok(i64::from_le_bytes(head))
Ok(i64::from_le_bytes(*head))
}

#[inline]
fn read_i32(&mut self) -> Result<i32, Error> {
let (head, rest) = get_split::<4>(self.data).ok_or_else(Error::eof)?;
let (head, rest) = split_first_chunk::<4>(self.data).ok_or_else(Error::eof)?;
self.data = rest;
Ok(i32::from_le_bytes(head))
Ok(i32::from_le_bytes(*head))
}

#[inline]
fn read_f32(&mut self) -> Result<[u8; 4], Error> {
let (head, rest) = get_split::<4>(self.data).ok_or_else(Error::eof)?;
let (head, rest) = split_first_chunk::<4>(self.data).ok_or_else(Error::eof)?;
self.data = rest;
Ok(head)
Ok(*head)
}

#[inline]
fn read_f64(&mut self) -> Result<[u8; 8], Error> {
let (head, rest) = get_split::<8>(self.data).ok_or_else(Error::eof)?;
let (head, rest) = split_first_chunk::<8>(self.data).ok_or_else(Error::eof)?;
self.data = rest;
Ok(head)
Ok(*head)
}

#[inline]
Expand Down
26 changes: 26 additions & 0 deletions src/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,32 @@ fn take<const N: usize>(data: &[u8]) -> [u8; N] {
unsafe { *(data.as_ptr() as *const [u8; N]) }
}

#[inline]
pub const fn split_first_chunk<const N: usize>(data: &[u8]) -> Option<(&[u8; N], &[u8])> {
#[inline]
const unsafe fn split_at_unchecked(data: &[u8], mid: usize) -> (&[u8], &[u8]) {
let len = data.len();
let ptr = data.as_ptr();
unsafe {
(
core::slice::from_raw_parts(ptr, mid),
core::slice::from_raw_parts(ptr.add(mid), len - mid),
)
}
}

if data.len() < N {
None
} else {
// SAFETY: We manually verified the bounds of the split.
let (first, tail) = unsafe { split_at_unchecked(data, N) };

// SAFETY: We explicitly check for the correct number of elements,
// and do not let the references outlive the slice.
Some((unsafe { &*(first.as_ptr().cast::<[u8; N]>()) }, tail))
}
}

#[inline]
pub(crate) fn get_split<const N: usize>(data: &[u8]) -> Option<([u8; N], &[u8])> {
data.get(N..).map(|d| (take::<N>(data), d))
Expand Down

0 comments on commit 63cccd9

Please sign in to comment.