Implement AVX512-accelerated JSON string escaping with runtime feature detection

Co-authored-by: Brooooooklyn <3468483+Brooooooklyn@users.noreply.github.com>

Author: Copilot

Cargo.toml

@@ -10,7 +10,7 @@ license = "MIT"
 publish = true
 readme = "README.md"
 repository = "https://github.com/oxc-project/oxc-sourcemap"
-rust-version = "1.85.0"
+rust-version = "1.89.0"
 description = "Basic sourcemap handling for Rust"
 
 # <https://doc.rust-lang.org/rustc/lints/listing/allowed-by-default.html>

benches/simple.rs

@@ -1,5 +1,29 @@
 use criterion::{Criterion, criterion_group, criterion_main};
-use oxc_sourcemap::{SourceMap, SourceMapBuilder};
+use oxc_sourcemap::{SourceMap, SourceMapBuilder, escape_json_string, escape_json_string_fallback};
+
+pub fn bench_json_escaping(c: &mut Criterion) {
+    let long_clean = "abcdefghijklmnopqrstuvwxyz".repeat(50);
+    let long_quotes = "\"test\"".repeat(100);
+    
+    let test_strings = vec![
+        ("short_clean", "simple string without escapes"),
+        ("short_quotes", "string with \"quotes\" and \\backslashes"),
+        ("short_control", "string with\ncontrol\tchars\r"),
+        ("long_clean", long_clean.as_str()), // Long string without escapes  
+        ("long_quotes", long_quotes.as_str()), // Many escapes
+        ("mixed", "mixed: \"quotes\", \\backslashes, \ncontrol\tchars, and regular text"),
+    ];
+
+    for (name, test_str) in test_strings {
+        c.bench_function(&format!("escape_fallback_{}", name), |b| {
+            b.iter(|| escape_json_string_fallback(test_str));
+        });
+        
+        c.bench_function(&format!("escape_avx512_{}", name), |b| {
+            b.iter(|| escape_json_string(test_str));
+        });
+    }
+}
 
 pub fn bench(c: &mut Criterion) {
     let input = r#"{
@@ -35,5 +59,5 @@ pub fn bench(c: &mut Criterion) {
     });
 }
 
-criterion_group!(sourcemap, bench);
+criterion_group!(sourcemap, bench, bench_json_escaping);
 criterion_main!(sourcemap);

src/encode.rs

@@ -330,7 +330,22 @@ impl<'a> PreAllocatedString<'a> {
     }
 }
 
-fn escape_json_string<S: AsRef<str>>(s: S) -> String {
+pub fn escape_json_string<S: AsRef<str>>(s: S) -> String {
+    let s = s.as_ref();
+    
+    // Use AVX512 acceleration if available on x86_64
+    #[cfg(target_arch = "x86_64")]
+    {
+        if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") {
+            return unsafe { escape_json_string_avx512(s) };
+        }
+    }
+    
+    // Fallback to serde_json implementation
+    escape_json_string_fallback(s)
+}
+
+pub fn escape_json_string_fallback<S: AsRef<str>>(s: S) -> String {
     let s = s.as_ref();
     let mut escaped_buf = Vec::with_capacity(s.len() * 2 + 2);
     // This call is infallible as only error it can return is if the writer errors.
@@ -340,6 +355,81 @@ fn escape_json_string<S: AsRef<str>>(s: S) -> String {
     unsafe { String::from_utf8_unchecked(escaped_buf) }
 }
 
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx512f,avx512bw")]
+unsafe fn escape_json_string_avx512(s: &str) -> String {
+    use std::arch::x86_64::*;
+    
+    let bytes = s.as_bytes();
+    let mut result = Vec::with_capacity(s.len() * 2 + 2);
+    
+    // Add opening quote
+    result.push(b'"');
+    
+    let mut i = 0;
+    
+    // Process 64-byte chunks with AVX512
+    while i + 64 <= bytes.len() {
+        unsafe {
+            let chunk = _mm512_loadu_si512(bytes.as_ptr().add(i) as *const __m512i);
+            
+            // Check for characters that need escaping
+            // ASCII control characters (0x00-0x1F), quote (0x22), backslash (0x5C)
+            let control_mask = _mm512_cmplt_epu8_mask(chunk, _mm512_set1_epi8(0x20));
+            let quote_mask = _mm512_cmpeq_epu8_mask(chunk, _mm512_set1_epi8(b'"' as i8));
+            let backslash_mask = _mm512_cmpeq_epu8_mask(chunk, _mm512_set1_epi8(b'\\' as i8));
+            
+            let escape_mask = control_mask | quote_mask | backslash_mask;
+            
+            if escape_mask == 0 {
+                // No characters need escaping, copy directly
+                result.extend_from_slice(&bytes[i..i + 64]);
+                i += 64;
+            } else {
+                // Process byte by byte for this chunk due to escaping needed
+                for j in 0..64 {
+                    let byte = bytes[i + j];
+                    append_escaped_byte(&mut result, byte);
+                }
+                i += 64;
+            }
+        }
+    }
+    
+    // Process remaining bytes
+    while i < bytes.len() {
+        append_escaped_byte(&mut result, bytes[i]);
+        i += 1;
+    }
+    
+    // Add closing quote
+    result.push(b'"');
+    
+    // Safety: We only write valid UTF-8 sequences
+    unsafe { String::from_utf8_unchecked(result) }
+}
+
+#[inline]
+fn append_escaped_byte(result: &mut Vec<u8>, byte: u8) {
+    match byte {
+        b'"' => result.extend_from_slice(b"\\\""),
+        b'\\' => result.extend_from_slice(b"\\\\"),
+        b'\x08' => result.extend_from_slice(b"\\b"),
+        b'\x0C' => result.extend_from_slice(b"\\f"),
+        b'\n' => result.extend_from_slice(b"\\n"),
+        b'\r' => result.extend_from_slice(b"\\r"),
+        b'\t' => result.extend_from_slice(b"\\t"),
+        0x00..=0x1F => {
+            // Other control characters as unicode escapes
+            result.extend_from_slice(b"\\u00");
+            let hex_chars = b"0123456789abcdef";
+            result.push(hex_chars[(byte >> 4) as usize]);
+            result.push(hex_chars[(byte & 0x0F) as usize]);
+        }
+        _ => result.push(byte),
+    }
+}
+
 #[test]
 fn test_escape_json_string() {
     const FIXTURES: &[(char, &str)] = &[
@@ -364,6 +454,61 @@ fn test_escape_json_string() {
     }
 }
 
+#[test]
+fn test_avx512_vs_fallback_consistency() {
+    let long_x = "x".repeat(100);
+    let long_quotes = "\"".repeat(100);
+    let long_control = "\n".repeat(100);
+    
+    let test_strings = vec![
+        "",
+        "simple",
+        "with\"quotes",
+        "with\\backslashes",
+        "with\ncontrol\tchars\r",
+        "🚀 unicode 虎",
+        &long_x, // Test longer strings to trigger AVX512 path
+        &long_quotes, // Many quotes to test escaping
+        &long_control, // Many control chars
+        "mixed content with \"quotes\", \\backslashes, and \ncontrol\tchars",
+    ];
+
+    for test_str in test_strings {
+        let fallback_result = escape_json_string_fallback(test_str);
+        let main_result = escape_json_string(test_str);
+        assert_eq!(
+            fallback_result, main_result,
+            "Results differ for input: {:?}",
+            test_str
+        );
+    }
+}
+
+#[test] 
+#[cfg(target_arch = "x86_64")]
+fn test_avx512_direct() {
+    // Test the AVX512 function directly if AVX512 is available
+    if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") {
+        let long_a = "a".repeat(128);
+        let test_strings = vec![
+            "simple test",
+            "with \"quotes\"",
+            "with\ncontrol\rchars\t",
+            long_a.as_str(), // Longer than one AVX512 chunk
+        ];
+        
+        for test_str in test_strings {
+            let avx512_result = unsafe { escape_json_string_avx512(test_str) };
+            let fallback_result = escape_json_string_fallback(test_str);
+            assert_eq!(
+                avx512_result, fallback_result,
+                "AVX512 and fallback results differ for: {:?}",
+                test_str
+            );
+        }
+    }
+}
+
 #[test]
 fn test_encode() {
     let input = r#"{

src/lib.rs

@@ -12,6 +12,7 @@ pub mod napi;
 
 pub use concat_sourcemap_builder::ConcatSourceMapBuilder;
 pub use decode::JSONSourceMap;
+pub use encode::{escape_json_string, escape_json_string_fallback};
 pub use error::Error;
 pub use sourcemap::SourceMap;
 pub use sourcemap_builder::SourceMapBuilder;