Skip to content

Commit

Permalink
impl_escape_unchecked() byte exact read
Browse files Browse the repository at this point in the history
  • Loading branch information
@ijl
ijl committed Feb 23, 2024
1 parent e04ea73 commit 5067ead
Show file tree
Hide file tree
Showing 5 changed files with 44 additions and 83 deletions.
33 changes: 0 additions & 33 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 0 additions & 1 deletion Cargo.toml
View file
Expand Up @@ -56,7 +56,6 @@ encoding_rs = { version = "0.8", default_features = false }
itoa = { version = "1", default_features = false }
itoap = { version = "1", features = ["std", "simd"] }
once_cell = { version = "1", default_features = false, features = ["race"] }
page_size = { version = "0.6" }
pyo3-ffi = { version = "^0.20.2", default_features = false, features = ["extension-module"]}
ryu = { version = "1", default_features = false }
serde = { version = "1", default_features = false }
Expand Down
2 changes: 1 addition & 1 deletion src/serialize/writer/json.rs
View file
Expand Up @@ -572,7 +572,7 @@ where
W: ?Sized + io::Write + WriteExt,
{
unsafe {
let num_reserved_bytes = value.len() * 8 + 32 + 3;
let num_reserved_bytes = value.len() * 8 + 32;
writer.reserve(num_reserved_bytes);

let written = crate::serialize::writer::simd::format_escaped_str_impl_128(
Expand Down
84 changes: 43 additions & 41 deletions src/serialize/writer/simd.rs
View file
Expand Up @@ -2,19 +2,25 @@
// Copyright 2023-2024 liuq19, ijl
// adapted from sonic-rs' src/util/string.rs

use crate::typeref::PAGE_SIZE;
use core::simd::cmp::{SimdPartialEq, SimdPartialOrd};

macro_rules! impl_escape_unchecked {
($src:expr, $dst:expr, $nb:expr, $omask:expr, $cn:expr) => {
($src:expr, $dst:expr, $nb:expr, $omask:expr, $cn:expr, $v:expr, $rotate:expr) => {
$nb -= $cn;
if $rotate == true {
for _ in 0..$cn {
$v = $v.rotate_elements_left::<1>();
}
}
$dst = $dst.add($cn);
$src = $src.add($cn);
$omask >>= $cn;
loop {
$nb -= 1;
if $rotate == true {
$v = $v.rotate_elements_left::<1>();
}
$omask = $omask >> 1;

if *($src) == b'"' {
core::ptr::copy_nonoverlapping(b"\\\"".as_ptr(), $dst, 2);
$dst = $dst.add(2);
Expand Down Expand Up @@ -48,46 +54,48 @@ macro_rules! impl_format_simd {
*dptr = b'"';
dptr = dptr.add(1);

while nb >= STRIDE {
let v = StrVector::from_slice(core::slice::from_raw_parts(sptr, STRIDE));
v.copy_to_slice(core::slice::from_raw_parts_mut(dptr, STRIDE));
let mut mask =
(v.simd_eq(blash) | v.simd_eq(quote) | v.simd_lt(x20)).to_bitmask() as u32;

if likely!(mask == 0) {
nb -= STRIDE;
dptr = dptr.add(STRIDE);
sptr = sptr.add(STRIDE);
} else {
let cn = mask.trailing_zeros() as usize;
impl_escape_unchecked!(sptr, dptr, nb, mask, cn);
{
const ROTATE: bool = false;
while nb >= STRIDE {
let mut v = StrVector::from_slice(core::slice::from_raw_parts(sptr, STRIDE));
let mut mask =
(v.simd_eq(blash) | v.simd_eq(quote) | v.simd_lt(x20)).to_bitmask() as u32;
v.copy_to_slice(core::slice::from_raw_parts_mut(dptr, STRIDE));

if likely!(mask == 0) {
nb -= STRIDE;
dptr = dptr.add(STRIDE);
sptr = sptr.add(STRIDE);
} else {
let cn = mask.trailing_zeros() as usize;
impl_escape_unchecked!(sptr, dptr, nb, mask, cn, v, ROTATE);
}
}
}

let mut v = if unlikely!(is_cross_page!(sptr)) {
{
const ROTATE: bool = true;
let mut v = StrVector::default();
v.as_mut_array()[..nb].copy_from_slice(core::slice::from_raw_parts(sptr, nb));
v
} else {
StrVector::from_slice(core::slice::from_raw_parts(sptr, STRIDE))
};
while nb > 0 {
v.copy_to_slice(core::slice::from_raw_parts_mut(dptr, STRIDE));
{
let vec_ptr = v.as_mut_array().as_mut_ptr();
for idx in 0..nb {
core::ptr::write(vec_ptr.add(idx), *sptr.add(idx));
}
}

let mut mask = (v.simd_eq(blash) | v.simd_eq(quote) | v.simd_lt(x20)).to_bitmask()
as u32
& (STRIDE_SATURATION >> (32 - STRIDE - nb));

if likely!(mask == 0) {
dptr = dptr.add(nb);
break;
} else {
let cn = mask.trailing_zeros() as usize;
let nb_start = nb;
impl_escape_unchecked!(sptr, dptr, nb, mask, cn);
let mut consumed = nb_start - nb;
while consumed != 0 {
v = v.rotate_elements_left::<1>();
consumed -= 1;
while nb > 0 {
v.copy_to_slice(core::slice::from_raw_parts_mut(dptr, STRIDE));

if likely!(mask == 0) {
dptr = dptr.add(nb);
break;
} else {
let cn = mask.trailing_zeros() as usize;
impl_escape_unchecked!(sptr, dptr, nb, mask, cn, v, ROTATE);
}
}
}
Expand All @@ -100,12 +108,6 @@ macro_rules! impl_format_simd {
};
}

macro_rules! is_cross_page {
($src:expr) => {
unsafe { (($src as usize & (PAGE_SIZE - 1)) + STRIDE) > PAGE_SIZE }
};
}

#[cold]
#[inline(never)]
fn write_unusual_escape(sptr: *const u8, dptr: *mut u8) -> *mut u8 {
Expand Down
7 changes: 0 additions & 7 deletions src/typeref.rs
View file
Expand Up @@ -75,9 +75,6 @@ pub static mut DESCR_STR: *mut PyObject = null_mut();
pub static mut VALUE_STR: *mut PyObject = null_mut();
pub static mut INT_ATTR_STR: *mut PyObject = null_mut();

#[cfg(feature = "unstable-simd")]
pub static mut PAGE_SIZE: usize = 0;

#[cfg(feature = "yyjson")]
pub const YYJSON_BUFFER_SIZE: usize = 1024 * 1024 * 8;

Expand Down Expand Up @@ -139,10 +136,6 @@ fn _init_typerefs_impl() -> bool {
unsafe {
debug_assert!(crate::opt::MAX_OPT < u16::MAX as i32);

#[cfg(feature = "unstable-simd")]
{
PAGE_SIZE = page_size::get();
}
assert!(crate::deserialize::KEY_MAP
.set(crate::deserialize::KeyMap::default())
.is_ok());
Expand Down

0 comments on commit 5067ead

Please sign in to comment.