musli/
reader.rs

1//! Trait governing how to read bytes.
2
3use core::array;
4use core::fmt;
5use core::marker;
6use core::ops::Range;
7use core::ptr;
8use core::slice;
9
10use crate::de::UnsizedVisitor;
11use crate::Context;
12
13mod sealed {
14    use super::{Limit, Reader};
15
16    pub trait Sealed {}
17
18    impl Sealed for &[u8] {}
19    impl Sealed for super::SliceReader<'_> {}
20    impl<'de, R> Sealed for Limit<R> where R: Reader<'de> {}
21    impl<'de, R> Sealed for &mut R where R: ?Sized + Reader<'de> {}
22}
23
24/// Trait governing how a source of bytes is read.
25///
26/// This requires the reader to be able to hand out contiguous references to the
27/// byte source through [`Reader::read_bytes`].
28pub trait Reader<'de>: self::sealed::Sealed {
29    /// Type borrowed from self.
30    ///
31    /// Why oh why would we want to do this over having a simple `&'this mut T`?
32    ///
33    /// We want to avoid recursive types, which will blow up the compiler. And
34    /// the above is a typical example of when that can go wrong. This ensures
35    /// that each call to `borrow_mut` dereferences the [`Reader`] at each step
36    /// to avoid constructing a large muted type, like `&mut &mut &mut
37    /// SliceReader<'de>`.
38    type Mut<'this>: Reader<'de>
39    where
40        Self: 'this;
41
42    /// Borrow the current reader.
43    fn borrow_mut(&mut self) -> Self::Mut<'_>;
44
45    /// Skip over the given number of bytes.
46    fn skip<C>(&mut self, cx: &C, n: usize) -> Result<(), C::Error>
47    where
48        C: ?Sized + Context;
49
50    /// Peek the next value.
51    fn peek(&mut self) -> Option<u8>;
52
53    /// Read a slice into the given buffer.
54    #[inline]
55    fn read<C>(&mut self, cx: &C, buf: &mut [u8]) -> Result<(), C::Error>
56    where
57        C: ?Sized + Context,
58    {
59        struct Visitor<'a>(&'a mut [u8]);
60
61        impl<'a, 'de, C> UnsizedVisitor<'de, C, [u8]> for Visitor<'a>
62        where
63            C: ?Sized + Context,
64        {
65            type Ok = ();
66
67            #[inline]
68            fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
69                write!(f, "bytes")
70            }
71
72            #[inline]
73            fn visit_borrowed(self, cx: &C, bytes: &'de [u8]) -> Result<Self::Ok, C::Error> {
74                self.visit_ref(cx, bytes)
75            }
76
77            #[inline]
78            fn visit_ref(self, _: &C, bytes: &[u8]) -> Result<Self::Ok, C::Error> {
79                self.0.copy_from_slice(bytes);
80                Ok(())
81            }
82        }
83
84        self.read_bytes(cx, buf.len(), Visitor(buf))
85    }
86
87    /// Read a slice out of the current reader.
88    fn read_bytes<C, V>(&mut self, cx: &C, n: usize, visitor: V) -> Result<V::Ok, C::Error>
89    where
90        C: ?Sized + Context,
91        V: UnsizedVisitor<'de, C, [u8]>;
92
93    /// Read a single byte.
94    #[inline]
95    fn read_byte<C>(&mut self, cx: &C) -> Result<u8, C::Error>
96    where
97        C: ?Sized + Context,
98    {
99        let [byte] = self.read_array::<C, 1>(cx)?;
100        Ok(byte)
101    }
102
103    /// Read an array out of the current reader.
104    #[inline]
105    fn read_array<C, const N: usize>(&mut self, cx: &C) -> Result<[u8; N], C::Error>
106    where
107        C: ?Sized + Context,
108    {
109        struct Visitor<const N: usize>([u8; N]);
110
111        impl<'de, const N: usize, C> UnsizedVisitor<'de, C, [u8]> for Visitor<N>
112        where
113            C: ?Sized + Context,
114        {
115            type Ok = [u8; N];
116
117            #[inline]
118            fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
119                write!(f, "bytes")
120            }
121
122            #[inline]
123            fn visit_borrowed(self, cx: &C, bytes: &'de [u8]) -> Result<Self::Ok, C::Error> {
124                self.visit_ref(cx, bytes)
125            }
126
127            #[inline]
128            fn visit_ref(mut self, cx: &C, bytes: &[u8]) -> Result<Self::Ok, C::Error> {
129                self.0.copy_from_slice(bytes);
130                cx.advance(bytes.len());
131                Ok(self.0)
132            }
133        }
134
135        self.read_bytes(cx, N, Visitor([0u8; N]))
136    }
137
138    /// Keep an accurate record of the position within the reader.
139    fn limit(self, limit: usize) -> Limit<Self>
140    where
141        Self: Sized,
142    {
143        Limit {
144            remaining: limit,
145            reader: self,
146        }
147    }
148}
149
150impl<'de> IntoReader<'de> for &'de [u8] {
151    type Reader = &'de [u8];
152
153    #[inline]
154    fn into_reader(self) -> Self::Reader {
155        self
156    }
157}
158
159impl<'a, 'de, R> IntoReader<'de> for &'a mut R
160where
161    R: ?Sized + Reader<'de>,
162{
163    type Reader = &'a mut R;
164
165    #[inline]
166    fn into_reader(self) -> Self::Reader {
167        self
168    }
169}
170
171impl<'de> Reader<'de> for &'de [u8] {
172    type Mut<'this> = &'this mut &'de [u8] where Self: 'this;
173
174    #[inline]
175    fn borrow_mut(&mut self) -> Self::Mut<'_> {
176        self
177    }
178
179    #[inline]
180    fn skip<C>(&mut self, cx: &C, n: usize) -> Result<(), C::Error>
181    where
182        C: ?Sized + Context,
183    {
184        if self.len() < n {
185            return Err(cx.message(SliceUnderflow {
186                n,
187                remaining: self.len(),
188            }));
189        }
190
191        let (_, tail) = self.split_at(n);
192        *self = tail;
193        cx.advance(n);
194        Ok(())
195    }
196
197    #[inline]
198    fn read<C>(&mut self, cx: &C, buf: &mut [u8]) -> Result<(), C::Error>
199    where
200        C: ?Sized + Context,
201    {
202        if self.len() < buf.len() {
203            return Err(cx.custom(SliceUnderflow::new(buf.len(), self.len())));
204        }
205
206        let (head, tail) = self.split_at(buf.len());
207        buf.copy_from_slice(head);
208        *self = tail;
209        cx.advance(buf.len());
210        Ok(())
211    }
212
213    #[inline]
214    fn read_bytes<C, V>(&mut self, cx: &C, n: usize, visitor: V) -> Result<V::Ok, C::Error>
215    where
216        C: ?Sized + Context,
217        V: UnsizedVisitor<'de, C, [u8]>,
218    {
219        if self.len() < n {
220            return Err(cx.custom(SliceUnderflow::new(n, self.len())));
221        }
222
223        let (head, tail) = self.split_at(n);
224        *self = tail;
225        let ok = visitor.visit_borrowed(cx, head)?;
226        cx.advance(n);
227        Ok(ok)
228    }
229
230    #[inline]
231    fn read_byte<C>(&mut self, cx: &C) -> Result<u8, C::Error>
232    where
233        C: ?Sized + Context,
234    {
235        let &[first, ref tail @ ..] = *self else {
236            return Err(cx.custom(SliceUnderflow::new(1, self.len())));
237        };
238
239        *self = tail;
240        cx.advance(1);
241        Ok(first)
242    }
243
244    #[inline]
245    fn read_array<C, const N: usize>(&mut self, cx: &C) -> Result<[u8; N], C::Error>
246    where
247        C: ?Sized + Context,
248    {
249        if self.len() < N {
250            return Err(cx.custom(SliceUnderflow::new(N, self.len())));
251        }
252
253        let (head, tail) = self.split_at(N);
254        *self = tail;
255        cx.advance(N);
256        Ok(array::from_fn(|n| head[n]))
257    }
258
259    #[inline]
260    fn peek(&mut self) -> Option<u8> {
261        self.first().copied()
262    }
263}
264
265/// Coerce a type into a [`Reader`].
266pub trait IntoReader<'de>: self::sealed::Sealed {
267    /// The reader type.
268    type Reader: Reader<'de>;
269
270    /// Convert the type into a reader.
271    fn into_reader(self) -> Self::Reader;
272}
273
274/// An efficient [`Reader`] wrapper around a slice.
275pub struct SliceReader<'de> {
276    range: Range<*const u8>,
277    _marker: marker::PhantomData<&'de [u8]>,
278}
279
280// SAFETY: `SliceReader` is effectively equivalent to `&'de [u8]`.
281unsafe impl Send for SliceReader<'_> {}
282// SAFETY: `SliceReader` is effectively equivalent to `&'de [u8]`.
283unsafe impl Sync for SliceReader<'_> {}
284
285impl<'de> SliceReader<'de> {
286    /// Construct a new instance around the specified slice.
287    #[inline]
288    pub fn new(slice: &'de [u8]) -> Self {
289        Self {
290            range: slice.as_ptr_range(),
291            _marker: marker::PhantomData,
292        }
293    }
294
295    /// Get the remaining contents of the reader as a slice.
296    ///
297    /// # Examples
298    ///
299    /// ```
300    /// use musli::Context;
301    /// use musli::reader::{Reader, SliceReader};
302    ///
303    /// fn process<C>(cx: &C) -> Result<(), C::Error>
304    /// where
305    ///     C: ?Sized + Context
306    /// {
307    ///     let mut reader = SliceReader::new(&[1, 2, 3, 4]);
308    ///     assert_eq!(reader.as_slice(), &[1, 2, 3, 4]);
309    ///     reader.skip(cx, 2)?;
310    ///     assert_eq!(reader.as_slice(), &[3, 4]);
311    ///     Ok(())
312    /// }
313    /// ```
314    pub fn as_slice(&self) -> &'de [u8] {
315        unsafe { slice::from_raw_parts(self.range.start, self.remaining()) }
316    }
317
318    /// Get remaining bytes in the reader.
319    ///
320    /// # Examples
321    ///
322    /// ```
323    /// use musli::Context;
324    /// use musli::reader::{Reader, SliceReader};
325    ///
326    /// fn process<C>(cx: &C) -> Result<(), C::Error>
327    /// where
328    ///     C: ?Sized + Context
329    /// {
330    ///     let mut reader = SliceReader::new(&[1, 2, 3, 4]);
331    ///     assert_eq!(reader.remaining(), 4);
332    ///     reader.skip(cx, 2);
333    ///     assert_eq!(reader.remaining(), 2);
334    ///     Ok(())
335    /// }
336    /// ```
337    pub fn remaining(&self) -> usize {
338        self.range.end as usize - self.range.start as usize
339    }
340}
341
342impl<'de> Reader<'de> for SliceReader<'de> {
343    type Mut<'this> = &'this mut Self where Self: 'this;
344
345    #[inline]
346    fn borrow_mut(&mut self) -> Self::Mut<'_> {
347        self
348    }
349
350    #[inline]
351    fn skip<C>(&mut self, cx: &C, n: usize) -> Result<(), C::Error>
352    where
353        C: ?Sized + Context,
354    {
355        self.range.start = bounds_check_add(cx, &self.range, n)?;
356        cx.advance(n);
357        Ok(())
358    }
359
360    #[inline]
361    fn read_bytes<C, V>(&mut self, cx: &C, n: usize, visitor: V) -> Result<V::Ok, C::Error>
362    where
363        C: ?Sized + Context,
364        V: UnsizedVisitor<'de, C, [u8]>,
365    {
366        let outcome = bounds_check_add(cx, &self.range, n)?;
367
368        let ok = unsafe {
369            let bytes = slice::from_raw_parts(self.range.start, n);
370            self.range.start = outcome;
371            visitor.visit_borrowed(cx, bytes)?
372        };
373
374        cx.advance(n);
375        Ok(ok)
376    }
377
378    #[inline]
379    fn peek(&mut self) -> Option<u8> {
380        if self.range.start == self.range.end {
381            return None;
382        }
383
384        // SAFETY: we've checked that the elements are in bound above.
385        unsafe { Some(ptr::read(self.range.start)) }
386    }
387
388    #[inline]
389    fn read<C>(&mut self, cx: &C, buf: &mut [u8]) -> Result<(), C::Error>
390    where
391        C: ?Sized + Context,
392    {
393        let outcome = bounds_check_add(cx, &self.range, buf.len())?;
394
395        unsafe {
396            ptr::copy_nonoverlapping(self.range.start, buf.as_mut_ptr(), buf.len());
397            self.range.start = outcome;
398        }
399
400        cx.advance(buf.len());
401        Ok(())
402    }
403}
404
405#[inline]
406fn bounds_check_add<C>(cx: &C, range: &Range<*const u8>, len: usize) -> Result<*const u8, C::Error>
407where
408    C: ?Sized + Context,
409{
410    let outcome = range.start.wrapping_add(len);
411
412    if outcome > range.end || outcome < range.start {
413        Err(cx.message(SliceUnderflow {
414            n: len,
415            remaining: (range.end as usize).wrapping_sub(range.start as usize),
416        }))
417    } else {
418        Ok(outcome)
419    }
420}
421
422/// Limit the number of bytes that can be read out of a reader to the specified limit.
423///
424/// Constructed through [Reader::limit].
425pub struct Limit<R> {
426    remaining: usize,
427    reader: R,
428}
429
430impl<R> Limit<R> {
431    /// Get the remaining data in the limited reader.
432    pub fn remaining(&self) -> usize {
433        self.remaining
434    }
435}
436
437impl<'de, R> Limit<R>
438where
439    R: Reader<'de>,
440{
441    fn bounds_check<C>(&mut self, cx: &C, n: usize) -> Result<(), C::Error>
442    where
443        C: ?Sized + Context,
444    {
445        match self.remaining.checked_sub(n) {
446            Some(remaining) => {
447                self.remaining = remaining;
448                Ok(())
449            }
450            None => Err(cx.message("Reader out of bounds")),
451        }
452    }
453}
454
455impl<'de, R> Reader<'de> for Limit<R>
456where
457    R: Reader<'de>,
458{
459    type Mut<'this> = &'this mut Self where Self: 'this;
460
461    #[inline]
462    fn borrow_mut(&mut self) -> Self::Mut<'_> {
463        self
464    }
465
466    #[inline]
467    fn skip<C>(&mut self, cx: &C, n: usize) -> Result<(), C::Error>
468    where
469        C: ?Sized + Context,
470    {
471        self.bounds_check(cx, n)?;
472        self.reader.skip(cx, n)
473    }
474
475    #[inline]
476    fn read_bytes<C, V>(&mut self, cx: &C, n: usize, visitor: V) -> Result<V::Ok, C::Error>
477    where
478        C: ?Sized + Context,
479        V: UnsizedVisitor<'de, C, [u8]>,
480    {
481        self.bounds_check(cx, n)?;
482        self.reader.read_bytes(cx, n, visitor)
483    }
484
485    #[inline]
486    fn peek(&mut self) -> Option<u8> {
487        if self.remaining > 0 {
488            self.reader.peek()
489        } else {
490            None
491        }
492    }
493
494    #[inline]
495    fn read<C>(&mut self, cx: &C, buf: &mut [u8]) -> Result<(), C::Error>
496    where
497        C: ?Sized + Context,
498    {
499        self.bounds_check(cx, buf.len())?;
500        self.reader.read(cx, buf)
501    }
502
503    #[inline]
504    fn read_byte<C>(&mut self, cx: &C) -> Result<u8, C::Error>
505    where
506        C: ?Sized + Context,
507    {
508        self.bounds_check(cx, 1)?;
509        self.reader.read_byte(cx)
510    }
511
512    #[inline]
513    fn read_array<C, const N: usize>(&mut self, cx: &C) -> Result<[u8; N], C::Error>
514    where
515        C: ?Sized + Context,
516    {
517        self.bounds_check(cx, N)?;
518        self.reader.read_array(cx)
519    }
520}
521
522// Forward implementations.
523
524impl<'de, R> Reader<'de> for &mut R
525where
526    R: ?Sized + Reader<'de>,
527{
528    type Mut<'this> = &'this mut R where Self: 'this;
529
530    #[inline]
531    fn borrow_mut(&mut self) -> Self::Mut<'_> {
532        self
533    }
534
535    #[inline]
536    fn skip<C>(&mut self, cx: &C, n: usize) -> Result<(), C::Error>
537    where
538        C: ?Sized + Context,
539    {
540        (**self).skip(cx, n)
541    }
542
543    #[inline]
544    fn read_bytes<C, V>(&mut self, cx: &C, n: usize, visitor: V) -> Result<V::Ok, C::Error>
545    where
546        C: ?Sized + Context,
547        V: UnsizedVisitor<'de, C, [u8]>,
548    {
549        (**self).read_bytes(cx, n, visitor)
550    }
551
552    #[inline]
553    fn peek(&mut self) -> Option<u8> {
554        (**self).peek()
555    }
556
557    #[inline]
558    fn read<C>(&mut self, cx: &C, buf: &mut [u8]) -> Result<(), C::Error>
559    where
560        C: ?Sized + Context,
561    {
562        (**self).read(cx, buf)
563    }
564
565    #[inline]
566    fn read_byte<C>(&mut self, cx: &C) -> Result<u8, C::Error>
567    where
568        C: ?Sized + Context,
569    {
570        (**self).read_byte(cx)
571    }
572
573    #[inline]
574    fn read_array<C, const N: usize>(&mut self, cx: &C) -> Result<[u8; N], C::Error>
575    where
576        C: ?Sized + Context,
577    {
578        (**self).read_array(cx)
579    }
580}
581
582/// Underflow when trying to read from a slice.
583#[derive(Debug)]
584pub(crate) struct SliceUnderflow {
585    n: usize,
586    remaining: usize,
587}
588
589impl SliceUnderflow {
590    pub(crate) fn new(n: usize, remaining: usize) -> Self {
591        Self { n, remaining }
592    }
593}
594
595impl fmt::Display for SliceUnderflow {
596    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
597        let SliceUnderflow { n, remaining } = self;
598
599        write!(
600            f,
601            "Tried to read {n} bytes from slice, with {remaining} byte remaining"
602        )
603    }
604}
605
606impl core::error::Error for SliceUnderflow {}