Skip to main content

substrate_data/wire/frame/
unpacker.rs

1//! [`Unpacker`]: reads a frame.
2
3use bumpalo::Bump;
4
5use super::sizes::{F32_BYTES, F64_BYTES};
6use super::strings::Strings;
7use super::{floats, lanes};
8use crate::binary::{
9    Decoder, MAX_DEPTH, VARINT_BITS, VARINT_LAST_SHIFT, VARINT_MAX_SHIFT, VARINT_MORE,
10    VARINT_SHIFT, read_varint, unzigzag,
11};
12use crate::schema::{Schema, read_ref_value};
13use crate::sync::Arc;
14
15/// A `T` out of a frame; `None` if the frame isn't one.
16#[must_use]
17pub fn read_frame<T: Schema>(bytes: &[u8]) -> Option<T> {
18    let mut u = Unpacker::open(bytes)?;
19    let v = T::unpack(&mut u)?;
20    u.at_end().then_some(v)
21}
22
23/// Reads a frame. Every read answers `None` for bytes that aren't what
24/// [`Packer`](super::Packer) writes there.
25#[derive(Debug)]
26pub struct Unpacker<'b> {
27    body: &'b [u8],
28    i: usize,
29    table: &'b [u8],
30    /// How far the table is read.
31    table_at: usize,
32    strings: Strings<'b>,
33    /// Strings handed out shared ([`Unpacker::shared_str`]), by table index:
34    /// a string a frame repeats is allocated once.
35    shared: Vec<Option<Arc<str>>>,
36    depth: usize,
37}
38
39impl<'b> Unpacker<'b> {
40    /// The frame's body and table found; the root value next.
41    #[must_use]
42    pub fn open(bytes: &'b [u8]) -> Option<Self> {
43        let (table_len, trailer) = read_varint_back(bytes)?;
44        let body_end = bytes.len().checked_sub(trailer)?.checked_sub(table_len)?;
45        Some(Unpacker {
46            body: &bytes[..body_end],
47            i: 0,
48            table: &bytes[body_end..bytes.len() - trailer],
49            table_at: 0,
50            strings: Strings::new(),
51            shared: Vec::new(),
52            depth: 0,
53        })
54    }
55
56    /// Whether the body is read through.
57    #[must_use]
58    pub fn at_end(&self) -> bool {
59        self.i == self.body.len()
60    }
61
62    fn take(&mut self, len: usize) -> Option<&'b [u8]> {
63        let end = self.i.checked_add(len).filter(|&e| e <= self.body.len())?;
64        let out = &self.body[self.i..end];
65        self.i = end;
66        Some(out)
67    }
68
69    /// An unsigned number ([`Packer::put_u`](super::Packer::put_u)).
70    pub fn uint(&mut self) -> Option<u64> {
71        read_varint(self.body, &mut self.i)
72    }
73
74    /// A signed number.
75    pub fn int(&mut self) -> Option<i64> {
76        self.uint().map(unzigzag)
77    }
78
79    /// A bool.
80    pub fn bool(&mut self) -> Option<bool> {
81        match self.take(1)? {
82            [0] => Some(false),
83            [1] => Some(true),
84            _ => None,
85        }
86    }
87
88    /// An `f64`.
89    pub fn f64(&mut self) -> Option<f64> {
90        Some(f64::from_le_bytes(self.take(F64_BYTES)?.try_into().ok()?))
91    }
92
93    /// An `f32`.
94    pub fn f32(&mut self) -> Option<f32> {
95        Some(f32::from_le_bytes(self.take(F32_BYTES)?.try_into().ok()?))
96    }
97
98    /// A count of things each at least `min` bytes: no bigger than what's left.
99    fn count_of(&mut self, min: usize) -> Option<usize> {
100        let n = usize::try_from(self.uint()?).ok()?;
101        (n <= (self.body.len() - self.i) / min.max(1)).then_some(n)
102    }
103
104    /// A list's count (every value takes a byte at least).
105    pub fn count(&mut self) -> Option<usize> {
106        self.count_of(1)
107    }
108
109    /// A list of `f64`s.
110    pub fn f64s(&mut self) -> Option<Vec<f64>> {
111        let (_, narrow, raw) = self.packed_floats(true)?;
112        Some(if narrow {
113            floats::widen(raw)
114        } else {
115            raw.as_chunks::<F64_BYTES>()
116                .0
117                .iter()
118                .map(|b| f64::from_le_bytes(*b))
119                .collect()
120        })
121    }
122
123    /// Bit-packed numbers ([`Packer::put_uints`](super::Packer::put_uints)), each made a `T` by `of`.
124    pub fn uints<T>(&mut self, of: impl Fn(u64) -> Option<T>) -> Option<Vec<T>> {
125        let (n, bits, raw) = self.packed_ints()?;
126        let mut out = Vec::with_capacity(n);
127        let mut ok = true;
128        lanes::unpack(raw, n, bits, |v| match of(v) {
129            Some(t) if ok => out.push(t),
130            _ => ok = false,
131        });
132        ok.then_some(out)
133    }
134
135    /// A list of `f32`s.
136    pub fn f32s(&mut self) -> Option<Vec<f32>> {
137        let (_, _, raw) = self.packed_floats(false)?;
138        Some(
139            raw.as_chunks::<F32_BYTES>()
140                .0
141                .iter()
142                .map(|b| f32::from_le_bytes(*b))
143                .collect(),
144        )
145    }
146
147    /// A packed list of ints, unread: how many, bits each, their bytes.
148    pub fn packed_ints(&mut self) -> Option<(usize, u32, &'b [u8])> {
149        let n = usize::try_from(self.uint()?).ok()?;
150        let bits = u32::from(*self.take(1)?.first()?);
151        if !(1..=u64::BITS).contains(&bits)
152            || n > (self.body.len() - self.i) * u8::BITS as usize / bits as usize
153        {
154            return None;
155        }
156        let bytes = self.take(lanes::packed_bytes(n, bits)?)?;
157        Some((n, bits, bytes))
158    }
159
160    /// A list of floats, unread: how many, whether they're `f32`s, their
161    /// bytes. A list of `f64` says whether it's narrowed (`wide`); one of
162    /// `f32` always is.
163    pub fn packed_floats(&mut self, wide: bool) -> Option<(usize, bool, &'b [u8])> {
164        let n = self.count()?;
165        let narrow = if wide { self.bool()? } else { true };
166        let width = if narrow { F32_BYTES } else { F64_BYTES };
167        Some((n, narrow, self.take(n.checked_mul(width)?)?))
168    }
169
170    /// A string, borrowed from the frame: the table read as far as it's needed.
171    pub fn str(&mut self) -> Option<&'b str> {
172        let i = self.string_index()?;
173        self.strings.get(i)
174    }
175
176    /// A string as an `Arc<str>`: one allocation for every time the frame has it.
177    pub fn shared_str(&mut self) -> Option<Arc<str>> {
178        let i = self.string_index()?;
179        if self.shared.len() <= i {
180            self.shared.resize(i + 1, None);
181        }
182        let s = self.strings.get(i)?;
183        Some(self.shared[i].get_or_insert_with(|| s.into()).clone())
184    }
185
186    /// A string's index, the table read up to it.
187    fn string_index(&mut self) -> Option<usize> {
188        let i = usize::try_from(self.uint()?).ok()?;
189        while self.strings.len() <= i {
190            let len = usize::try_from(read_varint(self.table, &mut self.table_at)?).ok()?;
191            let end = self
192                .table_at
193                .checked_add(len)
194                .filter(|&e| e <= self.table.len())?;
195            let s = crate::utf8::from_utf8(&self.table[self.table_at..end]).ok()?;
196            self.table_at = end;
197            self.strings.push(s);
198        }
199        Some(i)
200    }
201
202    /// Bytes, borrowed from the frame.
203    pub fn bytes(&mut self) -> Option<&'b [u8]> {
204        let len = self.count()?;
205        self.take(len)
206    }
207
208    /// A value held in its binary form ([`Packer::put_tagged`](super::Packer::put_tagged)).
209    pub fn tagged<T: Schema>(&mut self) -> Option<T> {
210        let bytes = self.bytes()?;
211        let mut d = Decoder::body(bytes);
212        if let Some(v) = T::decode(&mut d)
213            && d.at_end()
214        {
215            return Some(v);
216        }
217        let arena = Bump::new();
218        let mut any = Decoder::body(bytes);
219        let value = any.value_ref(&arena).ok()?;
220        any.at_end().then_some(())?;
221        read_ref_value(&value).ok()
222    }
223
224    /// A record starts: where it ends.
225    pub fn begin(&mut self) -> Option<usize> {
226        let len = self.count()?;
227        self.depth += 1;
228        (self.depth <= MAX_DEPTH).then_some(self.i + len)
229    }
230
231    /// A list ([`Packer::begin_list`](super::Packer::begin_list)): how many items, and where it ends.
232    pub fn list(&mut self) -> Option<(usize, usize)> {
233        let n = self.count()?;
234        let len = self.count()?;
235        Some((n, self.i + len))
236    }
237
238    /// Steps to `end` (a list's), past what's left of it.
239    pub fn skip_to(&mut self, end: usize) -> Option<()> {
240        (self.i <= end && end <= self.body.len()).then(|| self.i = end)
241    }
242
243    /// Where it's read up to.
244    #[must_use]
245    pub fn position(&self) -> usize {
246        self.i
247    }
248
249    /// A reader of the same frame from `at`, with a string cache of its own.
250    #[must_use]
251    pub fn fork(&self, at: usize) -> Unpacker<'b> {
252        Unpacker {
253            body: self.body,
254            i: at,
255            table: self.table,
256            table_at: 0,
257            strings: Strings::new(),
258            shared: Vec::new(),
259            depth: self.depth,
260        }
261    }
262
263    /// Whether the record ending at `end` has more fields.
264    #[must_use]
265    pub fn more(&self, end: usize) -> bool {
266        self.i < end
267    }
268
269    /// The record ending at `end` is done: fields this reader doesn't know skipped.
270    pub fn leave(&mut self, end: usize) -> Option<()> {
271        self.depth -= 1;
272        (self.i <= end).then(|| self.i = end)
273    }
274}
275
276/// A varint written backwards at the end of `bytes`: its value, and how
277/// many bytes it took.
278fn read_varint_back(bytes: &[u8]) -> Option<(usize, usize)> {
279    let mut n = 0u64;
280    let mut shift = 0;
281    for (taken, b) in bytes.iter().rev().enumerate() {
282        if shift >= VARINT_MAX_SHIFT {
283            return None;
284        }
285        let bits = u64::from(*b) & VARINT_BITS;
286        if shift == VARINT_LAST_SHIFT && bits > 1 {
287            return None;
288        }
289        n |= bits << shift;
290        if b & VARINT_MORE == 0 {
291            return Some((usize::try_from(n).ok()?, taken + 1));
292        }
293        shift += VARINT_SHIFT;
294    }
295    None
296}