use crate::encode::{FLAG_INLINE_SCHEMA, HEADER_LEN, MESSAGE_MAGIC};
use crate::error::{Error, Result};
use crate::message::Budget;
#[inline]
pub fn charge(budget: Option<&Budget>, bytes: u64) -> Result<()> {
match budget {
Some(b) => b.charge(bytes),
None => Ok(()),
}
}
pub fn read_str_budgeted<'b>(buf: &'b [u8], slot: u64, budget: Option<&Budget>) -> Result<&'b str> {
let off = read_u32(buf, slot)? as u64;
let len = read_u32(buf, off)? as u64;
charge(budget, 8 + len)?;
let bytes = read_slice(buf, off + 4, len)?;
std::str::from_utf8(bytes).map_err(|_| Error::BadUtf8)
}
pub fn read_bytes_budgeted<'b>(
buf: &'b [u8],
slot: u64,
budget: Option<&Budget>,
) -> Result<&'b [u8]> {
let off = read_u32(buf, slot)? as u64;
let len = read_u32(buf, off)? as u64;
charge(budget, 8 + len)?;
read_slice(buf, off + 4, len)
}
pub fn message_header(
schema_id: u128,
inline_schema: Option<&[u8]>,
capacity_hint: usize,
) -> Result<Vec<u8>> {
let mut buf = Vec::with_capacity(capacity_hint.max(HEADER_LEN));
buf.extend_from_slice(MESSAGE_MAGIC);
let flags: u16 = if inline_schema.is_some() {
FLAG_INLINE_SCHEMA
} else {
0
};
buf.extend_from_slice(&flags.to_le_bytes());
buf.extend_from_slice(&0u16.to_le_bytes()); buf.extend_from_slice(&schema_id.to_le_bytes()); buf.extend_from_slice(&0u32.to_le_bytes()); let schema_len = inline_schema.map(|s| s.len()).unwrap_or(0);
let schema_len = u32::try_from(schema_len).map_err(|_| Error::MessageTooLarge)?;
buf.extend_from_slice(&schema_len.to_le_bytes());
if let Some(s) = inline_schema {
buf.extend_from_slice(s);
}
while buf.len() % 8 != 0 {
buf.push(0);
}
Ok(buf)
}
pub fn finish_message(buf: &mut [u8], root_offset: u32) {
buf[24..28].copy_from_slice(&root_offset.to_le_bytes());
}
pub fn pos(buf: &[u8]) -> Result<u32> {
u32::try_from(buf.len()).map_err(|_| Error::MessageTooLarge)
}
pub fn pad_to(buf: &mut Vec<u8>, align: u32) -> Result<u32> {
let p = pos(buf)?;
let target = (p as u64 + align as u64 - 1) & !(align as u64 - 1);
let target = u32::try_from(target).map_err(|_| Error::MessageTooLarge)?;
buf.resize(buf.len() + (target - p) as usize, 0);
Ok(target)
}
pub fn alloc_block(buf: &mut Vec<u8>, size: u32, align: u32) -> Result<u32> {
let base = pad_to(buf, align)?;
buf.resize(buf.len() + size as usize, 0);
pos(buf)?;
Ok(base)
}
pub fn alloc_bytes(buf: &mut Vec<u8>, n: usize) -> Result<u32> {
let base = pos(buf)?;
buf.resize(buf.len().checked_add(n).ok_or(Error::MessageTooLarge)?, 0);
pos(buf)?;
Ok(base)
}
pub fn set_presence_bit(buf: &mut [u8], base: u32, bit: usize) -> Result<()> {
let at = base as usize + bit / 8;
let b = buf.get_mut(at).ok_or(Error::OutOfBounds)?;
*b |= 1 << (bit % 8);
Ok(())
}
fn put(buf: &mut [u8], at: u32, bytes: &[u8]) -> Result<()> {
let start = at as usize;
let end = start.checked_add(bytes.len()).ok_or(Error::OutOfBounds)?;
buf.get_mut(start..end)
.ok_or(Error::OutOfBounds)?
.copy_from_slice(bytes);
Ok(())
}
macro_rules! put_fns {
($($name:ident: $ty:ty),* $(,)?) => {$(
pub fn $name(buf: &mut [u8], at: u32, v: $ty) -> Result<()> {
put(buf, at, &v.to_le_bytes())
}
)*};
}
put_fns!(put_u16: u16, put_u32: u32, put_u64: u64, put_i16: i16, put_i32: i32, put_i64: i64, put_f32: f32, put_f64: f64);
pub fn put_u8(buf: &mut [u8], at: u32, v: u8) -> Result<()> {
put(buf, at, &[v])
}
pub fn put_i8(buf: &mut [u8], at: u32, v: i8) -> Result<()> {
put(buf, at, &[v as u8])
}
pub fn put_bool(buf: &mut [u8], at: u32, v: bool) -> Result<()> {
put(buf, at, &[v as u8])
}
pub fn patch_u32(buf: &mut [u8], at: u32, v: u32) -> Result<()> {
put_u32(buf, at, v)
}
pub fn write_blob(buf: &mut Vec<u8>, bytes: &[u8]) -> Result<u32> {
let len = u32::try_from(bytes.len()).map_err(|_| Error::MessageTooLarge)?;
let off = pad_to(buf, 4)?;
buf.extend_from_slice(&len.to_le_bytes());
buf.extend_from_slice(bytes);
pos(buf)?;
Ok(off)
}
pub fn begin_list(buf: &mut Vec<u8>, count: u32, elem_align: u32) -> Result<u32> {
let off = pad_to(buf, 4)?;
buf.extend_from_slice(&count.to_le_bytes());
pad_to(buf, elem_align)?;
Ok(off)
}
macro_rules! push_fns {
($($name:ident: $ty:ty),* $(,)?) => {$(
pub fn $name(buf: &mut Vec<u8>, v: $ty) {
buf.extend_from_slice(&v.to_le_bytes());
}
)*};
}
push_fns!(push_u16: u16, push_u32: u32, push_u64: u64, push_i16: i16, push_i32: i32, push_i64: i64, push_f32: f32, push_f64: f64);
pub fn push_u8(buf: &mut Vec<u8>, v: u8) {
buf.push(v);
}
pub fn push_i8(buf: &mut Vec<u8>, v: i8) {
buf.push(v as u8);
}
pub fn push_bool(buf: &mut Vec<u8>, v: bool) {
buf.push(v as u8);
}
macro_rules! push_slice_fns {
($($name:ident: $ty:ty = $n:literal),* $(,)?) => {$(
pub fn $name(buf: &mut Vec<u8>, vals: &[$ty]) {
let start = buf.len();
buf.resize(start + vals.len() * $n, 0);
for (chunk, v) in buf[start..].chunks_exact_mut($n).zip(vals) {
chunk.copy_from_slice(&v.to_le_bytes());
}
}
)*};
}
push_slice_fns!(
push_u16_slice: u16 = 2, push_u32_slice: u32 = 4, push_u64_slice: u64 = 8,
push_i16_slice: i16 = 2, push_i32_slice: i32 = 4, push_i64_slice: i64 = 8,
push_f32_slice: f32 = 4, push_f64_slice: f64 = 8,
);
pub fn push_u8_slice(buf: &mut Vec<u8>, vals: &[u8]) {
buf.extend_from_slice(vals);
}
pub fn push_i8_slice(buf: &mut Vec<u8>, vals: &[i8]) {
let start = buf.len();
buf.resize(start + vals.len(), 0);
for (b, &v) in buf[start..].iter_mut().zip(vals) {
*b = v as u8;
}
}
pub fn push_bool_slice(buf: &mut Vec<u8>, vals: &[bool]) {
let start = buf.len();
buf.resize(start + vals.len(), 0);
for (b, &v) in buf[start..].iter_mut().zip(vals) {
*b = v as u8;
}
}
pub fn read_slice(buf: &[u8], at: u64, len: u64) -> Result<&[u8]> {
let start = usize::try_from(at).map_err(|_| Error::OutOfBounds)?;
let len = usize::try_from(len).map_err(|_| Error::OutOfBounds)?;
let end = start.checked_add(len).ok_or(Error::OutOfBounds)?;
buf.get(start..end).ok_or(Error::OutOfBounds)
}
macro_rules! read_fns {
($($name:ident: $ty:ty = $n:literal),* $(,)?) => {$(
pub fn $name(buf: &[u8], at: u64) -> Result<$ty> {
Ok(<$ty>::from_le_bytes(read_slice(buf, at, $n)?.try_into().unwrap()))
}
)*};
}
read_fns!(
read_u16: u16 = 2, read_u32: u32 = 4, read_u64: u64 = 8,
read_i16: i16 = 2, read_i32: i32 = 4, read_i64: i64 = 8,
read_f32: f32 = 4, read_f64: f64 = 8,
);
pub fn read_u8(buf: &[u8], at: u64) -> Result<u8> {
Ok(read_slice(buf, at, 1)?[0])
}
pub fn read_i8(buf: &[u8], at: u64) -> Result<i8> {
Ok(read_u8(buf, at)? as i8)
}
pub fn read_bool(buf: &[u8], at: u64) -> Result<bool> {
Ok(read_u8(buf, at)? != 0)
}
pub fn read_str(buf: &[u8], slot: u64) -> Result<&str> {
let bytes = read_bytes(buf, slot)?;
std::str::from_utf8(bytes).map_err(|_| Error::BadUtf8)
}
pub fn read_bytes(buf: &[u8], slot: u64) -> Result<&[u8]> {
let off = read_u32(buf, slot)? as u64;
let len = read_u32(buf, off)? as u64;
read_slice(buf, off + 4, len)
}
pub fn list_header(buf: &[u8], slot: u64, elem_align: u32) -> Result<(u64, u32)> {
let off = read_u32(buf, slot)? as u64;
let count = read_u32(buf, off)?;
let a = elem_align as u64;
let elems = (off + 4 + a - 1) & !(a - 1);
Ok((elems, count))
}