Skip to main content

json_steroids/
writer.rs

1//! High-performance JSON writer
2//!
3//! Optimized for fast serialization with:
4//! - Pre-allocated buffer
5//! - Efficient string escaping
6//! - Optional pretty printing
7
8/// Lookup table: 0 = safe byte, 1 = needs escaping
9/// Covers bytes 0x00-0xFF. Bytes < 0x20, '"', and '\\' need escaping.
10static NEEDS_ESCAPE: [bool; 256] = {
11    let mut table = [false; 256];
12    let mut i = 0u8;
13    // Control characters 0x00-0x1F
14    loop {
15        table[i as usize] = true;
16        if i == 0x1F {
17            break;
18        }
19        i += 1;
20    }
21    table[b'"' as usize] = true;
22    table[b'\\' as usize] = true;
23    table
24};
25
26/// Internal writer trait for different formatting strategies
27pub trait Writer {
28    fn buffer(&self) -> &Vec<u8>;
29    fn buffer_mut(&mut self) -> &mut Vec<u8>;
30    fn into_buffer(self) -> Vec<u8>;
31
32    fn begin_object(&mut self);
33    fn end_object(&mut self);
34    fn begin_array(&mut self);
35    fn end_array(&mut self);
36    fn write_comma(&mut self);
37    fn write_key(&mut self, key: &str);
38    fn write_unescape_key(&mut self, s: &str);
39    fn write_string(&mut self, s: &str);
40    fn write_raw(&mut self, s: &str);
41    fn write_null(&mut self);
42    fn write_bool(&mut self, value: bool);
43    fn write_i8(&mut self, value: i8);
44    fn write_i16(&mut self, value: i16);
45    fn write_i32(&mut self, value: i32);
46    fn write_i64(&mut self, value: i64);
47    fn write_isize(&mut self, value: isize);
48    fn write_u8(&mut self, value: u8);
49    fn write_u16(&mut self, value: u16);
50    fn write_u32(&mut self, value: u32);
51    fn write_u64(&mut self, value: u64);
52    fn write_usize(&mut self, value: usize);
53    fn write_f32(&mut self, value: f32);
54    fn write_f64(&mut self, value: f64);
55}
56
57/// Compact JSON writer (no indentation)
58pub struct CompactWriter {
59    buffer: Vec<u8>,
60    itoa_buffer: itoa::Buffer,
61    ryu_buffer: ryu::Buffer,
62}
63
64/// Pretty-printed JSON writer (with indentation)
65pub struct PrettyWriter {
66    buffer: Vec<u8>,
67    indent: Vec<u8>,
68    depth: usize,
69    needs_newline: bool,
70}
71
72/// JSON writer with optimized string building
73pub struct JsonWriter<W: Writer = CompactWriter> {
74    inner: W,
75}
76
77impl CompactWriter {
78    #[inline]
79    pub fn new(capacity: usize) -> Self {
80        Self {
81            buffer: Vec::with_capacity(capacity),
82            itoa_buffer: itoa::Buffer::new(),
83            ryu_buffer: ryu::Buffer::new(),
84        }
85    }
86}
87
88impl Writer for CompactWriter {
89    #[inline]
90    fn buffer(&self) -> &Vec<u8> {
91        &self.buffer
92    }
93
94    #[inline]
95    fn buffer_mut(&mut self) -> &mut Vec<u8> {
96        &mut self.buffer
97    }
98
99    #[inline]
100    fn into_buffer(self) -> Vec<u8> {
101        self.buffer
102    }
103
104    #[inline]
105    fn begin_object(&mut self) {
106        self.buffer.push(b'{');
107    }
108
109    #[inline]
110    fn end_object(&mut self) {
111        self.buffer.push(b'}');
112    }
113
114    #[inline]
115    fn begin_array(&mut self) {
116        self.buffer.push(b'[');
117    }
118
119    #[inline]
120    fn end_array(&mut self) {
121        self.buffer.push(b']');
122    }
123
124    #[inline]
125    fn write_comma(&mut self) {
126        self.buffer.push(b',');
127    }
128
129    #[inline]
130    fn write_key(&mut self, key: &str) {
131        self.buffer.push(b'"');
132        write_escaped_string(&mut self.buffer, key);
133        self.buffer.extend_from_slice(b"\":");
134    }
135
136    #[inline]
137    fn write_unescape_key(&mut self, s: &str) {
138        self.buffer.push(b'"');
139        self.buffer.extend_from_slice(s.as_bytes());
140        self.buffer.extend_from_slice(b"\":");
141    }
142
143    #[inline]
144    fn write_string(&mut self, s: &str) {
145        self.buffer.push(b'"');
146        write_escaped_string(&mut self.buffer, s);
147        self.buffer.push(b'"');
148    }
149
150    #[inline]
151    fn write_raw(&mut self, s: &str) {
152        self.buffer.extend_from_slice(s.as_bytes());
153    }
154
155    #[inline]
156    fn write_null(&mut self) {
157        self.buffer.extend_from_slice(b"null");
158    }
159
160    #[inline]
161    fn write_bool(&mut self, value: bool) {
162        if value {
163            self.buffer.extend_from_slice(b"true");
164        } else {
165            self.buffer.extend_from_slice(b"false");
166        }
167    }
168
169    #[inline]
170    fn write_i8(&mut self, value: i8) {
171        self.buffer
172            .extend_from_slice(self.itoa_buffer.format(value).as_bytes());
173    }
174
175    #[inline]
176    fn write_i16(&mut self, value: i16) {
177        self.buffer
178            .extend_from_slice(self.itoa_buffer.format(value).as_bytes());
179    }
180
181    #[inline]
182    fn write_i32(&mut self, value: i32) {
183        self.buffer
184            .extend_from_slice(self.itoa_buffer.format(value).as_bytes());
185    }
186
187    #[inline]
188    fn write_i64(&mut self, value: i64) {
189        self.buffer
190            .extend_from_slice(self.itoa_buffer.format(value).as_bytes());
191    }
192
193    #[inline]
194    fn write_isize(&mut self, value: isize) {
195        self.buffer
196            .extend_from_slice(self.itoa_buffer.format(value).as_bytes());
197    }
198
199    #[inline]
200    fn write_u8(&mut self, value: u8) {
201        self.buffer
202            .extend_from_slice(self.itoa_buffer.format(value).as_bytes());
203    }
204
205    #[inline]
206    fn write_u16(&mut self, value: u16) {
207        self.buffer
208            .extend_from_slice(self.itoa_buffer.format(value).as_bytes());
209    }
210
211    #[inline]
212    fn write_u32(&mut self, value: u32) {
213        self.buffer
214            .extend_from_slice(self.itoa_buffer.format(value).as_bytes());
215    }
216
217    #[inline]
218    fn write_u64(&mut self, value: u64) {
219        self.buffer
220            .extend_from_slice(self.itoa_buffer.format(value).as_bytes());
221    }
222
223    #[inline]
224    fn write_usize(&mut self, value: usize) {
225        self.buffer
226            .extend_from_slice(self.itoa_buffer.format(value).as_bytes());
227    }
228
229    #[inline]
230    fn write_f32(&mut self, value: f32) {
231        self.buffer
232            .extend_from_slice(self.ryu_buffer.format(value).as_bytes());
233    }
234
235    #[inline]
236    fn write_f64(&mut self, value: f64) {
237        self.buffer
238            .extend_from_slice(self.ryu_buffer.format(value).as_bytes());
239    }
240}
241
242impl PrettyWriter {
243    #[inline]
244    pub fn new(capacity: usize, spaces: usize) -> Self {
245        Self {
246            buffer: Vec::with_capacity(capacity),
247            indent: vec![b' '; spaces],
248            depth: 0,
249            needs_newline: false,
250        }
251    }
252
253    #[inline]
254    fn write_indent(&mut self) {
255        if self.needs_newline {
256            self.buffer.push(b'\n');
257            for _ in 0..self.depth {
258                self.buffer.extend_from_slice(&self.indent);
259            }
260            self.needs_newline = false;
261        }
262    }
263}
264
265impl Writer for PrettyWriter {
266    #[inline]
267    fn buffer(&self) -> &Vec<u8> {
268        &self.buffer
269    }
270
271    #[inline]
272    fn buffer_mut(&mut self) -> &mut Vec<u8> {
273        &mut self.buffer
274    }
275
276    #[inline]
277    fn into_buffer(self) -> Vec<u8> {
278        self.buffer
279    }
280
281    #[inline]
282    fn begin_object(&mut self) {
283        self.write_indent();
284        self.buffer.push(b'{');
285        self.depth += 1;
286        self.needs_newline = true;
287    }
288
289    #[inline]
290    fn end_object(&mut self) {
291        self.depth -= 1;
292        self.needs_newline = true;
293        self.write_indent();
294        self.buffer.push(b'}');
295    }
296
297    #[inline]
298    fn begin_array(&mut self) {
299        self.write_indent();
300        self.buffer.push(b'[');
301        self.depth += 1;
302        self.needs_newline = true;
303    }
304
305    #[inline]
306    fn end_array(&mut self) {
307        self.depth -= 1;
308        self.needs_newline = true;
309        self.write_indent();
310        self.buffer.push(b']');
311    }
312
313    #[inline]
314    fn write_comma(&mut self) {
315        self.buffer.push(b',');
316        self.needs_newline = true;
317    }
318
319    #[inline]
320    fn write_key(&mut self, key: &str) {
321        self.write_indent();
322        self.buffer.push(b'"');
323        write_escaped_string(&mut self.buffer, key);
324        self.buffer.extend_from_slice(b"\": ");
325    }
326
327    #[inline]
328    fn write_unescape_key(&mut self, key: &str) {
329        self.write_indent();
330        self.buffer.push(b'"');
331        self.buffer.extend_from_slice(key.as_bytes());
332        self.buffer.extend_from_slice(b"\": ");
333    }
334
335    #[inline]
336    fn write_string(&mut self, s: &str) {
337        self.write_indent();
338        self.buffer.push(b'"');
339        write_escaped_string(&mut self.buffer, s);
340        self.buffer.push(b'"');
341    }
342
343    #[inline]
344    fn write_raw(&mut self, s: &str) {
345        self.write_indent();
346        self.buffer.extend_from_slice(s.as_bytes());
347    }
348
349    #[inline]
350    fn write_null(&mut self) {
351        self.write_indent();
352        self.buffer.extend_from_slice(b"null");
353    }
354
355    #[inline]
356    fn write_bool(&mut self, value: bool) {
357        self.write_indent();
358        if value {
359            self.buffer.extend_from_slice(b"true");
360        } else {
361            self.buffer.extend_from_slice(b"false");
362        }
363    }
364
365    #[inline]
366    fn write_i8(&mut self, value: i8) {
367        self.write_indent();
368        let mut buffer = itoa::Buffer::new();
369        self.buffer
370            .extend_from_slice(buffer.format(value).as_bytes());
371    }
372
373    #[inline]
374    fn write_i16(&mut self, value: i16) {
375        self.write_indent();
376        let mut buffer = itoa::Buffer::new();
377        self.buffer
378            .extend_from_slice(buffer.format(value).as_bytes());
379    }
380
381    #[inline]
382    fn write_i32(&mut self, value: i32) {
383        self.write_indent();
384        let mut buffer = itoa::Buffer::new();
385        self.buffer
386            .extend_from_slice(buffer.format(value).as_bytes());
387    }
388
389    #[inline]
390    fn write_i64(&mut self, value: i64) {
391        self.write_indent();
392        let mut buffer = itoa::Buffer::new();
393        self.buffer
394            .extend_from_slice(buffer.format(value).as_bytes());
395    }
396
397    #[inline]
398    fn write_isize(&mut self, value: isize) {
399        self.write_indent();
400        let mut buffer = itoa::Buffer::new();
401        self.buffer
402            .extend_from_slice(buffer.format(value).as_bytes());
403    }
404
405    #[inline]
406    fn write_u8(&mut self, value: u8) {
407        self.write_indent();
408        let mut buffer = itoa::Buffer::new();
409        self.buffer
410            .extend_from_slice(buffer.format(value).as_bytes());
411    }
412
413    #[inline]
414    fn write_u16(&mut self, value: u16) {
415        self.write_indent();
416        let mut buffer = itoa::Buffer::new();
417        self.buffer
418            .extend_from_slice(buffer.format(value).as_bytes());
419    }
420
421    #[inline]
422    fn write_u32(&mut self, value: u32) {
423        self.write_indent();
424        let mut buffer = itoa::Buffer::new();
425        self.buffer
426            .extend_from_slice(buffer.format(value).as_bytes());
427    }
428
429    #[inline]
430    fn write_u64(&mut self, value: u64) {
431        self.write_indent();
432        let mut buffer = itoa::Buffer::new();
433        self.buffer
434            .extend_from_slice(buffer.format(value).as_bytes());
435    }
436
437    #[inline]
438    fn write_usize(&mut self, value: usize) {
439        self.write_indent();
440        let mut buffer = itoa::Buffer::new();
441        self.buffer
442            .extend_from_slice(buffer.format(value).as_bytes());
443    }
444
445    #[inline]
446    fn write_f32(&mut self, value: f32) {
447        self.write_indent();
448        let mut buffer = ryu::Buffer::new();
449        self.buffer
450            .extend_from_slice(buffer.format(value).as_bytes());
451    }
452
453    #[inline]
454    fn write_f64(&mut self, value: f64) {
455        self.write_indent();
456        let mut buffer = ryu::Buffer::new();
457        self.buffer
458            .extend_from_slice(buffer.format(value).as_bytes());
459    }
460}
461
462/// Single-pass string escaping: copies clean byte runs in bulk,
463/// only pays the per-byte cost when an escape is actually needed.
464#[inline]
465fn write_escaped_string(buffer: &mut Vec<u8>, s: &str) {
466    let bytes = s.as_bytes();
467    let mut start = 0;
468
469    for i in 0..bytes.len() {
470        let byte = unsafe { *bytes.get_unchecked(i) };
471
472        if unsafe { *NEEDS_ESCAPE.get_unchecked(byte as usize) } {
473            // Write any accumulated clean bytes
474            if start < i {
475                buffer.extend_from_slice(&bytes[start..i]);
476            }
477
478            // Write the escape sequence
479            buffer.push(b'\\');
480            let escaped = match byte {
481                b'"' => b'"',
482                b'\\' => b'\\',
483                b'\n' => b'n',
484                b'\r' => b'r',
485                b'\t' => b't',
486                b'\x08' => b'b', // backspace
487                b'\x0C' => b'f', // form feed
488                _ => {
489                    // Unicode escape for other control characters
490                    buffer.push(b'u');
491                    buffer.push(b'0');
492                    buffer.push(b'0');
493                    let hex_digits = b"0123456789abcdef";
494                    buffer.push(hex_digits[(byte >> 4) as usize]);
495                    buffer.push(hex_digits[(byte & 0x0F) as usize]);
496                    start = i + 1;
497                    continue;
498                }
499            };
500            buffer.push(escaped);
501            start = i + 1;
502        }
503    }
504
505    // Write any remaining clean bytes
506    if start < bytes.len() {
507        buffer.extend_from_slice(&bytes[start..]);
508    }
509}
510
511impl JsonWriter<CompactWriter> {
512    /// Create a new compact JSON writer
513    #[inline]
514    pub fn new() -> Self {
515        Self {
516            inner: CompactWriter::new(2048),
517        }
518    }
519
520    /// Create a writer with pre-allocated capacity
521    #[inline]
522    pub fn with_capacity(capacity: usize) -> Self {
523        Self {
524            inner: CompactWriter::new(capacity),
525        }
526    }
527}
528
529impl JsonWriter<PrettyWriter> {
530    /// Create a new writer with the given indentation
531    #[inline]
532    pub fn with_indent(spaces: usize) -> Self {
533        Self {
534            inner: PrettyWriter::new(1024, spaces),
535        }
536    }
537}
538
539impl<W: Writer> JsonWriter<W> {
540    /// Get the result as a String
541    #[inline]
542    pub fn into_string(self) -> String {
543        let buffer = self.inner.into_buffer();
544        // Safety: we only write valid UTF-8
545        unsafe { String::from_utf8_unchecked(buffer) }
546    }
547
548    /// Get the result as bytes
549    #[inline]
550    pub fn as_bytes(&self) -> &[u8] {
551        self.inner.buffer()
552    }
553
554    /// Begin an object
555    #[inline]
556    pub fn begin_object(&mut self) {
557        self.inner.begin_object();
558    }
559
560    /// End an object
561    #[inline]
562    pub fn end_object(&mut self) {
563        self.inner.end_object();
564    }
565
566    /// Begin an array
567    #[inline]
568    pub fn begin_array(&mut self) {
569        self.inner.begin_array();
570    }
571
572    /// End an array
573    #[inline]
574    pub fn end_array(&mut self) {
575        self.inner.end_array();
576    }
577
578    /// Write a comma separator
579    #[inline]
580    pub fn write_comma(&mut self) {
581        self.inner.write_comma();
582    }
583
584    /// Write an object key
585    #[inline]
586    pub fn write_key(&mut self, key: &str) {
587        self.inner.write_key(key);
588    }
589
590    /// Write an object key
591    #[inline]
592    pub fn write_unescape_key(&mut self, key: &str) {
593        self.inner.write_unescape_key(key);
594    }
595
596    /// Write a string value with proper escaping (single-pass)
597    #[inline]
598    pub fn write_string(&mut self, s: &str) {
599        self.inner.write_string(s);
600    }
601
602    /// Write a raw string (no escaping, no quotes)
603    #[inline]
604    pub fn write_raw(&mut self, s: &str) {
605        self.inner.write_raw(s);
606    }
607
608    /// Write null
609    #[inline]
610    pub fn write_null(&mut self) {
611        self.inner.write_null();
612    }
613
614    /// Write a boolean
615    #[inline]
616    pub fn write_bool(&mut self, value: bool) {
617        self.inner.write_bool(value);
618    }
619
620    /// Write an i8
621    #[inline]
622    pub fn write_i8(&mut self, value: i8) {
623        self.inner.write_i8(value);
624    }
625
626    /// Write an i16
627    #[inline]
628    pub fn write_i16(&mut self, value: i16) {
629        self.inner.write_i16(value);
630    }
631
632    /// Write an i32
633    #[inline]
634    pub fn write_i32(&mut self, value: i32) {
635        self.inner.write_i32(value);
636    }
637
638    /// Write an i64
639    #[inline]
640    pub fn write_i64(&mut self, value: i64) {
641        self.inner.write_i64(value);
642    }
643
644    /// Write an isize
645    #[inline]
646    pub fn write_isize(&mut self, value: isize) {
647        self.inner.write_isize(value);
648    }
649
650    /// Write a u8
651    #[inline]
652    pub fn write_u8(&mut self, value: u8) {
653        self.inner.write_u8(value);
654    }
655
656    /// Write a u16
657    #[inline]
658    pub fn write_u16(&mut self, value: u16) {
659        self.inner.write_u16(value);
660    }
661
662    /// Write a u32
663    #[inline]
664    pub fn write_u32(&mut self, value: u32) {
665        self.inner.write_u32(value);
666    }
667
668    /// Write a u64
669    #[inline]
670    pub fn write_u64(&mut self, value: u64) {
671        self.inner.write_u64(value);
672    }
673
674    /// Write a usize
675    #[inline]
676    pub fn write_usize(&mut self, value: usize) {
677        self.inner.write_usize(value);
678    }
679
680    /// Write an f32
681    #[inline]
682    pub fn write_f32(&mut self, value: f32) {
683        self.inner.write_f32(value);
684    }
685
686    /// Write an f64
687    #[inline]
688    pub fn write_f64(&mut self, value: f64) {
689        self.inner.write_f64(value);
690    }
691}
692
693impl Default for JsonWriter<CompactWriter> {
694    fn default() -> Self {
695        Self::new()
696    }
697}
698
699#[cfg(test)]
700mod tests {
701    use super::*;
702
703    #[test]
704    fn test_write_string() {
705        let mut writer = JsonWriter::new();
706        writer.write_string("hello");
707        assert_eq!(writer.into_string(), r#""hello""#);
708    }
709
710    #[test]
711    fn test_write_escaped_string() {
712        let mut writer = JsonWriter::new();
713        writer.write_string("hello\nworld");
714        assert_eq!(writer.into_string(), r#""hello\nworld""#);
715    }
716
717    #[test]
718    fn test_write_object() {
719        let mut writer = JsonWriter::new();
720        writer.begin_object();
721        writer.write_key("name");
722        writer.write_string("test");
723        writer.end_object();
724        assert_eq!(writer.into_string(), r#"{"name":"test"}"#);
725    }
726
727    #[test]
728    fn test_write_array() {
729        let mut writer = JsonWriter::new();
730        writer.begin_array();
731        writer.write_i64(1);
732        writer.write_comma();
733        writer.write_i64(2);
734        writer.write_comma();
735        writer.write_i64(3);
736        writer.end_array();
737        assert_eq!(writer.into_string(), "[1,2,3]");
738    }
739
740    #[test]
741    fn test_pretty_print() {
742        let mut writer = JsonWriter::with_indent(2);
743        writer.begin_object();
744        writer.write_key("name");
745        writer.write_string("test");
746        writer.end_object();
747        let result = writer.into_string();
748        assert!(result.contains('\n'));
749    }
750}