1use 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 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 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
182pub trait TrySize {
184 fn try_size(&self) -> Result<usize, Overflow>;
185}
186
187pub 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#[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
257pub 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 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#[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 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}