use crate::error::{Error, Result};
use crate::header::*;
use crate::size::read_size;
use half::{bf16, f16};
use simdutf8::basic::from_utf8;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum TypedArrayClass {
Float,
Signed,
Unsigned,
Bool,
String,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct ObjectHeader {
pub(crate) key_type: u8,
pub(crate) byte_code: u8,
pub(crate) len: usize,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct TypedArrayHeader {
pub(crate) class: TypedArrayClass,
pub(crate) byte_code: u8,
pub(crate) len: usize,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct ComplexHeader {
pub(crate) class: u8,
pub(crate) is_array: bool,
pub(crate) byte_code: u8,
}
pub(crate) struct Reader<'a> {
input: &'a [u8],
pos: usize,
}
impl<'a> Reader<'a> {
pub(crate) fn new(input: &'a [u8]) -> Self {
Self { input, pos: 0 }
}
pub(crate) fn is_finished(&self) -> bool {
self.pos >= self.input.len()
}
pub(crate) fn read_header(&mut self) -> Result<u8> {
self.read_byte()
}
pub(crate) fn read_size(&mut self) -> Result<usize> {
read_size(self.input, &mut self.pos)
.map_err(|_| Error::InvalidSize)?
.try_into()
.map_err(|_| Error::InvalidSize)
}
pub(crate) fn read_string(&mut self) -> Result<&'a str> {
let len = self.read_size()?;
let bytes = self.read_exact(len)?;
from_utf8(bytes).map_err(|_| Error::InvalidType("invalid utf-8"))
}
pub(crate) fn read_signed(&mut self, byte_code: u8) -> Result<i128> {
let nbytes = byte_count_to_bytes(byte_code)?;
let bytes = self.read_exact(nbytes)?;
let mut buf = [0u8; 16];
buf[..nbytes].copy_from_slice(bytes);
if nbytes < 16 && (bytes[nbytes - 1] & 0x80) != 0 {
for b in &mut buf[nbytes..] {
*b = 0xff;
}
}
Ok(i128::from_le_bytes(buf))
}
pub(crate) fn read_unsigned(&mut self, byte_code: u8) -> Result<u128> {
let nbytes = byte_count_to_bytes(byte_code)?;
let bytes = self.read_exact(nbytes)?;
let mut buf = [0u8; 16];
buf[..nbytes].copy_from_slice(bytes);
Ok(u128::from_le_bytes(buf))
}
pub(crate) fn read_f32(&mut self) -> Result<f32> {
Ok(f32::from_le_bytes(self.read_exact(4)?.try_into().unwrap()))
}
pub(crate) fn read_f64(&mut self) -> Result<f64> {
Ok(f64::from_le_bytes(self.read_exact(8)?.try_into().unwrap()))
}
pub(crate) fn read_f16_as_f32(&mut self) -> Result<f32> {
let bytes = self.read_exact(2)?;
Ok(f16::from_bits(u16::from_le_bytes([bytes[0], bytes[1]])).to_f32())
}
pub(crate) fn read_bf16_as_f32(&mut self) -> Result<f32> {
let bytes = self.read_exact(2)?;
Ok(bf16::from_bits(u16::from_le_bytes([bytes[0], bytes[1]])).to_f32())
}
pub(crate) fn read_object_header(&mut self, header: u8) -> Result<ObjectHeader> {
if parse_type(header) != TYPE_OBJECT {
return Err(Error::InvalidType("expected object"));
}
Ok(ObjectHeader {
key_type: parse_subtype(header),
byte_code: parse_byte_count_code(header),
len: self.read_size()?,
})
}
pub(crate) fn read_typed_array_header(&mut self, header: u8) -> Result<TypedArrayHeader> {
if parse_type(header) != TYPE_TYPED_ARRAY {
return Err(Error::InvalidType("expected typed array"));
}
let class = match parse_subtype(header) {
ARRAY_FLOAT => TypedArrayClass::Float,
ARRAY_SIGNED => TypedArrayClass::Signed,
ARRAY_UNSIGNED => TypedArrayClass::Unsigned,
ARRAY_BOOL_OR_STRING => {
if (header & 0b0010_0000) == 0 {
TypedArrayClass::Bool
} else {
TypedArrayClass::String
}
}
_ => return Err(Error::InvalidHeader(header)),
};
Ok(TypedArrayHeader {
class,
byte_code: parse_byte_count_code(header),
len: self.read_size()?,
})
}
pub(crate) fn read_complex_header(&mut self) -> Result<ComplexHeader> {
let header = self.read_byte()?;
Ok(ComplexHeader {
class: (header >> 3) & 0b11,
is_array: (header & 0x01) != 0,
byte_code: (header >> 5) & 0b111,
})
}
pub(crate) fn read_matrix_extents(&mut self) -> Result<Vec<usize>> {
let header = self.read_header()?;
let info = self.read_typed_array_header(header)?;
let len = info.len;
let mut extents = Vec::with_capacity(len);
match info.class {
TypedArrayClass::Unsigned => {
for _ in 0..len {
let value = self.read_unsigned(info.byte_code)?;
extents.push(value.try_into().map_err(|_| Error::InvalidSize)?);
}
}
TypedArrayClass::Signed => {
for _ in 0..len {
let value = self.read_signed(info.byte_code)?;
if value < 0 {
return Err(Error::InvalidSize);
}
extents.push(value.try_into().map_err(|_| Error::InvalidSize)?);
}
}
_ => {
return Err(Error::InvalidType(
"matrix extents must be integer typed arrays",
));
}
}
Ok(extents)
}
pub(crate) fn read_exact(&mut self, n: usize) -> Result<&'a [u8]> {
if self.remaining() < n {
return Err(Error::Eof);
}
let start = self.pos;
self.pos += n;
Ok(&self.input[start..start + n])
}
pub(crate) fn remaining(&self) -> usize {
self.input.len().saturating_sub(self.pos)
}
fn peek_byte(&self) -> Result<u8> {
self.input.get(self.pos).copied().ok_or(Error::Eof)
}
fn read_byte(&mut self) -> Result<u8> {
let byte = self.peek_byte()?;
self.pos += 1;
Ok(byte)
}
}
pub(crate) fn byte_count_to_bytes(code: u8) -> Result<usize> {
match code {
0 => Ok(1),
1 => Ok(2),
2 => Ok(4),
3 => Ok(8),
4 => Ok(16),
_ => Err(Error::Unsupported("byte width > 16 not supported")),
}
}