Escape str using SIMD

ci/azure-macos.yml

@@ -23,7 +23,7 @@ steps:
     PATH=$HOME/.cargo/bin:$PATH \
     MACOSX_DEPLOYMENT_TARGET=$(macosx_deployment_target) \
     PYO3_CROSS_LIB_DIR=$(python -c "import sysconfig;print(sysconfig.get_config_var('LIBDIR'))") \
-    maturin build --release --strip --features=encoding_rs/simd-accel,no-panic,yyjson --interpreter $(interpreter) --target=universal2-apple-darwin
+    maturin build --release --strip --features=encoding_rs/simd-accel,no-panic,unstable-simd,yyjson --interpreter $(interpreter) --target=universal2-apple-darwin
   env:
     CC: "clang"
     CFLAGS: "-O2 -fstrict-aliasing -flto=full"

src/lib.rs

@@ -4,6 +4,7 @@
 #![cfg_attr(feature = "optimize", feature(optimize_attribute))]
 #![cfg_attr(feature = "strict_provenance", feature(strict_provenance))]
 #![cfg_attr(feature = "strict_provenance", warn(fuzzy_provenance_casts))]
+#![cfg_attr(feature = "unstable-simd", feature(portable_simd))]
 #![allow(unknown_lints)] // internal_features
 #![allow(internal_features)] // core_intrinsics
 #![allow(unused_unsafe)]

src/serialize/writer/json.rs

@@ -565,6 +565,85 @@ where
     }
 }
 
+#[cfg(all(
+    feature = "unstable-simd",
+    target_arch = "x86_64",
+    target_feature = "avx2"
+))]
+#[inline(always)]
+fn format_escaped_str<W>(writer: &mut W, value: &str) -> io::Result<()>
+where
+    W: ?Sized + io::Write + WriteExt,
+{
+    unsafe {
+        let num_reserved_bytes = value.len() * 8 + 32 + 3;
+        writer.reserve(num_reserved_bytes);
+
+        let written = crate::serialize::writer::simd::format_escaped_str_impl_256(
+            writer.as_mut_buffer_ptr(),
+            value.as_bytes().as_ptr(),
+            value.len(),
+        );
+        writer.set_written(written);
+    }
+    Ok(())
+}
+
+#[cfg(all(
+    feature = "unstable-simd",
+    target_arch = "x86_64",
+    not(target_feature = "avx2")
+))]
+#[inline(always)]
+fn format_escaped_str<W>(writer: &mut W, value: &str) -> io::Result<()>
+where
+    W: ?Sized + io::Write + WriteExt,
+{
+    unsafe {
+        let num_reserved_bytes = value.len() * 8 + 32 + 3;
+        writer.reserve(num_reserved_bytes);
+
+        if std::is_x86_feature_detected!("avx2") {
+            let written = crate::serialize::writer::simd::format_escaped_str_impl_256(
+                writer.as_mut_buffer_ptr(),
+                value.as_bytes().as_ptr(),
+                value.len(),
+            );
+            writer.set_written(written);
+        } else {
+            let written = crate::serialize::writer::simd::format_escaped_str_impl_128(
+                writer.as_mut_buffer_ptr(),
+                value.as_bytes().as_ptr(),
+                value.len(),
+            );
+            writer.set_written(written);
+        }
+    }
+    Ok(())
+}
+
+#[cfg(all(feature = "unstable-simd", not(target_arch = "x86_64")))]
+#[inline(always)]
+fn format_escaped_str<W>(writer: &mut W, value: &str) -> io::Result<()>
+where
+    W: ?Sized + io::Write + WriteExt,
+{
+    unsafe {
+        let num_reserved_bytes = value.len() * 8 + 32 + 3;
+        writer.reserve(num_reserved_bytes);
+
+        let written = crate::serialize::writer::simd::format_escaped_str_impl_128(
+            writer.as_mut_buffer_ptr(),
+            value.as_bytes().as_ptr(),
+            value.len(),
+        );
+
+        writer.set_written(written);
+    }
+    Ok(())
+}
+
+#[cfg(not(feature = "unstable-simd"))]
 #[inline(always)]
 fn format_escaped_str<W>(writer: &mut W, value: &str) -> io::Result<()>
 where

src/serialize/writer/mod.rs

@@ -1,9 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
 mod byteswriter;
+#[cfg(not(feature = "unstable-simd"))]
 mod escape;
 mod formatter;
 mod json;
+#[cfg(feature = "unstable-simd")]
+mod simd;
 
 pub use byteswriter::{BytesWriter, WriteExt};
 pub use json::{to_writer, to_writer_pretty};

src/serialize/writer/simd.rs

@@ -0,0 +1,149 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright 2023-2024 liuq19, ijl
+// adapted from sonic-rs' src/util/string.rs
+
+use std::simd::cmp::{SimdPartialEq, SimdPartialOrd};
+
+macro_rules! impl_escape_unchecked {
+    ($src:expr, $dst:expr, $nb:expr, $omask:expr, $cn:expr) => {
+        $nb -= $cn;
+        $dst = $dst.add($cn);
+        $src = $src.add($cn);
+        let mut mask = $omask << $cn;
+        loop {
+            $nb -= 1;
+            mask = mask << 1;
+            let replacement = if *($src) == b'"' {
+                (*b"\\\"\0\0\0\0\0\0", 2)
+            } else if *($src) == b'\\' {
+                (*b"\\\\\0\0\0\0\0\0", 2)
+            } else {
+                match *($src) {
+                    0 => (*b"\\u0000\0\0", 6),
+                    1 => (*b"\\u0001\0\0", 6),
+                    2 => (*b"\\u0002\0\0", 6),
+                    3 => (*b"\\u0003\0\0", 6),
+                    4 => (*b"\\u0004\0\0", 6),
+                    5 => (*b"\\u0005\0\0", 6),
+                    6 => (*b"\\u0006\0\0", 6),
+                    7 => (*b"\\u0007\0\0", 6),
+                    8 => (*b"\\b\0\0\0\0\0\0", 2),
+                    9 => (*b"\\t\0\0\0\0\0\0", 2),
+                    10 => (*b"\\n\0\0\0\0\0\0", 2),
+                    11 => (*b"\\u000b\0\0", 6),
+                    12 => (*b"\\f\0\0\0\0\0\0", 2),
+                    13 => (*b"\\r\0\0\0\0\0\0", 2),
+                    14 => (*b"\\u000e\0\0", 6),
+                    15 => (*b"\\u000f\0\0", 6),
+                    16 => (*b"\\u0010\0\0", 6),
+                    17 => (*b"\\u0011\0\0", 6),
+                    18 => (*b"\\u0012\0\0", 6),
+                    19 => (*b"\\u0013\0\0", 6),
+                    20 => (*b"\\u0014\0\0", 6),
+                    21 => (*b"\\u0015\0\0", 6),
+                    22 => (*b"\\u0016\0\0", 6),
+                    23 => (*b"\\u0017\0\0", 6),
+                    24 => (*b"\\u0018\0\0", 6),
+                    25 => (*b"\\u0019\0\0", 6),
+                    26 => (*b"\\u001a\0\0", 6),
+                    27 => (*b"\\u001b\0\0", 6),
+                    28 => (*b"\\u001c\0\0", 6),
+                    29 => (*b"\\u001d\0\0", 6),
+                    30 => (*b"\\u001e\0\0", 6),
+                    31 => (*b"\\u001f\0\0", 6),
+                    _ => unreachable!(),
+                }
+            };
+            std::ptr::copy_nonoverlapping(replacement.0.as_ptr(), $dst, 8);
+            $dst = $dst.add(replacement.1 as usize);
+            $src = $src.add(1);
+            if likely!(mask & (1 << (STRIDE - 1)) != 1) {
+                break;
+            }
+        }
+    };
+}
+macro_rules! impl_format_simd {
+    ($odptr:expr, $value_ptr:expr, $value_len:expr) => {
+        let mut dptr = $odptr;
+        let dstart = $odptr;
+        let mut sptr = $value_ptr;
+        let mut nb: usize = $value_len;
+
+        let blash = StrVector::from_array([b'\\'; STRIDE]);
+        let quote = StrVector::from_array([b'"'; STRIDE]);
+        let x20 = StrVector::from_array([32; STRIDE]);
+
+        unsafe {
+            *dptr = b'"';
+            dptr = dptr.add(1);
+
+            while nb >= STRIDE {
+                let v = StrVector::from_slice(std::slice::from_raw_parts(sptr, STRIDE));
+                v.copy_to_slice(std::slice::from_raw_parts_mut(dptr, STRIDE));
+                let 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);
+                }
+            }
+
+            while nb > 0 {
+                let v = StrVector::from_slice(std::slice::from_raw_parts(sptr, STRIDE));
+                v.copy_to_slice(std::slice::from_raw_parts_mut(dptr, STRIDE));
+                let mask = (v.simd_eq(blash) | v.simd_eq(quote) | v.simd_lt(x20)).to_bitmask()
+                    as u32
+                    & (STRIDE_SATURATION >> (STRIDE - nb));
+                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);
+                }
+            }
+
+            *dptr = b'"';
+            dptr = dptr.add(1);
+        }
+
+        return dptr as usize - dstart as usize;
+    };
+}
+
+#[inline(never)]
+#[cfg(not(target_feature = "avx2"))]
+#[cfg_attr(target_arch = "x86_64", cold)]
+pub unsafe fn format_escaped_str_impl_128(
+    odptr: *mut u8,
+    value_ptr: *const u8,
+    value_len: usize,
+) -> usize {
+    const STRIDE: usize = 16;
+    const STRIDE_SATURATION: u32 = u16::MAX as u32;
+    type StrVector = std::simd::u8x16;
+
+    impl_format_simd!(odptr, value_ptr, value_len);
+}
+
+#[cfg(target_arch = "x86_64")]
+#[inline(never)]
+#[cfg_attr(not(target_feature = "avx2"), target_feature(enable = "avx2"))]
+#[cfg_attr(not(target_feature = "avx2"), target_feature(enable = "bmi2"))]
+pub unsafe fn format_escaped_str_impl_256(
+    odptr: *mut u8,
+    value_ptr: *const u8,
+    value_len: usize,
+) -> usize {
+    const STRIDE: usize = 32;
+    const STRIDE_SATURATION: u32 = u32::MAX;
+    type StrVector = std::simd::u8x32;
+
+    impl_format_simd!(odptr, value_ptr, value_len);
+}

src/util.rs

@@ -70,6 +70,22 @@ macro_rules! unlikely {
     };
 }
 
+#[allow(unused_macros)]
+#[cfg(feature = "intrinsics")]
+macro_rules! likely {
+    ($exp:expr) => {
+        std::intrinsics::likely($exp)
+    };
+}
+
+#[allow(unused_macros)]
+#[cfg(not(feature = "intrinsics"))]
+macro_rules! likely {
+    ($exp:expr) => {
+        $exp
+    };
+}
+
 macro_rules! nonnull {
     ($exp:expr) => {
         unsafe { std::ptr::NonNull::new_unchecked($exp) }