Improve alignment

Signed-off-by: Heinz N. Gies <heinz@licenser.net>

Author: Licenser

src/impls/avx2/stage1.rs

@@ -44,14 +44,14 @@ impl Stage1Parse for SimdInput {
     type Utf8Validator = simdutf8::basic::imp::x86::avx2::ChunkedUtf8ValidatorImp;
     type SimdRepresentation = __m256i;
     #[cfg_attr(not(feature = "no-inline"), inline)]
-    // _mm256_loadu_si256 does not need alignment
+    // _mm256_loadu_si256 does not need alignment we allign our input so we can use _mm256_loadu_si256
     #[allow(clippy::cast_ptr_alignment)]
     #[target_feature(enable = "avx2")]
-    unsafe fn new(ptr: &[u8]) -> Self {
+    unsafe fn new(ptr: [u8; SIMDINPUT_LENGTH]) -> Self {
         unsafe {
             Self {
-                v0: _mm256_loadu_si256(ptr.as_ptr().cast::<__m256i>()),
-                v1: _mm256_loadu_si256(ptr.as_ptr().add(32).cast::<__m256i>()),
+                v0: _mm256_load_si256(ptr.as_ptr().cast::<__m256i>()),
+                v1: _mm256_load_si256(ptr.as_ptr().add(32).cast::<__m256i>()),
             }
         }
     }

src/impls/native/stage1.rs

@@ -292,13 +292,13 @@ pub(crate) struct SimdInput {
 impl Stage1Parse for SimdInput {
     type Utf8Validator = super::ChunkedUtf8ValidatorImp;
     type SimdRepresentation = V128;
-    unsafe fn new(ptr: &[u8]) -> Self {
+    unsafe fn new(ptr: [u8; SIMDINPUT_LENGTH]) -> Self {
         unsafe {
             SimdInput {
-                v0: *(ptr.as_ptr().cast::<V128>()),
-                v1: *(ptr.as_ptr().add(16).cast::<V128>()),
-                v2: *(ptr.as_ptr().add(32).cast::<V128>()),
-                v3: *(ptr.as_ptr().add(48).cast::<V128>()),
+                v0: ptr.as_ptr().cast::<V128>().read(),
+                v1: ptr.as_ptr().add(16).cast::<V128>().read(),
+                v2: ptr.as_ptr().add(32).cast::<V128>().read(),
+                v3: ptr.as_ptr().add(48).cast::<V128>().read(),
             }
         }
     }

src/impls/neon/stage1.rs

@@ -1,4 +1,4 @@
-use crate::Stage1Parse;
+use crate::{SIMDINPUT_LENGTH, Stage1Parse, macros::static_cast_i32};
 use std::arch::aarch64::{
     int8x16_t, int32x4_t, uint8x16_t, vaddq_s32, vandq_u8, vceqq_u8, vcleq_u8, vdupq_n_s8,
     vgetq_lane_u64, vld1q_u8, vmovq_n_u8, vpaddq_u8, vqtbl1q_u8, vreinterpretq_u8_s8,
@@ -57,13 +57,13 @@ impl Stage1Parse for SimdInput {
     type Utf8Validator = simdutf8::basic::imp::aarch64::neon::ChunkedUtf8ValidatorImp;
     type SimdRepresentation = int8x16_t;
     #[cfg_attr(not(feature = "no-inline"), inline)]
-    unsafe fn new(ptr: &[u8]) -> Self {
+    unsafe fn new(ptr: [u8; SIMDINPUT_LENGTH]) -> Self {
         unsafe {
             Self {
-                v0: vld1q_u8(ptr.as_ptr().cast::<u8>()),
-                v1: vld1q_u8(ptr.as_ptr().add(16).cast::<u8>()),
-                v2: vld1q_u8(ptr.as_ptr().add(32).cast::<u8>()),
-                v3: vld1q_u8(ptr.as_ptr().add(48).cast::<u8>()),
+                v0: vld1q_u8(ptr.as_ptr()),
+                v1: vld1q_u8(ptr.as_ptr().add(16)),
+                v2: vld1q_u8(ptr.as_ptr().add(32)),
+                v3: vld1q_u8(ptr.as_ptr().add(48)),
             }
         }
     }

src/impls/portable/stage1.rs

@@ -10,9 +10,9 @@ impl Stage1Parse for SimdInput {
     type Utf8Validator = simdutf8::basic::imp::portable::ChunkedUtf8ValidatorImp;
     type SimdRepresentation = u8x64;
     #[cfg_attr(not(feature = "no-inline"), inline)]
-    unsafe fn new(ptr: &[u8]) -> Self {
+    unsafe fn new(ptr: [u8; SIMDINPUT_LENGTH]) -> Self {
         Self {
-            v: u8x64::from_array(*ptr.as_ptr().cast::<[u8; 64]>()),
+            v: u8x64::from_array(ptr),
         }
     }
 

src/impls/simd128/stage1.rs

@@ -18,7 +18,7 @@ impl Stage1Parse for SimdInput {
 
     #[cfg_attr(not(feature = "no-inline"), inline)]
     #[allow(clippy::cast_ptr_alignment)]
-    unsafe fn new(ptr: &[u8]) -> Self {
+    unsafe fn new(ptr: [u8; SIMDINPUT_LENGTH]) -> Self {
         Self {
             v0: v128_load(ptr.as_ptr().cast::<v128>()),
             v1: v128_load(ptr.as_ptr().add(16).cast::<v128>()),

src/impls/sse42/stage1.rs

@@ -48,13 +48,13 @@ impl Stage1Parse for SimdInput {
     #[target_feature(enable = "sse4.2")]
     #[cfg_attr(not(feature = "no-inline"), inline)]
     #[allow(clippy::cast_ptr_alignment)]
-    unsafe fn new(ptr: &[u8]) -> Self {
+    unsafe fn new(ptr: [u8; SIMDINPUT_LENGTH]) -> Self {
         unsafe {
             Self {
-                v0: _mm_loadu_si128(ptr.as_ptr().cast::<arch::__m128i>()),
-                v1: _mm_loadu_si128(ptr.as_ptr().add(16).cast::<arch::__m128i>()),
-                v2: _mm_loadu_si128(ptr.as_ptr().add(32).cast::<arch::__m128i>()),
-                v3: _mm_loadu_si128(ptr.as_ptr().add(48).cast::<arch::__m128i>()),
+                v0: _mm_load_si128(ptr.as_ptr().cast::<arch::__m128i>()),
+                v1: _mm_load_si128(ptr.as_ptr().add(16).cast::<arch::__m128i>()),
+                v2: _mm_load_si128(ptr.as_ptr().add(32).cast::<arch::__m128i>()),
+                v3: _mm_load_si128(ptr.as_ptr().add(48).cast::<arch::__m128i>()),
             }
         }
     }

src/lib.rs

@@ -150,7 +150,7 @@ pub(crate) trait Stage1Parse {
     type Utf8Validator: ChunkedUtf8Validator;
     type SimdRepresentation;
 
-    unsafe fn new(ptr: &[u8]) -> Self;
+    unsafe fn new(ptr: [u8; SIMDINPUT_LENGTH]) -> Self;
 
     unsafe fn compute_quote_mask(quote_bits: u64) -> u64;
 
@@ -665,6 +665,7 @@ impl Deserializer<'_> {
     #[cfg_attr(not(feature = "no-inline"), inline)]
     pub(crate) unsafe fn find_structural_bits(
         input: &[u8],
+        len: usize,
         structural_indexes: &mut Vec<u32>,
     ) -> std::result::Result<(), ErrorType> {
         // This is a nasty hack, we don't have a chunked implementation for native rust
@@ -675,18 +676,23 @@ impl Deserializer<'_> {
         };
         #[cfg(not(feature = "portable"))]
         unsafe {
-            Self::_find_structural_bits::<impls::native::SimdInput>(input, structural_indexes)
+            Self::_find_structural_bits::<impls::native::SimdInput>(input, len, structural_indexes)
         }
     }
 
     #[cfg(all(feature = "portable", not(feature = "runtime-detection")))]
     #[cfg_attr(not(feature = "no-inline"), inline)]
     pub(crate) unsafe fn find_structural_bits(
         input: &[u8],
+        len: usize,
         structural_indexes: &mut Vec<u32>,
     ) -> std::result::Result<(), ErrorType> {
         unsafe {
-            Self::_find_structural_bits::<impls::portable::SimdInput>(input, structural_indexes)
+            Self::_find_structural_bits::<impls::portable::SimdInput>(
+                input,
+                len,
+                structural_indexes,
+            )
         }
     }
 
@@ -698,9 +704,12 @@ impl Deserializer<'_> {
     #[cfg_attr(not(feature = "no-inline"), inline)]
     pub(crate) unsafe fn find_structural_bits(
         input: &[u8],
+        len: usize,
         structural_indexes: &mut Vec<u32>,
     ) -> std::result::Result<(), ErrorType> {
-        unsafe { Self::_find_structural_bits::<impls::avx2::SimdInput>(input, structural_indexes) }
+        unsafe {
+            Self::_find_structural_bits::<impls::avx2::SimdInput>(input, len, structural_indexes)
+        }
     }
 
     #[cfg(all(
@@ -712,28 +721,35 @@ impl Deserializer<'_> {
     #[cfg_attr(not(feature = "no-inline"), inline)]
     pub(crate) unsafe fn find_structural_bits(
         input: &[u8],
+        len: usize,
         structural_indexes: &mut Vec<u32>,
     ) -> std::result::Result<(), ErrorType> {
-        unsafe { Self::_find_structural_bits::<impls::sse42::SimdInput>(input, structural_indexes) }
+        unsafe {
+            Self::_find_structural_bits::<impls::sse42::SimdInput>(input, len, structural_indexes)
+        }
     }
 
     #[cfg(all(target_arch = "aarch64", not(feature = "portable")))]
     #[cfg_attr(not(feature = "no-inline"), inline)]
     pub(crate) unsafe fn find_structural_bits(
-        input: &[u8],
+        input: &AlignedBuf,
+        len: usize,
         structural_indexes: &mut Vec<u32>,
     ) -> std::result::Result<(), ErrorType> {
-        unsafe { Self::_find_structural_bits::<impls::neon::SimdInput>(input, structural_indexes) }
+        unsafe {
+            Self::_find_structural_bits::<impls::neon::SimdInput>(input, len, structural_indexes)
+        }
     }
 
     #[cfg(all(target_feature = "simd128", not(feature = "portable")))]
     #[cfg_attr(not(feature = "no-inline"), inline)]
     pub(crate) unsafe fn find_structural_bits(
         input: &[u8],
+        len: usize,
         structural_indexes: &mut Vec<u32>,
     ) -> std::result::Result<(), ErrorType> {
         unsafe {
-            Self::_find_structural_bits::<impls::simd128::SimdInput>(input, structural_indexes)
+            Self::_find_structural_bits::<impls::simd128::SimdInput>(input, len, structural_indexes)
         }
     }
 }
@@ -795,7 +811,7 @@ impl<'de> Deserializer<'de> {
         buffer: &mut Buffers,
         tape: &mut Vec<Node<'de>>,
     ) -> Result<()> {
-        const LOTS_OF_ZOERS: [u8; SIMDINPUT_LENGTH] = [0; SIMDINPUT_LENGTH];
+        const LOTS_OF_ZOERS: [u8; SIMDINPUT_LENGTH] = [0x20; SIMDINPUT_LENGTH];
         let len = input.len();
         let simd_safe_len = len + SIMDINPUT_LENGTH;
 
@@ -830,7 +846,7 @@ impl<'de> Deserializer<'de> {
             // safety: all bytes are initialized
             input_buffer.set_len(simd_safe_len);
 
-            Self::find_structural_bits(input, &mut buffer.structural_indexes)
+            Self::find_structural_bits(input_buffer, input.len(), &mut buffer.structural_indexes)
                 .map_err(Error::generic)?;
         };
 
@@ -881,10 +897,11 @@ impl<'de> Deserializer<'de> {
     #[cfg_attr(not(feature = "no-inline"), inline)]
     #[allow(clippy::cast_possible_truncation)]
     pub(crate) unsafe fn _find_structural_bits<S: Stage1Parse>(
-        input: &[u8],
+        input: &AlignedBuf,
+        len: usize,
         structural_indexes: &mut Vec<u32>,
     ) -> std::result::Result<(), ErrorType> {
-        let len = input.len();
+        // let len = input.len();
         // 8 is a heuristic number to estimate it turns out a rate of 1/8 structural characters
         // leads almost never to relocations.
         structural_indexes.clear();
@@ -916,18 +933,18 @@ impl<'de> Deserializer<'de> {
         // expensive carryless multiply in the previous step with this work
         let mut structurals: u64 = 0;
 
-        let lenminus64: usize = if len < 64 { 0 } else { len - 64 };
+        // let lenminus64: usize = if len < 64 { 0 } else { len - 64 };
         let mut idx: usize = 0;
         let mut error_mask: u64 = 0; // for unescaped characters within strings (ASCII code points < 0x20)
 
-        while idx < lenminus64 {
+        while idx <= len / SIMDINPUT_LENGTH {
             /*
             #ifndef _MSC_VER
               __builtin_prefetch(buf + idx + 128);
             #endif
              */
-            let chunk = unsafe { input.get_kinda_unchecked(idx..idx + 64) };
-            unsafe { utf8_validator.update_from_chunks(chunk) };
+            let chunk: [u8; SIMDINPUT_LENGTH] = unsafe { input.load_register(idx) };
+            unsafe { utf8_validator.update_from_chunks(&chunk) };
 
             let input = unsafe { S::new(chunk) };
             // detect odd sequences of backslashes
@@ -946,7 +963,7 @@ impl<'de> Deserializer<'de> {
 
             // take the previous iterations structural bits, not our current iteration,
             // and flatten
-            unsafe { S::flatten_bits(structural_indexes, idx as u32, structurals) };
+            unsafe { S::flatten_bits(structural_indexes, (idx * 64) as u32, structurals) };
 
             let mut whitespace: u64 = 0;
             unsafe { input.find_whitespace_and_structurals(&mut whitespace, &mut structurals) };
@@ -959,60 +976,15 @@ impl<'de> Deserializer<'de> {
                 quote_bits,
                 &mut prev_iter_ends_pseudo_pred,
             );
-            idx += SIMDINPUT_LENGTH;
+            idx += 1;
         }
 
-        // we use a giant copy-paste which is ugly.
-        // but otherwise the string needs to be properly padded or else we
-        // risk invalidating the UTF-8 checks.
-        if idx < len {
-            let mut tmpbuf: [u8; SIMDINPUT_LENGTH] = [0x20; SIMDINPUT_LENGTH];
-            unsafe {
-                tmpbuf
-                    .as_mut_ptr()
-                    .copy_from(input.as_ptr().add(idx), len - idx);
-            };
-            unsafe { utf8_validator.update_from_chunks(&tmpbuf) };
-
-            let input = unsafe { S::new(&tmpbuf) };
-
-            // detect odd sequences of backslashes
-            let odd_ends: u64 =
-                input.find_odd_backslash_sequences(&mut prev_iter_ends_odd_backslash);
-
-            // detect insides of quote pairs ("quote_mask") and also our quote_bits
-            // themselves
-            let mut quote_bits: u64 = 0;
-            let quote_mask: u64 = input.find_quote_mask_and_bits(
-                odd_ends,
-                &mut prev_iter_inside_quote,
-                &mut quote_bits,
-                &mut error_mask,
-            );
-
-            // take the previous iterations structural bits, not our current iteration,
-            // and flatten
-            unsafe { S::flatten_bits(structural_indexes, idx as u32, structurals) };
-
-            let mut whitespace: u64 = 0;
-            unsafe { input.find_whitespace_and_structurals(&mut whitespace, &mut structurals) };
-
-            // fixup structurals to reflect quotes and add pseudo-structural characters
-            structurals = S::finalize_structurals(
-                structurals,
-                whitespace,
-                quote_mask,
-                quote_bits,
-                &mut prev_iter_ends_pseudo_pred,
-            );
-            idx += SIMDINPUT_LENGTH;
-        }
         // This test isn't in upstream, for some reason the error mask is et for then.
         if prev_iter_inside_quote != 0 {
             return Err(ErrorType::Syntax);
         }
         // finally, flatten out the remaining structurals from the last iteration
-        unsafe { S::flatten_bits(structural_indexes, idx as u32, structurals) };
+        unsafe { S::flatten_bits(structural_indexes, (idx * 64) as u32, structurals) };
 
         // a valid JSON file cannot have zero structural indexes - we should have
         // found something (note that we compare to 1 as we always add the root!)
@@ -1051,22 +1023,40 @@ impl AlignedBuf {
     /// Creates a new buffer that is  aligned with the simd register size
     #[must_use]
     pub fn with_capacity(capacity: usize) -> Self {
-        let Ok(layout) = Layout::from_size_align(capacity, SIMDJSON_PADDING) else {
-            Self::capacity_overflow()
+        let offset = capacity % SIMDINPUT_LENGTH;
+        let capacity = if offset == 0 {
+            capacity
+        } else {
+            capacity + SIMDINPUT_LENGTH - offset
         };
+
         if mem::size_of::<usize>() < 8 && capacity > isize::MAX as usize {
             Self::capacity_overflow()
         }
+        let layout = match Layout::from_size_align(capacity, SIMDINPUT_LENGTH) {
+            Ok(layout) => layout,
+            Err(_) => Self::capacity_overflow(),
+        };
+
+        let inner = match unsafe { NonNull::new(alloc(layout)) } {
+            Some(ptr) => ptr,
+            None => handle_alloc_error(layout),
+        };
+        Self {
+            layout,
+            capacity,
+            len: 0,
+            inner,
+        }
+    }
+
+    unsafe fn load_register(&self, idx: usize) -> [u8; SIMDINPUT_LENGTH] {
         unsafe {
-            let Some(inner) = NonNull::new(alloc(layout)) else {
-                handle_alloc_error(layout)
-            };
-            Self {
-                layout,
-                capacity,
-                len: 0,
-                inner,
-            }
+            self.inner
+                .as_ptr()
+                .cast::<[u8; SIMDINPUT_LENGTH]>()
+                .add(idx)
+                .read()
         }
     }