use crate::{Element, Error, ListView};
const WORD: usize = size_of::<u64>();
const WIDTH_MASK: u8 = 0b11;
const KIND_SHIFT: u32 = 2;
const KIND_MASK: u8 = 0b11;
const COUNT_SHIFT: u32 = 4;
const COUNT_FOLLOWS: usize = 15;
const TABLE: u8 = 0;
const TAGGED: u8 = 1;
const PACKED: u8 = 2;
const MAX_FRAME: usize = u32::MAX as usize;
const MAX_DEPTH: u32 = 64;
const MAX_VARINT: usize = 10;
#[derive(Clone, Copy)]
pub struct Packing {
pub width: u8,
pub total: bool,
}
#[inline]
fn read_uint(bytes: &[u8], at: usize, width: usize) -> Result<usize, Error> {
let value = if at <= bytes.len().saturating_sub(WORD) {
let word = u64::from_le_bytes(
bytes[at..at + WORD]
.try_into()
.expect("the slice is exactly a word"),
);
word & (u64::MAX >> (8 * (WORD - width)))
} else {
let end = at.checked_add(width).ok_or(Error::UnexpectedEof)?;
let slice = bytes.get(at..end).ok_or(Error::UnexpectedEof)?;
let mut word = [0; WORD];
word[..width].copy_from_slice(slice);
u64::from_le_bytes(word)
};
usize::try_from(value).map_err(|_| Error::UnexpectedEof)
}
fn write_uint(out: &mut [u8], value: usize, width: usize) {
out[..width].copy_from_slice(&(value as u64).to_le_bytes()[..width]);
}
fn varint_len(value: u32) -> usize {
let bits = (u32::BITS - value.leading_zeros()).max(1) as usize;
bits.div_ceil(7)
}
fn put_varint(out: &mut [u8], mut value: u32) {
let mut at = 0;
loop {
let byte = (value & 0x7f) as u8;
value >>= 7;
if value == 0 {
out[at] = byte;
return;
}
out[at] = byte | 0x80;
at += 1;
}
}
fn push_varint(out: &mut Vec<u8>, mut value: u64) {
loop {
let byte = (value & 0x7f) as u8;
value >>= 7;
if value == 0 {
out.push(byte);
return;
}
out.push(byte | 0x80);
}
}
fn read_varint(bytes: &[u8], at: usize) -> Result<(u64, usize), Error> {
let mut value = 0;
let mut shift = 0;
for size in 1..=MAX_VARINT {
let byte = *bytes
.get(at.checked_add(size - 1).ok_or(Error::UnexpectedEof)?)
.ok_or(Error::UnexpectedEof)?;
if shift == u64::BITS - 1 && byte & 0x7f > 1 {
return Err(Error::UnexpectedEof);
}
value |= u64::from(byte & 0x7f) << shift;
if byte & 0x80 == 0 {
return Ok((value, size));
}
shift += 7;
}
Err(Error::UnexpectedEof)
}
fn prefix(width: usize, follows: bool, kind: u8) -> usize {
1 + width + if follows { width } else { 0 } + usize::from(kind == PACKED)
}
fn fits(value: usize, width: usize) -> bool {
width == WORD || value >> (8 * width) == 0
}
fn header_size(width: usize, count: usize, kind: u8, tag: u32) -> usize {
prefix(width, count >= COUNT_FOLLOWS, kind)
+ match kind {
PACKED => 0,
TAGGED => count * width + varint_len(tag),
_ => count * width,
}
}
pub struct Writer<'out> {
output: &'out mut Vec<u8>,
}
#[must_use = "the frame is only completed by Writer::end_frame"]
pub struct FrameMark {
start: usize,
count: usize,
kind: u8,
tag: u32,
stride: usize,
}
macro_rules! write_scalars {
($($name:ident: $ty:ty,)*) => {
$(
pub fn $name(&mut self, value: $ty) {
self.output.extend_from_slice(&value.to_le_bytes());
}
)*
};
}
impl<'out> Writer<'out> {
pub fn new(output: &'out mut Vec<u8>) -> Self {
Self { output }
}
pub fn begin_frame(&mut self, count: usize) -> FrameMark {
self.begin(count, TABLE, 0, 0)
}
pub fn begin_variant(&mut self, tag: u32, count: usize) -> FrameMark {
self.begin(count, TAGGED, tag, 0)
}
pub fn begin_list<E: Element + ?Sized>(&mut self, count: usize) -> FrameMark {
match E::PACKING {
Some(packing) => self.begin(count, PACKED, 0, usize::from(packing.width)),
None => self.begin_frame(count),
}
}
fn begin(&mut self, count: usize, kind: u8, tag: u32, stride: usize) -> FrameMark {
assert!(count <= MAX_FRAME, "the frame is too large to encode");
let start = self.output.len();
let inline = count.min(COUNT_FOLLOWS) as u8;
self.output
.push((kind << KIND_SHIFT) | (inline << COUNT_SHIFT));
self.output
.resize(start + header_size(1, count, kind, tag), 0);
FrameMark {
start,
count,
kind,
tag,
stride,
}
}
#[inline]
fn width(&self, mark: &FrameMark) -> usize {
1 << (self.output[mark.start] & WIDTH_MASK)
}
#[cold]
fn widen_to_fit(&mut self, mark: &FrameMark) -> usize {
loop {
self.widen(mark);
let width = self.width(mark);
if fits(self.output.len() - mark.start, width) {
return width;
}
}
}
fn widen(&mut self, mark: &FrameMark) {
let width = self.width(mark);
let wider = width * 2;
let header = header_size(width, mark.count, mark.kind, mark.tag);
let shift = header_size(wider, mark.count, mark.kind, mark.tag) - header;
let from = mark.start + header;
let data = self.output.len() - from;
self.output.resize(self.output.len() + shift, 0);
self.output.copy_within(from..from + data, from + shift);
if mark.kind != PACKED {
let follows = mark.count >= COUNT_FOLLOWS;
let table = mark.start + prefix(width, follows, mark.kind);
let wide = mark.start + prefix(wider, follows, mark.kind);
for index in (0..mark.count).rev() {
let offset = read_uint(self.output, table + index * width, width)
.expect("the table is inside the header");
let offset = if offset == 0 { 0 } else { offset + shift };
write_uint(&mut self.output[wide + index * wider..], offset, wider);
}
}
let ctrl = self.output[mark.start] & !WIDTH_MASK;
self.output[mark.start] = ctrl | wider.trailing_zeros() as u8;
}
#[inline]
pub fn begin_entry(&mut self, mark: &FrameMark, index: usize) {
if mark.kind == PACKED {
return;
}
let mut width = self.width(mark);
if !fits(self.output.len() - mark.start, width) {
width = self.widen_to_fit(mark);
}
let at = self.output.len() - mark.start;
let table = mark.start + prefix(width, mark.count >= COUNT_FOLLOWS, mark.kind);
write_uint(&mut self.output[table + index * width..], at, width);
}
pub fn end_frame(&mut self, mark: FrameMark) {
let mut width = self.width(&mark);
if !fits(self.output.len() - mark.start, width) {
width = self.widen_to_fit(&mark);
}
let length = self.output.len() - mark.start;
assert!(length <= MAX_FRAME, "the frame is too large to encode");
let mut at = mark.start + 1;
write_uint(&mut self.output[at..], length, width);
at += width;
if mark.count >= COUNT_FOLLOWS {
write_uint(&mut self.output[at..], mark.count, width);
at += width;
}
match mark.kind {
PACKED => self.output[at] = mark.stride as u8,
TAGGED => put_varint(&mut self.output[at + mark.count * width..], mark.tag),
_ => {}
}
}
pub fn write_str(&mut self, value: &str) {
self.write_bytes(value.as_bytes());
}
pub fn write_bytes(&mut self, value: &[u8]) {
push_varint(self.output, value.len() as u64);
self.output.extend_from_slice(value);
}
pub fn write_bool(&mut self, value: bool) {
self.output.push(u8::from(value));
}
pub fn write_char(&mut self, value: char) {
self.write_u32(value as u32);
}
pub fn write_number(&mut self, value: u32) {
push_varint(self.output, u64::from(value));
}
write_scalars! {
write_u8: u8,
write_u16: u16,
write_u32: u32,
write_u64: u64,
write_u128: u128,
write_i8: i8,
write_i16: i16,
write_i32: i32,
write_i64: i64,
write_i128: i128,
write_f32: f32,
write_f64: f64,
}
}
#[derive(Clone, Copy)]
pub(crate) struct Frame<'buf> {
bytes: &'buf [u8],
table: u32,
count: u32,
data: u32,
step: u8,
kind: u8,
}
fn parse(bytes: &[u8]) -> Result<Frame<'_>, Error> {
let ctrl = *bytes.first().ok_or(Error::UnexpectedEof)?;
let width = 1 << (ctrl & WIDTH_MASK);
let kind = (ctrl >> KIND_SHIFT) & KIND_MASK;
let length = read_uint(bytes, 1, width)?;
let frame = bytes.get(..length).ok_or(Error::UnexpectedEof)?;
if kind == KIND_MASK || length > MAX_FRAME {
return Err(Error::InvalidFrame);
}
let inline = usize::from(ctrl >> COUNT_SHIFT);
let follows = inline == COUNT_FOLLOWS;
let count = if follows {
read_uint(frame, 1 + width, width)?
} else {
inline
};
let table = prefix(width, follows, kind);
let step = if kind == PACKED {
usize::from(*frame.get(table - 1).ok_or(Error::UnexpectedEof)?)
} else {
width
};
let entries = count
.checked_mul(step)
.and_then(|span| span.checked_add(table))
.ok_or(Error::UnexpectedEof)?;
if entries > length {
return Err(Error::InvalidFrame);
}
let data = if kind == TAGGED {
let (_, size) = read_varint(frame, entries)?;
entries + size
} else {
entries
};
Ok(Frame {
bytes: frame,
table: table as u32,
count: count as u32,
data: data as u32,
step: step as u8,
kind,
})
}
impl<'buf> Frame<'buf> {
pub(crate) fn read(bytes: &'buf [u8]) -> Result<Self, Error> {
parse(bytes)
}
pub(crate) fn trusted(bytes: &'buf [u8]) -> Self {
parse(bytes).expect("a frame's header was checked when it was read")
}
pub(crate) fn bytes(&self) -> &'buf [u8] {
self.bytes
}
pub(crate) fn count(&self) -> usize {
self.count as usize
}
pub(crate) fn kind(&self) -> u8 {
self.kind
}
pub(crate) fn stride(&self) -> usize {
usize::from(self.step)
}
pub(crate) fn elements(&self) -> &'buf [u8] {
&self.bytes[self.table as usize..]
}
fn tag_at(&self) -> usize {
self.table as usize + self.count as usize * usize::from(self.step)
}
pub(crate) fn tag(&self) -> Result<u32, Error> {
let (tag, _) = read_varint(self.bytes, self.tag_at())?;
u32::try_from(tag).map_err(|_| Error::InvalidFrame)
}
#[inline]
pub(crate) fn entry(&self, index: usize) -> Result<Option<&'buf [u8]>, Error> {
if index >= self.count as usize {
return Ok(None);
}
let step = usize::from(self.step);
let at = self.table as usize + index * step;
let offset = read_uint(self.bytes, at, step)?;
if offset == 0 {
return Ok(None);
}
if offset < self.data as usize {
return Err(Error::UnexpectedEof);
}
self.bytes
.get(offset..)
.map(Some)
.ok_or(Error::UnexpectedEof)
}
}
#[derive(Clone, Copy)]
pub struct Message<'buf> {
frame: Frame<'buf>,
depth: u32,
validate: bool,
}
macro_rules! read_scalars {
($($name:ident: $ty:ty,)*) => {
$(
pub fn $name(&self, slot: u32) -> Result<$ty, Error> {
let bytes = self.slot(slot)?;
let value = bytes.get(..size_of::<$ty>()).ok_or(Error::UnexpectedEof)?;
Ok(<$ty>::from_le_bytes(
value.try_into().expect("the slice is exactly sized"),
))
}
)*
};
}
impl<'buf> Message<'buf> {
pub(crate) fn root(bytes: &'buf [u8]) -> Result<Self, Error> {
let frame = Frame::read(bytes)?;
if frame.bytes().len() != bytes.len() {
return Err(Error::TrailingBytes);
}
Ok(Self {
frame,
depth: 0,
validate: true,
})
}
#[doc(hidden)]
pub fn trusted(bytes: &'buf [u8]) -> Self {
Self {
frame: Frame::trusted(bytes),
depth: 0,
validate: false,
}
}
pub fn bytes(&self) -> &'buf [u8] {
self.frame.bytes()
}
pub fn validates(&self) -> bool {
self.validate
}
fn slot(&self, slot: u32) -> Result<&'buf [u8], Error> {
self.frame.entry(slot as usize)?.ok_or(Error::MissingField)
}
pub fn read_str(&self, slot: u32) -> Result<&'buf str, Error> {
str::from_utf8(self.read_bytes(slot)?).map_err(|_| Error::InvalidUtf8)
}
pub fn read_bytes(&self, slot: u32) -> Result<&'buf [u8], Error> {
let bytes = self.slot(slot)?;
let (length, size) = read_varint(bytes, 0)?;
let length = usize::try_from(length).map_err(|_| Error::UnexpectedEof)?;
let end = size.checked_add(length).ok_or(Error::UnexpectedEof)?;
bytes.get(size..end).ok_or(Error::UnexpectedEof)
}
pub fn read_bool(&self, slot: u32) -> Result<bool, Error> {
match self.read_u8(slot)? {
0 => Ok(false),
1 => Ok(true),
_ => Err(Error::InvalidBool),
}
}
pub fn read_char(&self, slot: u32) -> Result<char, Error> {
char::from_u32(self.read_u32(slot)?).ok_or(Error::InvalidChar)
}
pub fn read_number(&self, slot: u32) -> Result<u32, Error> {
let (value, _) = read_varint(self.slot(slot)?, 0)?;
u32::try_from(value).map_err(|_| Error::UnknownVariant)
}
pub fn is_present(&self, slot: u32) -> Result<bool, Error> {
Ok(self.frame.entry(slot as usize)?.is_some())
}
pub fn expect_message(&self) -> Result<(), Error> {
if self.frame.kind() != TABLE {
return Err(Error::InvalidFrame);
}
Ok(())
}
pub fn read_tag(&self) -> Result<u32, Error> {
if self.frame.kind() != TAGGED {
return Err(Error::InvalidFrame);
}
self.frame.tag()
}
pub fn read_message(&self, slot: u32) -> Result<Message<'buf>, Error> {
self.read_nested(slot, TABLE)
}
pub fn read_variant(&self, slot: u32) -> Result<Message<'buf>, Error> {
self.read_nested(slot, TAGGED)
}
fn read_nested(&self, slot: u32, kind: u8) -> Result<Message<'buf>, Error> {
if self.depth == MAX_DEPTH {
return Err(Error::RecursionLimit);
}
let frame = Frame::read(self.slot(slot)?)?;
if frame.kind() != kind {
return Err(Error::InvalidFrame);
}
Ok(Self {
frame,
depth: self.depth + 1,
validate: self.validate,
})
}
pub fn read_list<E: Element + ?Sized>(&self, slot: u32) -> Result<ListView<'buf, E>, Error> {
if self.depth == MAX_DEPTH {
return Err(Error::RecursionLimit);
}
let frame = Frame::read(self.slot(slot)?)?;
let (kind, checked) = match E::PACKING {
Some(packing) => {
if frame.stride() != usize::from(packing.width) {
return Err(Error::InvalidFrame);
}
(PACKED, packing.total)
}
None => (TABLE, false),
};
if frame.kind() != kind {
return Err(Error::InvalidFrame);
}
let list = ListView::new(frame);
if self.validate && !checked {
list.validate(self.depth)?;
}
Ok(list)
}
pub(crate) fn nested(frame: Frame<'buf>, depth: u32, validate: bool) -> Self {
Self {
frame,
depth,
validate,
}
}
read_scalars! {
read_u8: u8,
read_u16: u16,
read_u32: u32,
read_u64: u64,
read_u128: u128,
read_i8: i8,
read_i16: i16,
read_i32: i32,
read_i64: i64,
read_i128: i128,
read_f32: f32,
read_f64: f64,
}
}