Skip to main content

sark_json/writing/
mod.rs

1//! JSON sizing and encoding machinery.
2use std::{mem, ops::ControlFlow};
3
4use escape::Split;
5use exact::Overflow;
6use o3::{
7    alloc,
8    buffer::write::{self, SpareExt as _, consume, spans},
9    mem::Rejected,
10};
11use ryu;
12use sark_http::{
13    field::{Value, literal},
14    payload::response::{
15        Encode, bounded,
16        dispatch::IntoOutput,
17        exact,
18        headers::{Entry, Name, head::Template, set::Set},
19        status::Code,
20    },
21};
22
23use crate::{Backing, Plain, Raw, Text, TextValue};
24
25mod escape;
26
27#[repr(transparent)]
28#[must_use = "continue the writer chain"]
29pub struct Writer<'a, W: write::Sink> {
30    inner: &'a mut W,
31}
32
33#[repr(C)]
34#[must_use = "preserve the writer and its exact failure"]
35pub struct Failure<'a, W: write::Sink> {
36    writer: Writer<'a, W>,
37    error: W::Error,
38}
39
40impl<'a, W: write::Sink> Writer<'a, W> {
41    pub fn new(inner: &'a mut W) -> Self {
42        Self { inner }
43    }
44
45    /// Writes bytes verbatim without JSON validation, quoting, or escaping.
46    pub fn put(self, src: &[u8]) -> ControlFlow<Failure<'a, W>, Self> {
47        match self.inner.write_slice(src) {
48            Ok(()) => ControlFlow::Continue(self),
49            Err(error) => ControlFlow::Break(Failure {
50                writer: self,
51                error,
52            }),
53        }
54    }
55
56    pub fn put_str<T: TextValue + ?Sized>(self, value: &T) -> ControlFlow<Failure<'a, W>, Self> {
57        let mut value = value.as_bytes();
58        match Split::escape(&mut value) {
59            Some(split) => self.put(b"\"")?.put_rest(value, Some(split)),
60            None => self.put_plain(&Plain::admitted(value)),
61        }
62    }
63
64    fn put_rest<'s>(
65        mut self,
66        mut value: &'s [u8],
67        mut next: Option<Split<'s>>,
68    ) -> ControlFlow<Failure<'a, W>, Self> {
69        while let Some(split) = next {
70            if !split.prefix.is_empty() {
71                self = self.put(split.prefix)?;
72            }
73            self = self.put(split.token.bytes())?;
74            next = Split::escape(&mut value);
75        }
76        if !value.is_empty() {
77            self = self.put(value)?;
78        }
79        self.put(b"\"")
80    }
81
82    pub fn put_plain<T: Backing>(self, value: &Plain<T>) -> ControlFlow<Failure<'a, W>, Self> {
83        match self
84            .inner
85            .write_slices([b"\"".as_slice(), value.as_slice(), b"\"".as_slice()])
86        {
87            Ok(()) => ControlFlow::Continue(self),
88            Err(error) => ControlFlow::Break(Failure {
89                writer: self,
90                error,
91            }),
92        }
93    }
94
95    pub fn put_u64(self, value: u64) -> ControlFlow<Failure<'a, W>, Self> {
96        let mut digits;
97        let bytes = if value < 100 {
98            let pair = &Digits::PAIRS[value as usize];
99            if value < 10 { &pair[1..] } else { pair }
100        } else {
101            digits = Digits { bytes: [0; 20] };
102            digits.write(value)
103        };
104        self.put(bytes)
105    }
106
107    pub fn put_i64(mut self, value: i64) -> ControlFlow<Failure<'a, W>, Self> {
108        if value < 0 {
109            self = self.put(b"-")?;
110        }
111        self.put_u64(value.unsigned_abs())
112    }
113
114    /// Writes a JSON number, or `null` for NaN and either infinity.
115    pub fn put_f64(self, value: f64) -> ControlFlow<Failure<'a, W>, Self> {
116        if value.is_finite() {
117            self.put(ryu::Buffer::new().format_finite(value).as_bytes())
118        } else {
119            self.put(b"null")
120        }
121    }
122}
123
124impl<'a, W: write::Sink> Failure<'a, W> {
125    pub fn into_error(self) -> W::Error {
126        self.error
127    }
128}
129
130const _: () = {
131    assert!(usize::BITS <= u64::BITS);
132    assert!(mem::size_of::<Writer<'static, write::Slice<'static>>>() == mem::size_of::<usize>());
133    assert!(
134        mem::size_of::<Option<Writer<'static, write::Slice<'static>>>>() == mem::size_of::<usize>()
135    );
136    assert!(
137        mem::size_of::<Failure<'static, write::Slice<'static>>>() == mem::size_of::<usize>() * 3
138    );
139    assert!(
140        mem::size_of::<
141            ControlFlow<
142                Failure<'static, write::Slice<'static>>,
143                Writer<'static, write::Slice<'static>>,
144            >,
145        >() == mem::size_of::<usize>() * 3
146    );
147};
148
149struct Digits {
150    bytes: [u8; 20],
151}
152
153impl Digits {
154    const PAIRS: [[u8; 2]; 100] = {
155        let mut pairs = [[0; 2]; 100];
156        let mut value = 0;
157        while value < pairs.len() {
158            pairs[value] = [b'0' + (value / 10) as u8, b'0' + (value % 10) as u8];
159            value += 1;
160        }
161        pairs
162    };
163
164    fn write(&mut self, mut value: u64) -> &[u8] {
165        let bytes = &mut self.bytes;
166        let mut start = bytes.len();
167        while value >= 100 {
168            start -= 2;
169            bytes[start..start + 2].copy_from_slice(&Self::PAIRS[(value % 100) as usize]);
170            value /= 100;
171        }
172        start -= 2;
173        bytes[start..start + 2].copy_from_slice(&Self::PAIRS[value as usize]);
174        &bytes[start + usize::from(value < 10)..]
175    }
176}
177
178pub trait Size {
179    fn json_size(&self) -> usize;
180}
181
182/// Sizing that can reject an unrepresentable encoded length.
183pub trait TrySize {
184    fn try_size(&self) -> Result<usize, Overflow>;
185}
186
187/// A checked upper bound that can be obtained without exact sizing.
188pub trait TryMaxSize {
189    type Error;
190
191    fn try_max_size(&self) -> Result<usize, Self::Error>;
192}
193
194impl<T: Size + ?Sized> TrySize for T {
195    fn try_size(&self) -> Result<usize, Overflow> {
196        Ok(self.json_size())
197    }
198}
199
200/// Generated sequence code keeps each value alive while borrowing its encoder.
201#[doc(hidden)]
202#[derive(Debug)]
203pub struct Item<T>(pub T);
204
205impl<T: Copy + IntoIterator> IntoIterator for &Item<T> {
206    type Item = T::Item;
207    type IntoIter = T::IntoIter;
208
209    fn into_iter(self) -> Self::IntoIter {
210        IntoIterator::into_iter(self.0)
211    }
212}
213
214impl<'inner, T: ?Sized> IntoIterator for Item<&&'inner T>
215where
216    &'inner T: IntoIterator,
217{
218    type Item = <&'inner T as IntoIterator>::Item;
219    type IntoIter = <&'inner T as IntoIterator>::IntoIter;
220
221    fn into_iter(self) -> Self::IntoIter {
222        IntoIterator::into_iter(*self.0)
223    }
224}
225
226#[doc(hidden)]
227pub trait Nested<'a> {
228    type Value: TrySize + ?Sized + 'a;
229
230    fn into_ref(self) -> &'a Self::Value;
231}
232
233impl<'a, T: TrySize + ?Sized> Nested<'a> for Item<&'a T> {
234    type Value = T;
235
236    fn into_ref(self) -> &'a Self::Value {
237        self.0
238    }
239}
240
241impl<'a, T: TrySize> Nested<'a> for &'a Item<T> {
242    type Value = T;
243
244    fn into_ref(self) -> &'a Self::Value {
245        &self.0
246    }
247}
248
249#[derive(Clone, Copy, Debug, PartialEq, Eq)]
250pub enum EncodeError {
251    Size(Overflow),
252    Allocate(alloc::Error),
253    Capacity(consume::Error),
254    Length(spans::Gap),
255}
256
257/// One nonempty sequence of object members, without enclosing braces.
258/// Nonempty generated objects using decoding or `encode` implement this contract.
259/// `FIRST` omits the leading comma; the size always excludes that comma and both braces.
260pub trait Members {
261    fn try_fields(&self) -> Result<usize, Overflow>;
262
263    fn try_members<'a, W: write::Sink, const FIRST: bool>(
264        &self,
265        writer: Writer<'a, W>,
266    ) -> ControlFlow<Failure<'a, W>, Writer<'a, W>>;
267}
268
269impl<T: Members + ?Sized> Members for &T {
270    fn try_fields(&self) -> Result<usize, Overflow> {
271        T::try_fields(*self)
272    }
273
274    fn try_members<'a, W: write::Sink, const FIRST: bool>(
275        &self,
276        writer: Writer<'a, W>,
277    ) -> ControlFlow<Failure<'a, W>, Writer<'a, W>> {
278        T::try_members::<W, FIRST>(*self, writer)
279    }
280}
281
282impl Size for u64 {
283    fn json_size(&self) -> usize {
284        self.checked_ilog10().unwrap_or(0) as usize + 1
285    }
286}
287
288impl Size for usize {
289    fn json_size(&self) -> usize {
290        (*self as u64).json_size()
291    }
292}
293
294impl Size for i64 {
295    fn json_size(&self) -> usize {
296        self.unsigned_abs().json_size() + usize::from(*self < 0)
297    }
298}
299
300impl Size for f64 {
301    fn json_size(&self) -> usize {
302        if self.is_finite() {
303            ryu::Buffer::new().format_finite(*self).len()
304        } else {
305            4
306        }
307    }
308}
309
310impl TrySize for [u8] {
311    fn try_size(&self) -> Result<usize, Overflow> {
312        let mut value = self;
313        let mut len = Overflow::try_add(2, value.len())?;
314        while let Some(split) = Split::escape(&mut value) {
315            len = Overflow::try_add(len, split.token.overhead())?;
316        }
317        Ok(len)
318    }
319}
320
321impl<T: Backing> TrySize for Text<T> {
322    fn try_size(&self) -> Result<usize, Overflow> {
323        self.as_slice().try_size()
324    }
325}
326
327impl<T: Backing> TrySize for Plain<T> {
328    fn try_size(&self) -> Result<usize, Overflow> {
329        Overflow::try_add(2, self.len())
330    }
331}
332
333impl<T: Backing> TrySize for Raw<T> {
334    fn try_size(&self) -> Result<usize, Overflow> {
335        Ok(self.len())
336    }
337}
338
339impl TrySize for str {
340    fn try_size(&self) -> Result<usize, Overflow> {
341        self.as_bytes().try_size()
342    }
343}
344
345impl TrySize for String {
346    fn try_size(&self) -> Result<usize, Overflow> {
347        self.as_str().try_size()
348    }
349}
350
351pub trait TryEncode: TrySize + Sized {
352    fn try_write<'a, W: write::Sink>(
353        &self,
354        writer: Writer<'a, W>,
355    ) -> ControlFlow<Failure<'a, W>, Writer<'a, W>>;
356
357    /// Admits and writes one exact suffix, preserving the original prefix on failure.
358    fn try_append(&self, out: &mut Vec<u8>) -> Result<(), EncodeError> {
359        let expected = self.try_size().map_err(EncodeError::Size)?;
360        consume::Error::check(expected, u32::MAX as usize).map_err(EncodeError::Capacity)?;
361        let mut staged =
362            write::Staged::try_grow(out, expected as u32).map_err(EncodeError::Allocate)?;
363        let mut spare = staged.spare();
364        if let ControlFlow::Break(failure) = self.try_write(Writer::new(&mut spare)) {
365            return Err(EncodeError::Capacity(failure.into_error()));
366        }
367        spare.try_complete().map_err(EncodeError::Length)?;
368        staged.commit();
369        Ok(())
370    }
371
372    fn try_encode(&self) -> Result<Vec<u8>, EncodeError> {
373        let mut out = Vec::new();
374        self.try_append(&mut out)?;
375        Ok(out)
376    }
377}
378
379impl TryEncode for u64 {
380    fn try_write<'a, W: write::Sink>(
381        &self,
382        writer: Writer<'a, W>,
383    ) -> ControlFlow<Failure<'a, W>, Writer<'a, W>> {
384        writer.put_u64(*self)
385    }
386}
387
388impl TryEncode for usize {
389    fn try_write<'a, W: write::Sink>(
390        &self,
391        writer: Writer<'a, W>,
392    ) -> ControlFlow<Failure<'a, W>, Writer<'a, W>> {
393        (*self as u64).try_write(writer)
394    }
395}
396
397pub struct Body<T>(T);
398
399/// A checked JSON response retaining only its bound and original encoder owner.
400#[repr(transparent)]
401pub struct Response<T>(bounded::Plan<Body<T>>);
402
403static HEAD: Template = Template::new(&[Entry::new(
404    literal!(Name::try_bytes(b"content-type")),
405    literal!(Value::try_bytes(b"application/json")),
406)]);
407
408impl<T> Body<T> {
409    pub fn new(value: T) -> Self {
410        Self(value)
411    }
412}
413
414impl<T: TryEncode + TryMaxSize> Body<T> {
415    /// Prepares an OK JSON response, checking its upper bound once before emission.
416    /// The static JSON header imposes no lifetime requirement on the body.
417    pub fn try_response(self) -> Result<Response<T>, Rejected<Self, T::Error>> {
418        Ok(Response(bounded::Plan::try_new(self)?))
419    }
420}
421
422impl<'req, 'd, T: TryEncode + TryMaxSize> IntoOutput<'req, 'd> for Response<T> {
423    type Output = bounded::Output<'req, bounded::Plan<Body<T>>, 0>;
424
425    fn into_output(self) -> Self::Output {
426        bounded::Output::new(Code::OK, &HEAD, Set::new(), self.0)
427    }
428}
429
430impl<T: TryEncode> exact::TryBody for Body<T> {
431    fn try_len(&self) -> Result<usize, Overflow> {
432        self.0.try_size()
433    }
434}
435
436impl<T: TryEncode> Encode for Body<T> {
437    fn encode<W>(self, out: &mut W) -> Result<(), consume::Error>
438    where
439        W: write::Sink<Error = consume::Error>,
440    {
441        let ControlFlow::Break(failure) = self.0.try_write(Writer::new(out)) else {
442            return Ok(());
443        };
444        Err(failure.into_error())
445    }
446}
447
448impl<T: TryEncode + TryMaxSize> bounded::TryBody for Body<T> {
449    type Error = T::Error;
450
451    fn try_max_len(&self) -> Result<usize, Self::Error> {
452        self.0.try_max_size()
453    }
454}