Re-use the Neon code between the generator and the parser (#6)

* Reuse the SIMD code between the generator and the parser.

Author: samyron

ext/json/ext/generator/generator.c

@@ -284,8 +284,29 @@ static inline FORCE_INLINE char *copy_remaining_bytes(search_state *search, unsi
     return s;
 }
 
+static inline FORCE_INLINE int has_next_vector(void *state, size_t width)
+{
+    search_state *search = (search_state *) state;
+    return search->ptr + width <= search->end;
+}
+
+static inline FORCE_INLINE const char *ptr(void *state)
+{
+    return ((search_state *) state)->ptr;
+}
+
+static inline FORCE_INLINE void advance_by(void *state, size_t count)
+{
+    ((search_state *) state)->ptr += count;
+}
+
 #ifdef HAVE_SIMD_NEON
 
+static inline FORCE_INLINE void set_match_mask(void *state, uint64_t mask)
+{
+    ((search_state *) state)->matches_mask = mask;
+}
+
 static inline FORCE_INLINE unsigned char neon_next_match(search_state *search)
 {
     uint64_t mask = search->matches_mask;
@@ -304,20 +325,6 @@ static inline FORCE_INLINE unsigned char neon_next_match(search_state *search)
     return 1;
 }
 
-static inline FORCE_INLINE uint64_t neon_rules_update(const char *ptr)
-{
-    uint8x16_t chunk = vld1q_u8((const unsigned char *)ptr);
-
-    // Trick: c < 32 || c == 34 can be factored as c ^ 2 < 33
-    // https://lemire.me/blog/2025/04/13/detect-control-characters-quotes-and-backslashes-efficiently-using-swar/
-    const uint8x16_t too_low_or_dbl_quote = vcltq_u8(veorq_u8(chunk, vdupq_n_u8(2)), vdupq_n_u8(33));
-
-    uint8x16_t has_backslash = vceqq_u8(chunk, vdupq_n_u8('\\'));
-    uint8x16_t needs_escape  = vorrq_u8(too_low_or_dbl_quote, has_backslash);
-
-    return neon_match_mask(needs_escape);
-}
-
 static inline unsigned char search_escape_basic_neon(search_state *search)
 {
     if (RB_UNLIKELY(search->has_matches)) {
@@ -372,26 +379,30 @@ static inline unsigned char search_escape_basic_neon(search_state *search)
     * no bytes need to be escaped and we can continue to the next chunk. If the mask is not 0 then we
     * have at least one byte that needs to be escaped.
     */
-    while (search->ptr + sizeof(uint8x16_t) <= search->end) {
-        uint64_t mask = neon_rules_update(search->ptr);
 
-        if (!mask) {
-            search->ptr += sizeof(uint8x16_t);
-            continue;
-        }
-        search->matches_mask = mask;
+    // Do not extract this structor to a global variable or make it static without profiling and
+    // verifying that it does not cause performance regressions. clang 17 and gcc 14 currently
+    // inline these methods if this structure is defined like this.
+    NeonStringScanIterator iterator = {
+        .has_next_vector = has_next_vector,
+        .ptr = ptr,
+        .advance_by = advance_by,
+        .set_match_mask = set_match_mask
+    };
+
+    if (string_scan_simd_neon(&iterator, search)) {
         search->has_matches = true;
         search->chunk_base = search->ptr;
         search->chunk_end = search->ptr + sizeof(uint8x16_t);
         return neon_next_match(search);
     }
-
+    
     // There are fewer than 16 bytes left. 
     unsigned long remaining = (search->end - search->ptr);
     if (remaining >= SIMD_MINIMUM_THRESHOLD) {
         char *s = copy_remaining_bytes(search, sizeof(uint8x16_t), remaining);
 
-        uint64_t mask = neon_rules_update(s);
+        uint64_t mask = compute_chunk_mask_neon(s);
 
         if (!mask) {
             // Nothing to escape, ensure search_flush doesn't do anything by setting 
@@ -420,10 +431,11 @@ static inline unsigned char search_escape_basic_neon(search_state *search)
 
 #ifdef HAVE_SIMD_SSE2
 
-#define _mm_cmpge_epu8(a, b) _mm_cmpeq_epi8(_mm_max_epu8(a, b), a)
-#define _mm_cmple_epu8(a, b) _mm_cmpge_epu8(b, a)
-#define _mm_cmpgt_epu8(a, b) _mm_xor_si128(_mm_cmple_epu8(a, b), _mm_set1_epi8(-1))
-#define _mm_cmplt_epu8(a, b) _mm_cmpgt_epu8(b, a)
+static inline FORCE_INLINE void set_match_mask(void *state, int mask)
+{
+    ((search_state *) state)->matches_mask = mask;
+}
+
 
 static inline FORCE_INLINE unsigned char sse2_next_match(search_state *search)
 {
@@ -449,18 +461,6 @@ static inline FORCE_INLINE unsigned char sse2_next_match(search_state *search)
 #define TARGET_SSE2
 #endif
 
-static inline TARGET_SSE2 FORCE_INLINE int sse2_update(const char *ptr)
-{
-    __m128i chunk         = _mm_loadu_si128((__m128i const*)ptr);
-
-    // Trick: c < 32 || c == 34 can be factored as c ^ 2 < 33
-    // https://lemire.me/blog/2025/04/13/detect-control-characters-quotes-and-backslashes-efficiently-using-swar/
-    __m128i too_low_or_dbl_quote = _mm_cmplt_epu8(_mm_xor_si128(chunk, _mm_set1_epi8(2)), _mm_set1_epi8(33));
-    __m128i has_backslash = _mm_cmpeq_epi8(chunk, _mm_set1_epi8('\\'));
-    __m128i needs_escape  = _mm_or_si128(too_low_or_dbl_quote, has_backslash);
-    return _mm_movemask_epi8(needs_escape);
-}
-
 static inline TARGET_SSE2 FORCE_INLINE unsigned char search_escape_basic_sse2(search_state *search)
 {
     if (RB_UNLIKELY(search->has_matches)) {
@@ -479,17 +479,17 @@ static inline TARGET_SSE2 FORCE_INLINE unsigned char search_escape_basic_sse2(se
         }
     }
 
-    while (search->ptr + sizeof(__m128i) <= search->end) {
-        int needs_escape_mask = sse2_update(search->ptr);
-
-        if (needs_escape_mask == 0) {
-            search->ptr += sizeof(__m128i);
-            continue;
-        }
+    SSE2StringScanIterator iterator = {
+        .has_next_vector = has_next_vector,
+        .ptr = ptr,
+        .advance_by = advance_by,
+        .set_match_mask = set_match_mask
+    };
 
+    if (string_scan_simd_sse2(&iterator, search)) {
         search->has_matches = true;
-        search->matches_mask = needs_escape_mask;
         search->chunk_base = search->ptr;
+        search->chunk_end = search->ptr + sizeof(__m128i);
         return sse2_next_match(search);
     }
 
@@ -498,7 +498,7 @@ static inline TARGET_SSE2 FORCE_INLINE unsigned char search_escape_basic_sse2(se
     if (remaining >= SIMD_MINIMUM_THRESHOLD) {
         char *s = copy_remaining_bytes(search, sizeof(__m128i), remaining);
 
-        int needs_escape_mask = sse2_update(s);
+        int needs_escape_mask = compute_chunk_mask_sse2(s);
 
         if (needs_escape_mask == 0) {
             // Nothing to escape, ensure search_flush doesn't do anything by setting 

ext/json/ext/parser/parser.c

@@ -846,7 +846,7 @@ static inline VALUE json_push_value(JSON_ParserState *state, JSON_ParserConfig *
     return value;
 }
 
-static const bool string_scan[256] = {
+static const bool string_scan_table[256] = {
     // ASCII Control Characters
      1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
      1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
@@ -865,97 +865,76 @@ static const bool string_scan[256] = {
 #define FORCE_INLINE
 #endif
 
-static inline bool FORCE_INLINE string_scan_basic(JSON_ParserState *state)
-{
-    while (state->cursor < state->end) {
-        if (RB_UNLIKELY(string_scan[(unsigned char)*state->cursor])) {
-            return 1;
-        }
-        *state->cursor++;
-    }
-    return 0;
-}
-
-static bool (*string_scan_impl)(JSON_ParserState *);
+// Not used at the moment. 
+static SIMD_Implementation simd_impl = SIMD_NONE;
 
 #ifdef HAVE_SIMD
-#ifdef HAVE_SIMD_NEON
 
-static inline FORCE_INLINE bool string_scan_neon(JSON_ParserState *state)
+static inline FORCE_INLINE int has_next_vector(void *state, size_t width)
 {
-    while (state->cursor + sizeof(uint8x16_t) <= state->end) {
-        uint8x16_t chunk = vld1q_u8((const unsigned char *)state->cursor);
-
-        // Trick: c < 32 || c == 34 can be factored as c ^ 2 < 33
-        // https://lemire.me/blog/2025/04/13/detect-control-characters-quotes-and-backslashes-efficiently-using-swar/
-        const uint8x16_t too_low_or_dbl_quote = vcltq_u8(veorq_u8(chunk, vdupq_n_u8(2)), vdupq_n_u8(33));
-
-        uint8x16_t has_backslash = vceqq_u8(chunk, vdupq_n_u8('\\'));
-        uint8x16_t needs_escape  = vorrq_u8(too_low_or_dbl_quote, has_backslash);
-        uint64_t mask = neon_match_mask(needs_escape);
-        if (mask) {
-            state->cursor += trailing_zeros64(mask) >> 2;
-            return true;
-        }
-
-        state->cursor += sizeof(uint8x16_t);
-    }
-
-    if (state->cursor < state->end) {
-        return string_scan_basic(state);
-    }
-    return 0;
+    JSON_ParserState *s = state;
+    return s->cursor + width <= s->end;
 }
-#endif /* HAVE_SIMD_NEON */
-
-#ifdef HAVE_SIMD_SSE2
 
-#define _mm_cmpge_epu8(a, b) _mm_cmpeq_epi8(_mm_max_epu8(a, b), a)
-#define _mm_cmple_epu8(a, b) _mm_cmpge_epu8(b, a)
-#define _mm_cmpgt_epu8(a, b) _mm_xor_si128(_mm_cmple_epu8(a, b), _mm_set1_epi8(-1))
-#define _mm_cmplt_epu8(a, b) _mm_cmpgt_epu8(b, a)
+static inline FORCE_INLINE const char *ptr(void *state)                           { return ((JSON_ParserState *) state)->cursor; }
+static inline FORCE_INLINE void        advance_by(void *state, size_t count)      { ((JSON_ParserState *) state)->cursor += count; }
 
-#if defined(__clang__) || defined(__GNUC__)
-#define TARGET_SSE2 __attribute__((target("sse2")))
-#else
-#define TARGET_SSE2
-#endif
+#ifdef HAVE_SIMD_NEON
+static inline FORCE_INLINE void        set_match_mask(void *state, uint64_t mask) { advance_by(state, trailing_zeros64(mask) >> 2); }
+#elif defined(HAVE_SIMD_SSE2)
+static inline FORCE_INLINE void        set_match_mask(void *state, int mask)      { advance_by(state, trailing_zeros(mask)); }
+#endif /* HAVE_SIMD_NEON */
+#endif /* HAVE_SIMD */
 
-static inline TARGET_SSE2 FORCE_INLINE bool string_scan_sse2(JSON_ParserState *state)
+static inline bool FORCE_INLINE string_scan(JSON_ParserState *state)
 {
-    while (state->cursor + sizeof(__m128i) <= state->end) {
-        __m128i chunk = _mm_loadu_si128((__m128i const*)state->cursor);
-
-        // Trick: c < 32 || c == 34 can be factored as c ^ 2 < 33
-        // https://lemire.me/blog/2025/04/13/detect-control-characters-quotes-and-backslashes-efficiently-using-swar/
-        __m128i too_low_or_dbl_quote = _mm_cmplt_epu8(_mm_xor_si128(chunk, _mm_set1_epi8(2)), _mm_set1_epi8(33));
-        __m128i has_backslash = _mm_cmpeq_epi8(chunk, _mm_set1_epi8('\\'));
-        __m128i needs_escape  = _mm_or_si128(too_low_or_dbl_quote, has_backslash);
-        int mask = _mm_movemask_epi8(needs_escape);
-        if (mask) {
-            state->cursor += trailing_zeros(mask);
-            return true;
-        }
+#ifdef HAVE_SIMD
+#if defined(HAVE_SIMD_NEON)
+    // Do not extract this structor to a global variable or make it static without profiling and
+    // verifying that it does not cause performance regressions. clang 17 and gcc 14 currently
+    // inline these methods if this structure is defined like this.
+    NeonStringScanIterator iterator = {
+        .has_next_vector = has_next_vector,
+        .ptr = ptr,
+        .advance_by = advance_by,
+        .set_match_mask = set_match_mask
+    };
 
-        state->cursor += sizeof(__m128i);
+    if (string_scan_simd_neon(&iterator, state)) {
+        return 1;
+    }
+#elif defined(HAVE_SIMD_SSE2)
+    if (simd_impl == SIMD_SSE2) {
+        SSE2StringScanIterator iterator = {
+            .has_next_vector = has_next_vector,
+            .ptr = ptr,
+            .advance_by = advance_by,
+            .set_match_mask = set_match_mask
+        };
+
+        if (string_scan_simd_sse2(&iterator, state)) {
+            return 1;
+        }
     }
+#endif /* HAVE_SIMD_NEON or HAVE_SIMD_SSE2 */
+#endif /* HAVE_SIMD */
 
-    if (state->cursor < state->end) {
-        return string_scan_basic(state);
+    while (state->cursor < state->end) {
+        if (RB_UNLIKELY(string_scan_table[(unsigned char)*state->cursor])) {
+            return 1;
+        }
+        *state->cursor++;
     }
     return 0;
 }
-#endif
-
-#endif /* HAVE_SIMD */
 
 static inline VALUE json_parse_string(JSON_ParserState *state, JSON_ParserConfig *config, bool is_name)
 {
     state->cursor++;
     const char *start = state->cursor;
     bool escaped = false;
 
-    while (RB_UNLIKELY(string_scan_impl(state))) {
+    while (RB_UNLIKELY(string_scan(state))) {
         switch (*state->cursor) {
             case '"': {
                 VALUE string = json_decode_string(state, config, start, state->cursor, escaped, is_name);
@@ -1504,21 +1483,7 @@ void Init_parser(void)
     utf8_encindex = rb_utf8_encindex();
     enc_utf8 = rb_utf8_encoding();
 
-switch(find_simd_implementation()) {
 #ifdef HAVE_SIMD
-#ifdef HAVE_SIMD_NEON
-        case SIMD_NEON:
-            string_scan_impl = string_scan_neon;
-            break;
-#endif /* HAVE_SIMD_NEON */
-#ifdef HAVE_SIMD_SSE2
-        case SIMD_SSE2:
-            string_scan_impl = string_scan_sse2;
-            break;
-#endif /* HAVE_SIMD_SSE2 */
-#endif /* HAVE_SIMD */
-        default:
-            string_scan_impl = string_scan_basic;
-            break;
-    }
+    simd_impl = find_simd_implementation();
+#endif
 }