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/// Generated encoding-only objects implement this contract. `FIRST` omits the
259/// 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 Size for u64 {
270    fn json_size(&self) -> usize {
271        self.checked_ilog10().unwrap_or(0) as usize + 1
272    }
273}
274
275impl Size for usize {
276    fn json_size(&self) -> usize {
277        (*self as u64).json_size()
278    }
279}
280
281impl Size for i64 {
282    fn json_size(&self) -> usize {
283        self.unsigned_abs().json_size() + usize::from(*self < 0)
284    }
285}
286
287impl Size for f64 {
288    fn json_size(&self) -> usize {
289        if self.is_finite() {
290            ryu::Buffer::new().format_finite(*self).len()
291        } else {
292            4
293        }
294    }
295}
296
297impl TrySize for [u8] {
298    fn try_size(&self) -> Result<usize, Overflow> {
299        let mut value = self;
300        let mut len = Overflow::try_add(2, value.len())?;
301        while let Some(split) = Split::escape(&mut value) {
302            len = Overflow::try_add(len, split.token.overhead())?;
303        }
304        Ok(len)
305    }
306}
307
308impl<T: Backing> TrySize for Text<T> {
309    fn try_size(&self) -> Result<usize, Overflow> {
310        self.as_slice().try_size()
311    }
312}
313
314impl<T: Backing> TrySize for Plain<T> {
315    fn try_size(&self) -> Result<usize, Overflow> {
316        Overflow::try_add(2, self.len())
317    }
318}
319
320impl<T: Backing> TrySize for Raw<T> {
321    fn try_size(&self) -> Result<usize, Overflow> {
322        Ok(self.len())
323    }
324}
325
326impl TrySize for str {
327    fn try_size(&self) -> Result<usize, Overflow> {
328        self.as_bytes().try_size()
329    }
330}
331
332impl TrySize for String {
333    fn try_size(&self) -> Result<usize, Overflow> {
334        self.as_str().try_size()
335    }
336}
337
338pub trait TryEncode: TrySize + Sized {
339    fn try_write<'a, W: write::Sink>(
340        &self,
341        writer: Writer<'a, W>,
342    ) -> ControlFlow<Failure<'a, W>, Writer<'a, W>>;
343
344    /// Admits and writes one exact suffix, preserving the original prefix on failure.
345    fn try_append(&self, out: &mut Vec<u8>) -> Result<(), EncodeError> {
346        let expected = self.try_size().map_err(EncodeError::Size)?;
347        consume::Error::check(expected, u32::MAX as usize).map_err(EncodeError::Capacity)?;
348        let mut staged =
349            write::Staged::try_grow(out, expected as u32).map_err(EncodeError::Allocate)?;
350        let mut spare = staged.spare();
351        if let ControlFlow::Break(failure) = self.try_write(Writer::new(&mut spare)) {
352            return Err(EncodeError::Capacity(failure.into_error()));
353        }
354        spare.try_complete().map_err(EncodeError::Length)?;
355        staged.commit();
356        Ok(())
357    }
358
359    fn try_encode(&self) -> Result<Vec<u8>, EncodeError> {
360        let mut out = Vec::new();
361        self.try_append(&mut out)?;
362        Ok(out)
363    }
364}
365
366impl TryEncode for u64 {
367    fn try_write<'a, W: write::Sink>(
368        &self,
369        writer: Writer<'a, W>,
370    ) -> ControlFlow<Failure<'a, W>, Writer<'a, W>> {
371        writer.put_u64(*self)
372    }
373}
374
375impl TryEncode for usize {
376    fn try_write<'a, W: write::Sink>(
377        &self,
378        writer: Writer<'a, W>,
379    ) -> ControlFlow<Failure<'a, W>, Writer<'a, W>> {
380        (*self as u64).try_write(writer)
381    }
382}
383
384pub struct Body<T>(T);
385
386/// A checked JSON response retaining only its bound and original encoder owner.
387#[repr(transparent)]
388pub struct Response<T>(bounded::Plan<Body<T>>);
389
390static HEAD: Template = Template::new(&[Entry::new(
391    literal!(Name::try_bytes(b"content-type")),
392    literal!(Value::try_bytes(b"application/json")),
393)]);
394
395impl<T> Body<T> {
396    pub fn new(value: T) -> Self {
397        Self(value)
398    }
399}
400
401impl<T: TryEncode + TryMaxSize> Body<T> {
402    /// Prepares an OK JSON response, checking its upper bound once before emission.
403    /// The static JSON header imposes no lifetime requirement on the body.
404    pub fn try_response(self) -> Result<Response<T>, Rejected<Self, T::Error>> {
405        Ok(Response(bounded::Plan::try_new(self)?))
406    }
407}
408
409impl<'req, 'd, T: TryEncode + TryMaxSize> IntoOutput<'req, 'd> for Response<T> {
410    type Output = bounded::Output<'req, bounded::Plan<Body<T>>, 0>;
411
412    fn into_output(self) -> Self::Output {
413        bounded::Output::new(Code::OK, &HEAD, Set::new(), self.0)
414    }
415}
416
417impl<T: TryEncode> exact::TryBody for Body<T> {
418    fn try_len(&self) -> Result<usize, Overflow> {
419        self.0.try_size()
420    }
421}
422
423impl<T: TryEncode> Encode for Body<T> {
424    fn encode<W>(self, out: &mut W) -> Result<(), consume::Error>
425    where
426        W: write::Sink<Error = consume::Error>,
427    {
428        let ControlFlow::Break(failure) = self.0.try_write(Writer::new(out)) else {
429            return Ok(());
430        };
431        Err(failure.into_error())
432    }
433}
434
435impl<T: TryEncode + TryMaxSize> bounded::TryBody for Body<T> {
436    type Error = T::Error;
437
438    fn try_max_len(&self) -> Result<usize, Self::Error> {
439        self.0.try_max_size()
440    }
441}