bitcode/
consume.rs

1use crate::coder::Result;
2use crate::error::{err, error};
3
4/// Attempts to claim `bytes` bytes out of `input`.
5pub fn consume_bytes<'a>(input: &mut &'a [u8], bytes: usize) -> Result<&'a [u8]> {
6    if bytes > input.len() {
7        return err("EOF");
8    }
9    let (bytes, remaining) = input.split_at(bytes);
10    *input = remaining;
11    Ok(bytes)
12}
13
14/// Attempts to claim one byte out of `input`.
15pub fn consume_byte(input: &mut &[u8]) -> Result<u8> {
16    Ok(consume_bytes(input, 1)?[0])
17}
18
19/// Like `consume_bytes` but consumes `[u8; N]` instead of `u8`.
20pub fn consume_byte_arrays<'a, const N: usize>(
21    input: &mut &'a [u8],
22    length: usize,
23) -> Result<&'a [[u8; N]]> {
24    // Avoid * overflow by using / instead.
25    if input.len() / N < length {
26        return err("EOF");
27    }
28
29    // Safety: input.len() >= mid since we've checked it above.
30    let mid = length * N;
31    let (bytes, remaining) = unsafe { (input.get_unchecked(..mid), input.get_unchecked(mid..)) };
32
33    *input = remaining;
34    Ok(bytemuck::cast_slice(bytes))
35}
36
37/// Check if `input` is empty or return error.
38pub fn expect_eof(input: &[u8]) -> Result<()> {
39    if cfg!(not(fuzzing)) && !input.is_empty() {
40        err("Expected EOF")
41    } else {
42        Ok(())
43    }
44}
45
46/// Returns `Ok(length * x)` if it does not overflow.
47pub fn mul_length(length: usize, x: usize) -> Result<usize> {
48    length
49        .checked_mul(x)
50        .ok_or_else(|| error("length overflow"))
51}