musli_web/
buf.rs

1use core::cell::Cell;
2use core::fmt;
3use core::mem;
4use core::num::NonZeroUsize;
5use core::ops::Range;
6
7use alloc::vec::Vec;
8use musli::mode::Binary;
9use musli::{Encode, storage};
10
11#[derive(Debug)]
12pub(crate) struct InvalidFrame {
13    frame: Range<usize>,
14    size: usize,
15}
16
17impl fmt::Display for InvalidFrame {
18    #[inline]
19    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
20        write!(
21            f,
22            "Invalid frame {}-{} (in {} bytes)",
23            self.frame.start, self.frame.end, self.size
24        )
25    }
26}
27
28/// A length-prefixed buffer which keeps track of the start of each frame and
29/// allows them to be iterated over.
30#[derive(Default)]
31pub(crate) struct Buf {
32    start: Option<NonZeroUsize>,
33    buffer: Vec<u8>,
34    read: Cell<usize>,
35}
36
37impl Buf {
38    /// Write data to the current frame, or start a new frame if no frame is
39    /// being written.
40    ///
41    /// This needs to be paired with a call to [`Buf::done`] to complete an
42    /// outgoing frame.
43    ///
44    /// If a new frame is started, a new start point is recorded.
45    #[inline]
46    pub(crate) fn write<T>(&mut self, value: T) -> Result<(), storage::Error>
47    where
48        T: Encode<Binary>,
49    {
50        if self.start.is_none() {
51            let bytes = 0u32.to_ne_bytes();
52            self.buffer.extend_from_slice(&bytes);
53            self.start = NonZeroUsize::new(self.buffer.len());
54        }
55
56        storage::to_writer(&mut self.buffer, &value)?;
57        Ok(())
58    }
59
60    /// Check if the buffer is empty.
61    #[inline]
62    pub(crate) fn is_empty(&self) -> bool {
63        // NB: Read should never exceed the length of the buffer.
64        debug_assert!(self.read.get() <= self.buffer.len());
65        self.read.get() >= self.buffer.len()
66    }
67
68    fn len_at_mut(&mut self, at: usize) -> Option<&mut [u8; 4]> {
69        match self.buffer.get_mut(at..at + mem::size_of::<u32>())? {
70            bytes if bytes.len() == mem::size_of::<u32>() => {
71                Some(unsafe { &mut *bytes.as_mut_ptr().cast() })
72            }
73            _ => None,
74        }
75    }
76
77    /// Mark an outgoing frame as done from the previous start point.
78    ///
79    /// If no start point is recorded, calling this method does nothing.
80    #[inline]
81    pub(crate) fn done(&mut self) {
82        if let Some(start) = self.start.take() {
83            let l = u32::try_from(self.buffer.len().saturating_sub(start.get()))
84                .unwrap_or(u32::MAX)
85                .to_ne_bytes();
86
87            let Some(len) = self.len_at_mut(start.get().saturating_sub(mem::size_of::<u32>()))
88            else {
89                return;
90            };
91
92            *len = l;
93        }
94    }
95
96    /// Reset the buffer to the previous start point.
97    ///
98    /// If no start point is set, this method does nothing.
99    #[inline]
100    pub(crate) fn reset(&mut self) {
101        if let Some(start) = self.start {
102            self.buffer.truncate(start.get());
103        }
104    }
105
106    #[inline]
107    pub(crate) fn clear(&mut self) {
108        self.start = None;
109        self.buffer.clear();
110        self.read.set(0);
111    }
112
113    /// Get the next frame starting at the given location.
114    #[inline]
115    pub(crate) fn read(&self) -> Result<Option<&[u8]>, InvalidFrame> {
116        let read = self.read.get();
117
118        if self.buffer.len() == read {
119            return Ok(None);
120        }
121
122        let Some(tail) = self.buffer.get(read..) else {
123            return Err(InvalidFrame {
124                frame: 0..read,
125                size: self.buffer.len(),
126            });
127        };
128
129        let Some((head, tail)) = tail.split_at_checked(mem::size_of::<u32>()) else {
130            return Err(InvalidFrame {
131                frame: 0..read,
132                size: self.buffer.len(),
133            });
134        };
135
136        let frame = read..read + mem::size_of::<u32>();
137
138        let &[a, b, c, d] = head else {
139            return Err(InvalidFrame {
140                frame: frame.clone(),
141                size: self.buffer.len(),
142            });
143        };
144
145        let Ok(len) = usize::try_from(u32::from_ne_bytes([a, b, c, d])) else {
146            return Err(InvalidFrame {
147                frame: frame.clone(),
148                size: self.buffer.len(),
149            });
150        };
151
152        let Some(out) = tail.get(..len) else {
153            return Err(InvalidFrame {
154                frame: frame.start..frame.end + len,
155                size: self.buffer.len(),
156            });
157        };
158
159        self.read.set(
160            read.saturating_add(mem::size_of::<u32>())
161                .saturating_add(len),
162        );
163        Ok(Some(out))
164    }
165}