value_trait/
generator.rs

1// This is mostly taken from json-rust's codegen
2// as it seems to perform well and it makes sense to see
3// if we can adopt the approach
4//
5// https://github.com/maciejhirsz/json-rust/blob/master/src/codegen.rs
6
7macro_rules! stry {
8    ($e:expr) => {
9        match $e {
10            ::std::result::Result::Ok(val) => val,
11            ::std::result::Result::Err(err) => return ::std::result::Result::Err(err),
12        }
13    };
14}
15
16use std::io;
17use std::io::Write;
18use std::ptr;
19
20const QU: u8 = b'"';
21const BS: u8 = b'\\';
22const BB: u8 = b'b';
23const TT: u8 = b't';
24const NN: u8 = b'n';
25const FF: u8 = b'f';
26const RR: u8 = b'r';
27const UU: u8 = b'u';
28const __: u8 = 0;
29
30// Look up table for characters that need escaping in a product string
31pub(crate) static ESCAPED: [u8; 256] = [
32    // 0   1   2   3   4   5   6   7   8   9   A   B   C   D   E   F
33    UU, UU, UU, UU, UU, UU, UU, UU, BB, TT, NN, UU, FF, RR, UU, UU, // 0
34    UU, UU, UU, UU, UU, UU, UU, UU, UU, UU, UU, UU, UU, UU, UU, UU, // 1
35    __, __, QU, __, __, __, __, __, __, __, __, __, __, __, __, __, // 2
36    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 3
37    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 4
38    __, __, __, __, __, __, __, __, __, __, __, __, BS, __, __, __, // 5
39    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 6
40    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 7
41    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 8
42    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 9
43    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // A
44    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // B
45    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // C
46    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // D
47    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // E
48    __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // F
49];
50
51// taken from https://github.com/serde-rs/json/blob/4354fc3eb2232ee0ba9a9a23acce107a980a6dc0/src/ser.rs#L1790
52// This is called rarely
53#[inline(never)]
54fn u_encode<W>(w: &mut W, byte: u8) -> io::Result<()>
55where
56    W: Write,
57{
58    static HEX_DIGITS: [u8; 16] = *b"0123456789abcdef";
59    let bytes = [
60        b'\\',
61        b'u',
62        b'0',
63        b'0',
64        HEX_DIGITS[(byte >> 4) as usize],
65        HEX_DIGITS[(byte & 0xF) as usize],
66    ];
67    w.write_all(&bytes)
68}
69
70/// Base generator trait
71pub trait BaseGenerator {
72    /// The writer
73    type T: Write;
74
75    /// returns the writer
76    fn get_writer(&mut self) -> &mut Self::T;
77
78    /// Write a slice
79    /// # Errors
80    /// if the write fails
81    #[inline]
82    fn write(&mut self, slice: &[u8]) -> io::Result<()> {
83        self.get_writer().write_all(slice)
84    }
85
86    /// Write a char
87    /// # Errors
88    /// if the write fails
89    #[inline]
90    fn write_char(&mut self, ch: u8) -> io::Result<()> {
91        self.get_writer().write_all(&[ch])
92    }
93
94    /// write with minimum
95    /// # Errors
96    /// if the write fails
97    fn write_min(&mut self, slice: &[u8], min: u8) -> io::Result<()>;
98
99    /// writes new line
100    /// # Errors
101    /// if the write fails
102    #[inline]
103    fn new_line(&mut self) -> io::Result<()> {
104        Ok(())
105    }
106
107    /// indents one step
108    #[inline]
109    fn indent(&mut self) {}
110
111    /// dedents one step
112    #[inline]
113    fn dedent(&mut self) {}
114
115    /// writes a string
116    /// # Errors
117    /// if the write fails
118    #[inline]
119    fn write_string(&mut self, string: &str) -> io::Result<()> {
120        stry!(self.write_char(b'"'));
121        stry!(self.write_string_content(string));
122        self.write_char(b'"')
123    }
124
125    /// writes a string
126    /// # Errors
127    /// if the write fails
128    #[inline]
129    fn write_string_content(&mut self, string: &str) -> io::Result<()> {
130        let mut string = string.as_bytes();
131        unsafe {
132            // Looking at the table above the lower 5 bits are entirely
133            // quote characters that gives us a bitmask of 0x1f for that
134            // region, only quote (`"`) and backslash (`\`) are not in
135            // this range.
136            stry!(self.write_str_simd(&mut string));
137        }
138        write_string_rust(self.get_writer(), &mut string)
139    }
140
141    /// writes a simple string (usually short and non escaped)
142    /// This means we can skip the simd accelerated writing which is
143    /// expensive on short strings.
144    /// # Errors
145    /// if the write fails
146    #[inline]
147    fn write_simple_string(&mut self, string: &str) -> io::Result<()> {
148        self.write(br#"""#)?;
149        write_string_rust(self.get_writer(), &mut string.as_bytes())?;
150        self.write(br#"""#)
151    }
152    /// writes a simple string content  (usually short and non escaped)
153    /// This means we can skip the simd accelerated writing which is
154    /// expensive on short strings.
155    /// # Errors
156    /// if the write fails
157    #[inline]
158    fn write_simple_str_content(&mut self, string: &str) -> io::Result<()> {
159        let mut string = string.as_bytes();
160        // Legacy code to handle the remainder of the code
161        write_string_rust(self.get_writer(), &mut string)
162    }
163
164    /// writes a float value
165    /// # Errors
166    /// if the write fails
167    #[inline]
168    fn write_float(&mut self, num: f64) -> io::Result<()> {
169        // JSON has no NaN or infinity; write `null` like serde_json does.
170        if !num.is_finite() {
171            return self.write(b"null");
172        }
173        let mut buffer = ryu::Buffer::new();
174        let s = buffer.format_finite(num);
175        self.get_writer().write_all(s.as_bytes())
176    }
177
178    /// writes an integer value
179    /// # Errors
180    /// if the write fails
181    #[inline]
182    fn write_int<I: itoa::Integer>(&mut self, num: I) -> io::Result<()> {
183        let mut buffer = itoa::Buffer::new();
184        let s = buffer.format(num);
185        self.get_writer().write_all(s.as_bytes())
186    }
187
188    /// # Safety
189    /// This function is unsafe because it may use simd instructions
190    /// # Errors
191    ///  if the write fails
192    #[cfg(all(
193        feature = "runtime-detection",
194        any(target_arch = "x86_64", target_arch = "x86"),
195    ))]
196    unsafe fn write_str_simd(&mut self, string: &mut &[u8]) -> io::Result<()> {
197        unsafe { write_str_simd_fastest(self.get_writer(), string) }
198    }
199    #[cfg(all(target_feature = "avx2", not(feature = "runtime-detection")))]
200    #[inline]
201    /// Writes a string with simd-acceleration
202    /// # Safety
203    /// This function is unsafe because it uses simd instructions
204    /// # Errors
205    ///  if the write fails
206    unsafe fn write_str_simd(&mut self, string: &mut &[u8]) -> io::Result<()> {
207        write_str_simd_avx2(self.get_writer(), string)
208    }
209
210    #[cfg(all(
211        target_feature = "sse2",
212        not(target_feature = "avx2"),
213        not(feature = "runtime-detection")
214    ))]
215    #[inline]
216    /// Writes a string with simd-acceleration
217    /// # Safety
218    /// This function is unsafe because it uses simd instructions
219    /// # Errors
220    ///  if the write fails
221    unsafe fn write_str_simd(&mut self, string: &mut &[u8]) -> io::Result<()> {
222        write_str_simd_sse2(self.get_writer(), string)
223    }
224
225    #[cfg(not(any(
226        all(
227            feature = "runtime-detection",
228            any(target_arch = "x86_64", target_arch = "x86")
229        ),
230        feature = "portable",
231        target_feature = "avx2",
232        target_feature = "sse2",
233        target_feature = "simd128",
234        target_arch = "aarch64",
235    )))]
236    #[inline]
237    /// Writes a string with simd-acceleration (not really, as the architecture doesn't support it)
238    /// # Safety
239    /// This function is unsafe because it uses simd instructions
240    /// # Errors
241    ///  if the write fails
242    unsafe fn write_str_simd(&mut self, string: &mut &[u8]) -> io::Result<()> {
243        write_string_rust(self.get_writer(), string)
244    }
245
246    #[cfg(target_arch = "aarch64")]
247    #[inline]
248    /// Writes a string with simd-acceleration
249    /// # Safety
250    /// This function is unsafe because it uses simd instructions
251    /// # Errors
252    ///  if the write fails
253    unsafe fn write_str_simd(&mut self, string: &mut &[u8]) -> io::Result<()> {
254        use std::arch::aarch64::{
255            uint8x16_t, vandq_u8, vceqq_u8, vdupq_n_u8, veorq_u8, vgetq_lane_u16, vld1q_u8,
256            vorrq_u8, vpaddq_u8, vreinterpretq_u16_u8,
257        };
258        use std::mem;
259
260        #[inline]
261        unsafe fn bit_mask() -> uint8x16_t {
262            unsafe {
263                mem::transmute([
264                    0x01_u8, 0x02, 0x4, 0x8, 0x10, 0x20, 0x40, 0x80, 0x01, 0x02, 0x4, 0x8, 0x10,
265                    0x20, 0x40, 0x80,
266                ])
267            }
268        }
269
270        #[inline]
271        unsafe fn neon_movemask(input: uint8x16_t) -> u16 {
272            unsafe {
273                let simd_input: uint8x16_t = vandq_u8(input, bit_mask());
274                let tmp: uint8x16_t = vpaddq_u8(simd_input, simd_input);
275                let tmp = vpaddq_u8(tmp, tmp);
276                let tmp = vpaddq_u8(tmp, tmp);
277
278                vgetq_lane_u16(vreinterpretq_u16_u8(tmp), 0)
279            }
280        }
281
282        let writer = self.get_writer();
283        // The case where we have a 16+ byte block
284        // we repeate the same logic as above but with
285        // only 16 bytes
286        let mut idx = 0;
287        unsafe {
288            let zero = vdupq_n_u8(0);
289            let lower_quote_range = vdupq_n_u8(0x1F);
290            let quote = vdupq_n_u8(b'"');
291            let backslash = vdupq_n_u8(b'\\');
292            while string.len() - idx > 16 {
293                // Load 16 bytes of data;
294                let data: uint8x16_t = vld1q_u8(string.as_ptr().add(idx));
295                // Test the data against being backslash and quote.
296                let bs_or_quote = vorrq_u8(vceqq_u8(data, backslash), vceqq_u8(data, quote));
297                // Now mask the data with the quote range (0x1F).
298                let in_quote_range = vandq_u8(data, lower_quote_range);
299                // then test of the data is unchanged. aka: xor it with the
300                // Any field that was inside the quote range it will be zero
301                // now.
302                let is_unchanged = veorq_u8(data, in_quote_range);
303                let in_range = vceqq_u8(is_unchanged, zero);
304                let quote_bits = neon_movemask(vorrq_u8(bs_or_quote, in_range));
305                if quote_bits == 0 {
306                    idx += 16;
307                } else {
308                    let quote_dist = quote_bits.trailing_zeros() as usize;
309                    stry!(writer.write_all(&string[0..idx + quote_dist]));
310                    let ch = string[idx + quote_dist];
311                    match ESCAPED[ch as usize] {
312                        b'u' => stry!(u_encode(writer, ch)),
313                        escape => stry!(writer.write_all(&[b'\\', escape])),
314                    }
315
316                    *string = &string[idx + quote_dist + 1..];
317                    idx = 0;
318                }
319            }
320        }
321        stry!(writer.write_all(&string[0..idx]));
322        *string = &string[idx..];
323        Ok(())
324    }
325
326    #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
327    #[inline]
328    /// Writes a string with simd-acceleration
329    /// # Safety
330    /// This function is unsafe because it uses simd instructions
331    /// # Errors
332    ///  if the write fails
333    unsafe fn write_str_simd(&mut self, string: &mut &[u8]) -> io::Result<()> {
334        let writer = self.get_writer();
335        use std::arch::wasm32::{
336            u8x16_bitmask, u8x16_eq, u8x16_splat, v128, v128_and, v128_load, v128_or, v128_xor,
337        };
338
339        // The case where we have a 16+ byte block
340        // we repeat the same logic as above but with
341        // only 16 bytes
342        let mut idx = 0;
343        let zero = u8x16_splat(0);
344        let lower_quote_range = u8x16_splat(0x1F);
345        let quote = u8x16_splat(b'"');
346        let backslash = u8x16_splat(b'\\');
347        while string.len() - idx > 16 {
348            // Load 16 bytes of data;
349            let data = v128_load(string.as_ptr().add(idx).cast::<v128>());
350            // Test the data against being backslash and quote.
351            let bs_or_quote = v128_or(u8x16_eq(data, backslash), u8x16_eq(data, quote));
352            // Now mask the data with the quote range (0x1F).
353            let in_quote_range = v128_and(data, lower_quote_range);
354            // then test of the data is unchanged. aka: xor it with the
355            // Any field that was inside the quote range it will be zero
356            // now.
357            let is_unchanged = v128_xor(data, in_quote_range);
358            let in_range = u8x16_eq(is_unchanged, zero);
359            let quote_bits = u8x16_bitmask(v128_or(bs_or_quote, in_range));
360            if quote_bits == 0 {
361                idx += 16;
362            } else {
363                let quote_dist = quote_bits.trailing_zeros() as usize;
364                stry!(writer.write_all(&string[0..idx + quote_dist]));
365                let ch = string[idx + quote_dist];
366                match ESCAPED[ch as usize] {
367                    b'u' => stry!(u_encode(writer, ch)),
368                    escape => stry!(writer.write_all(&[b'\\', escape])),
369                }
370
371                *string = &string[idx + quote_dist + 1..];
372                idx = 0;
373            }
374        }
375        stry!(writer.write_all(&string[0..idx]));
376        *string = &string[idx..];
377        Ok(())
378    }
379}
380
381#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
382unsafe fn write_str_simd_fastest<W>(writer: &mut W, string: &mut &[u8]) -> io::Result<()>
383where
384    W: Write,
385{
386    // This is not possible right now since we can't get `W` to be part of the static
387    // use std::sync::atomic::{AtomicPtr, Ordering};
388    // type FnRaw = *mut ();
389
390    // type WriteStrFn<T> = for<'a, 'b, 'c> unsafe fn(&'a mut T, &'b mut &'c [u8]) -> io::Result<()>;
391    // static FN: AtomicPtr<()> = AtomicPtr::new(get_fastest::<W> as FnRaw);
392    // #[inline]
393    // fn get_fastest_available_implementation<W>() -> WriteStrFn<W>
394    // where
395    //     W: Write,
396    // {
397    //     if std::is_x86_feature_detected!("avx2") {
398    //         write_str_simd_avx2
399    //     } else if std::is_x86_feature_detected!("sse4.2") {
400    //         write_str_simd_sse2
401    //     } else {
402    //         write_str_simd_rust
403    //     }
404    // }
405
406    // #[inline]
407    // unsafe fn get_fastest<'invoke, 'de, W>(writer: &mut W, string: &mut &[u8]) -> io::Result<()>
408    // where
409    //     W: Write,
410    //     'de: 'invoke,
411    // {
412    //     let fun = get_fastest_available_implementation();
413    //     FN.store(fun as FnRaw, Ordering::Relaxed);
414    //     (fun)(writer, string)
415    // }
416    // let fun = FN.load(Ordering::Relaxed);
417    // mem::transmute::<FnRaw, WriteStrFn>(fun)(writer, string)
418
419    if std::is_x86_feature_detected!("avx2") {
420        unsafe { write_str_simd_avx2(writer, string) }
421    } else if std::is_x86_feature_detected!("sse2") {
422        unsafe { write_str_simd_sse2(writer, string) }
423    } else {
424        #[cfg(not(feature = "portable"))]
425        return write_string_rust(writer, string);
426        #[cfg(feature = "portable")]
427        return write_str_simd_portable(writer, string);
428    }
429}
430#[inline]
431fn write_string_container<W>(writer: &mut W, string: &[u8], mut start: usize) -> io::Result<()>
432where
433    W: Write,
434{
435    stry!(writer.write_all(&string[..start]));
436
437    for (index, ch) in string.iter().enumerate().skip(start) {
438        let escape = ESCAPED[*ch as usize];
439        if escape > 0 {
440            stry!(writer.write_all(&string[start..index]));
441            if escape == b'u' {
442                stry!(u_encode(writer, *ch));
443            } else {
444                stry!(writer.write_all(&[b'\\', escape]));
445            }
446            start = index + 1;
447        }
448    }
449    writer.write_all(&string[start..])
450}
451
452#[inline]
453fn write_string_rust<W>(writer: &mut W, string: &mut &[u8]) -> io::Result<()>
454where
455    W: Write,
456{
457    // Legacy code to handle the remainder of the code
458    for (index, ch) in string.iter().enumerate() {
459        if ESCAPED[*ch as usize] > 0 {
460            write_string_container(writer, string, index)?;
461            *string = &string[string.len()..];
462            return Ok(());
463        }
464    }
465    writer.write_all(string)?;
466    *string = &string[string.len()..];
467    Ok(())
468}
469
470#[cfg(feature = "portable")]
471#[inline]
472/// Writes a string with simd-acceleration
473/// # Safety
474/// This function is unsafe because it uses simd instructions
475/// # Errors
476///  if the write fails
477unsafe fn write_str_simd_portable<W>(writer: &mut W, string: &mut &[u8]) -> io::Result<()>
478where
479    W: Write,
480{
481    use std::simd::{SimdPartialEq, ToBitMask, u8x32};
482
483    let mut idx = 0;
484    let zero = u8x32::splat(0);
485    let lower_quote_range = u8x32::splat(0x1F_u8);
486    let quote = u8x32::splat(b'"');
487    let backslash = u8x32::splat(b'\\');
488    while string.len() - idx >= 32 {
489        // Load 32 bytes of data;
490        let data = u8x32::from_slice(&string[idx..]);
491        // Test the data against being backslash and quote.
492        let bs_or_quote = data.simd_eq(backslash) | data.simd_eq(quote);
493        // Now mask the data with the quote range (0x1F).
494        let in_quote_range = data & lower_quote_range;
495        // then test of the data is unchanged. aka: xor it with the
496        // Any field that was inside the quote range it will be zero
497        // now.
498        let is_unchanged = data ^ in_quote_range;
499        let in_range = is_unchanged.simd_eq(zero);
500        let quote_bits = (bs_or_quote | in_range).to_bitmask();
501        if quote_bits == 0 {
502            idx += 32;
503        } else {
504            let quote_dist = quote_bits.trailing_zeros() as usize;
505            stry!(writer.write_all(string.get_unchecked(0..idx + quote_dist)));
506
507            let ch = string[idx + quote_dist];
508            match ESCAPED[ch as usize] {
509                b'u' => stry!(u_encode(writer, ch)),
510                escape => stry!(writer.write_all(&[b'\\', escape])),
511            };
512
513            *string = string.get_unchecked(idx + quote_dist + 1..);
514            idx = 0;
515        }
516    }
517    stry!(writer.write_all(&string[0..idx]));
518    *string = string.get_unchecked(idx..);
519    Ok(())
520}
521
522#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
523#[target_feature(enable = "avx2")]
524#[inline]
525/// Writes a string with simd-acceleration
526/// # Safety
527/// This function is unsafe because it uses simd instructions
528/// # Errors
529///  if the write fails
530unsafe fn write_str_simd_avx2<W>(writer: &mut W, string: &mut &[u8]) -> io::Result<()>
531where
532    W: Write,
533{
534    #[cfg(target_arch = "x86")]
535    use std::arch::x86::{
536        __m256i, _mm256_and_si256, _mm256_cmpeq_epi8, _mm256_loadu_si256, _mm256_movemask_epi8,
537        _mm256_or_si256, _mm256_set1_epi8, _mm256_xor_si256,
538    };
539    #[cfg(target_arch = "x86_64")]
540    use std::arch::x86_64::{
541        __m256i, _mm256_and_si256, _mm256_cmpeq_epi8, _mm256_loadu_si256, _mm256_movemask_epi8,
542        _mm256_or_si256, _mm256_set1_epi8, _mm256_xor_si256,
543    };
544
545    unsafe {
546        let mut idx = 0;
547        let zero = _mm256_set1_epi8(0);
548        let lower_quote_range = _mm256_set1_epi8(0x1F_i8);
549        #[allow(clippy::cast_possible_wrap)] // it's a const, it's fine
550        let quote = _mm256_set1_epi8(b'"' as i8);
551        #[allow(clippy::cast_possible_wrap)] // it's a const, it's fine
552        let backslash = _mm256_set1_epi8(b'\\' as i8);
553        while string.len() - idx >= 32 {
554            // Load 32 bytes of data; _mm256_loadu_si256 does not require alignment
555            #[allow(clippy::cast_ptr_alignment)]
556            let data: __m256i = _mm256_loadu_si256(string.as_ptr().add(idx).cast::<__m256i>());
557            // Test the data against being backslash and quote.
558            let bs_or_quote = _mm256_or_si256(
559                _mm256_cmpeq_epi8(data, backslash),
560                _mm256_cmpeq_epi8(data, quote),
561            );
562            // Now mask the data with the quote range (0x1F).
563            let in_quote_range = _mm256_and_si256(data, lower_quote_range);
564            // then test of the data is unchanged. aka: xor it with the
565            // Any field that was inside the quote range it will be zero
566            // now.
567            let is_unchanged = _mm256_xor_si256(data, in_quote_range);
568            let in_range = _mm256_cmpeq_epi8(is_unchanged, zero);
569            let quote_bits = _mm256_movemask_epi8(_mm256_or_si256(bs_or_quote, in_range));
570            if quote_bits == 0 {
571                idx += 32;
572            } else {
573                let quote_dist = quote_bits.trailing_zeros() as usize;
574                stry!(writer.write_all(string.get_unchecked(0..idx + quote_dist)));
575
576                let ch = string[idx + quote_dist];
577                match ESCAPED[ch as usize] {
578                    b'u' => stry!(u_encode(writer, ch)),
579                    escape => stry!(writer.write_all(&[b'\\', escape])),
580                }
581
582                *string = string.get_unchecked(idx + quote_dist + 1..);
583                idx = 0;
584            }
585        }
586        stry!(writer.write_all(&string[0..idx]));
587        *string = string.get_unchecked(idx..);
588    }
589
590    Ok(())
591}
592
593#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
594#[target_feature(enable = "sse2")]
595#[inline]
596/// Writes a string with simd-acceleration
597/// # Safety
598/// This function is unsafe because it uses simd instructions
599/// # Errors
600///  if the write fails
601unsafe fn write_str_simd_sse2<W>(writer: &mut W, string: &mut &[u8]) -> io::Result<()>
602where
603    W: Write,
604{
605    #[cfg(target_arch = "x86")]
606    use std::arch::x86::{
607        __m128i, _mm_and_si128, _mm_cmpeq_epi8, _mm_loadu_si128, _mm_movemask_epi8, _mm_or_si128,
608        _mm_set1_epi8, _mm_xor_si128,
609    };
610    #[cfg(target_arch = "x86_64")]
611    use std::arch::x86_64::{
612        __m128i, _mm_and_si128, _mm_cmpeq_epi8, _mm_loadu_si128, _mm_movemask_epi8, _mm_or_si128,
613        _mm_set1_epi8, _mm_xor_si128,
614    };
615
616    unsafe {
617        let mut idx = 0;
618        let zero = _mm_set1_epi8(0);
619        let lower_quote_range = _mm_set1_epi8(0x1F_i8);
620        #[allow(clippy::cast_possible_wrap)] // it's a const, it's fine
621        let quote = _mm_set1_epi8(b'"' as i8);
622        #[allow(clippy::cast_possible_wrap)] // it's a const, it's fine
623        let backslash = _mm_set1_epi8(b'\\' as i8);
624        while string.len() - idx > 16 {
625            // Load 16 bytes of data; _mm_loadu_si128 does not require alignment
626            #[allow(clippy::cast_ptr_alignment)]
627            let data: __m128i = _mm_loadu_si128(string.as_ptr().add(idx).cast::<__m128i>());
628            // Test the data against being backslash and quote.
629            let bs_or_quote =
630                _mm_or_si128(_mm_cmpeq_epi8(data, backslash), _mm_cmpeq_epi8(data, quote));
631            // Now mask the data with the quote range (0x1F).
632            let in_quote_range = _mm_and_si128(data, lower_quote_range);
633            // then test of the data is unchanged. aka: xor it with the
634            // Any field that was inside the quote range it will be zero
635            // now.
636            let is_unchanged = _mm_xor_si128(data, in_quote_range);
637            let in_range = _mm_cmpeq_epi8(is_unchanged, zero);
638            let quote_bits = _mm_movemask_epi8(_mm_or_si128(bs_or_quote, in_range));
639            if quote_bits == 0 {
640                idx += 16;
641            } else {
642                let quote_dist = quote_bits.trailing_zeros() as usize;
643                stry!(writer.write_all(&string[0..idx + quote_dist]));
644
645                let ch = string[idx + quote_dist];
646                match ESCAPED[ch as usize] {
647                    b'u' => stry!(u_encode(writer, ch)),
648                    escape => stry!(writer.write_all(&[b'\\', escape])),
649                }
650
651                *string = &string[idx + quote_dist + 1..];
652                idx = 0;
653            }
654        }
655        stry!(writer.write_all(&string[0..idx]));
656        *string = &string[idx..];
657    }
658
659    Ok(())
660}
661
662///  Simple dump Generator
663pub struct DumpGenerator {
664    code: Vec<u8>,
665}
666
667impl Default for DumpGenerator {
668    fn default() -> Self {
669        Self {
670            code: Vec::with_capacity(1024),
671        }
672    }
673}
674
675impl DumpGenerator {
676    /// Creates a new generator
677    #[must_use]
678    pub fn new() -> Self {
679        Self::default()
680    }
681
682    /// Returns the data as a String
683    #[must_use]
684    pub fn consume(self) -> String {
685        // Original strings were unicode, numbers are all ASCII,
686        // therefore this is safe.
687        unsafe { String::from_utf8_unchecked(self.code) }
688    }
689}
690
691impl BaseGenerator for DumpGenerator {
692    type T = Vec<u8>;
693
694    #[inline]
695    fn write(&mut self, slice: &[u8]) -> io::Result<()> {
696        extend_from_slice(&mut self.code, slice);
697        Ok(())
698    }
699    #[inline]
700    fn write_char(&mut self, ch: u8) -> io::Result<()> {
701        self.code.push(ch);
702        Ok(())
703    }
704
705    #[inline]
706    fn get_writer(&mut self) -> &mut Vec<u8> {
707        &mut self.code
708    }
709
710    #[inline]
711    fn write_min(&mut self, _: &[u8], min: u8) -> io::Result<()> {
712        self.code.push(min);
713        Ok(())
714    }
715}
716
717/// Pretty Generator
718pub struct PrettyGenerator {
719    code: Vec<u8>,
720    dent: u16,
721    spaces_per_indent: u16,
722}
723
724impl PrettyGenerator {
725    /// Creates a new pretty printing generator
726    #[must_use]
727    pub fn new(spaces: u16) -> Self {
728        Self {
729            code: Vec::with_capacity(1024),
730            dent: 0,
731            spaces_per_indent: spaces,
732        }
733    }
734
735    /// Returns the data as a String
736    #[must_use]
737    pub fn consume(self) -> String {
738        unsafe { String::from_utf8_unchecked(self.code) }
739    }
740}
741
742impl BaseGenerator for PrettyGenerator {
743    type T = Vec<u8>;
744    #[inline]
745    fn write(&mut self, slice: &[u8]) -> io::Result<()> {
746        extend_from_slice(&mut self.code, slice);
747        Ok(())
748    }
749
750    #[inline]
751    fn write_char(&mut self, ch: u8) -> io::Result<()> {
752        self.code.push(ch);
753        Ok(())
754    }
755
756    #[inline]
757    fn get_writer(&mut self) -> &mut Vec<u8> {
758        &mut self.code
759    }
760
761    #[inline]
762    fn write_min(&mut self, slice: &[u8], _: u8) -> io::Result<()> {
763        extend_from_slice(&mut self.code, slice);
764        Ok(())
765    }
766
767    fn new_line(&mut self) -> io::Result<()> {
768        self.code.push(b'\n');
769        self.code.resize(
770            self.code.len() + (self.dent * self.spaces_per_indent) as usize,
771            b' ',
772        );
773        Ok(())
774    }
775
776    fn indent(&mut self) {
777        self.dent += 1;
778    }
779
780    fn dedent(&mut self) {
781        self.dent -= 1;
782    }
783}
784
785/// Writer Generator
786pub struct WriterGenerator<'w, W: 'w + Write> {
787    writer: &'w mut W,
788}
789
790impl<'w, W> WriterGenerator<'w, W>
791where
792    W: 'w + Write,
793{
794    /// Creates a new generator
795    pub fn new(writer: &'w mut W) -> Self {
796        WriterGenerator { writer }
797    }
798}
799
800impl<W> BaseGenerator for WriterGenerator<'_, W>
801where
802    W: Write,
803{
804    type T = W;
805
806    #[inline]
807    fn get_writer(&mut self) -> &mut W {
808        self.writer
809    }
810
811    #[inline]
812    fn write_min(&mut self, _: &[u8], min: u8) -> io::Result<()> {
813        self.writer.write_all(&[min])
814    }
815}
816
817/// Pretty Writer Generator
818pub struct PrettyWriterGenerator<'w, W>
819where
820    W: 'w + Write,
821{
822    writer: &'w mut W,
823    dent: u16,
824    spaces_per_indent: u16,
825}
826
827impl<'w, W> PrettyWriterGenerator<'w, W>
828where
829    W: 'w + Write,
830{
831    /// Creates a new generator
832    pub fn new(writer: &'w mut W, spaces_per_indent: u16) -> Self {
833        PrettyWriterGenerator {
834            writer,
835            dent: 0,
836            spaces_per_indent,
837        }
838    }
839}
840
841impl<W> BaseGenerator for PrettyWriterGenerator<'_, W>
842where
843    W: Write,
844{
845    type T = W;
846
847    #[inline]
848    fn get_writer(&mut self) -> &mut W {
849        self.writer
850    }
851
852    #[inline]
853    fn write_min(&mut self, slice: &[u8], _: u8) -> io::Result<()> {
854        self.writer.write_all(slice)
855    }
856
857    fn new_line(&mut self) -> io::Result<()> {
858        stry!(self.write_char(b'\n'));
859        for _ in 0..(self.dent * self.spaces_per_indent) {
860            stry!(self.write_char(b' '));
861        }
862        Ok(())
863    }
864
865    fn indent(&mut self) {
866        self.dent += 1;
867    }
868
869    fn dedent(&mut self) {
870        self.dent -= 1;
871    }
872}
873
874// From: https://github.com/dtolnay/fastwrite/blob/master/src/lib.rs#L68
875//
876// LLVM is not able to lower `Vec::extend_from_slice` into a memcpy, so this
877// helps eke out that last bit of performance.
878#[inline]
879pub(crate) fn extend_from_slice(dst: &mut Vec<u8>, src: &[u8]) {
880    let dst_len = dst.len();
881    let src_len = src.len();
882
883    dst.reserve(src_len);
884
885    unsafe {
886        // We would have failed if `reserve` overflowed\
887        ptr::copy_nonoverlapping(src.as_ptr(), dst.as_mut_ptr().add(dst_len), src_len);
888        dst.set_len(dst_len + src_len);
889    }
890}
891
892#[cfg(test)]
893mod tests {
894
895    #[test]
896    fn test_write_string_rust() {
897        let mut writer = Vec::new();
898        let mut string = "Hello, World!".as_bytes();
899        super::write_string_rust(&mut writer, &mut string).expect("failed to write string");
900        assert_eq!(writer, "Hello, World!".as_bytes());
901        assert!(string.is_empty(), "write_string_rust should advance string");
902    }
903    #[test]
904    fn test_write_string_rust2() {
905        let mut writer = Vec::new();
906        let mut string = "Hello, \"World!\"".as_bytes();
907        super::write_string_rust(&mut writer, &mut string).expect("failed to write string");
908        assert_eq!(writer, "Hello, \\\"World!\\\"".as_bytes());
909        assert!(string.is_empty(), "write_string_rust should advance string");
910    }
911
912    #[test]
913    fn test_write_string_long_with_escapes() {
914        use super::{BaseGenerator, DumpGenerator};
915        // Long string with escapes past the 32-byte SIMD boundary exercises
916        // write_string_content -> write_str_simd -> write_string_rust and
917        // catches the double-write bug if the fallback doesn't advance the string.
918        let input = "abcdefghijklmnopqrstuvwxyz012345\"escape here\"";
919        let mut g = DumpGenerator::new();
920        g.write_string(input).expect("write failed");
921        assert_eq!(
922            g.consume(),
923            "\"abcdefghijklmnopqrstuvwxyz012345\\\"escape here\\\"\""
924        );
925    }
926    #[test]
927    fn test_write_float_non_finite_is_null() {
928        use super::BaseGenerator;
929        for f in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
930            let mut g = super::DumpGenerator::new();
931            g.write_float(f).expect("failed to write float");
932            assert_eq!(g.consume(), "null", "{f}");
933        }
934    }
935}
value_trait/generator.rs

value_trait/
generator.rs
