use crate::error::{Error, Result};
use crate::ext::Complex;
use crate::header::*;
use crate::size::{
read_size, read_size_from_reader, size_encoded_len, write_size, write_size_to_writer,
};
#[cfg(target_endian = "little")]
use core::ptr;
use half::{bf16, f16};
use std::io::{Read, Write};
#[inline]
fn write_typed_array_header_numeric(out: &mut Vec<u8>, class: u8, byte_code: u8, len: usize) {
out.push(make_header(TYPE_TYPED_ARRAY, class, byte_code));
write_size(len as u64, out);
}
#[inline]
fn write_typed_array_header_bool(out: &mut Vec<u8>, len: usize) {
let header = ((ARRAY_BOOL_OR_STRING & 0b11) << 3) | (TYPE_TYPED_ARRAY & 0b111);
out.push(header);
write_size(len as u64, out);
}
#[inline]
fn write_typed_array_header_string(out: &mut Vec<u8>, len: usize) {
let header = ((1u8) << 5) | ((ARRAY_BOOL_OR_STRING & 0b11) << 3) | (TYPE_TYPED_ARRAY & 0b111);
out.push(header);
write_size(len as u64, out);
}
pub trait BeveTypedSlice: Sized {
const CLASS: u8;
const BYTE_CODE: u8;
const ELEM_SIZE: usize;
fn write_one_le(v: &Self, out: &mut Vec<u8>);
}
macro_rules! impl_beve_typed_int {
($t:ty, $class:expr, $code:expr) => {
impl BeveTypedSlice for $t {
const CLASS: u8 = $class;
const BYTE_CODE: u8 = $code;
const ELEM_SIZE: usize = core::mem::size_of::<$t>();
#[inline]
fn write_one_le(v: &Self, out: &mut Vec<u8>) {
out.extend_from_slice(&v.to_le_bytes());
}
}
};
}
impl_beve_typed_int!(i8, ARRAY_SIGNED, 0);
impl_beve_typed_int!(i16, ARRAY_SIGNED, 1);
impl_beve_typed_int!(i32, ARRAY_SIGNED, 2);
impl_beve_typed_int!(i64, ARRAY_SIGNED, 3);
impl_beve_typed_int!(i128, ARRAY_SIGNED, 4);
impl_beve_typed_int!(u8, ARRAY_UNSIGNED, 0);
impl_beve_typed_int!(u16, ARRAY_UNSIGNED, 1);
impl_beve_typed_int!(u32, ARRAY_UNSIGNED, 2);
impl_beve_typed_int!(u64, ARRAY_UNSIGNED, 3);
impl_beve_typed_int!(u128, ARRAY_UNSIGNED, 4);
impl BeveTypedSlice for f32 {
const CLASS: u8 = ARRAY_FLOAT;
const BYTE_CODE: u8 = 2; const ELEM_SIZE: usize = core::mem::size_of::<f32>();
#[inline]
fn write_one_le(v: &Self, out: &mut Vec<u8>) {
out.extend_from_slice(&v.to_le_bytes());
}
}
impl BeveTypedSlice for f64 {
const CLASS: u8 = ARRAY_FLOAT;
const BYTE_CODE: u8 = 3; const ELEM_SIZE: usize = core::mem::size_of::<f64>();
#[inline]
fn write_one_le(v: &Self, out: &mut Vec<u8>) {
out.extend_from_slice(&v.to_le_bytes());
}
}
impl BeveTypedSlice for bf16 {
const CLASS: u8 = ARRAY_FLOAT;
const BYTE_CODE: u8 = 0; const ELEM_SIZE: usize = core::mem::size_of::<bf16>();
#[inline]
fn write_one_le(v: &Self, out: &mut Vec<u8>) {
out.extend_from_slice(&v.to_le_bytes());
}
}
impl BeveTypedSlice for f16 {
const CLASS: u8 = ARRAY_FLOAT;
const BYTE_CODE: u8 = 1; const ELEM_SIZE: usize = core::mem::size_of::<f16>();
#[inline]
fn write_one_le(v: &Self, out: &mut Vec<u8>) {
out.extend_from_slice(&v.to_le_bytes());
}
}
pub fn write_typed_slice<T: BeveTypedSlice>(out: &mut Vec<u8>, slice: &[T]) {
let payload = core::mem::size_of_val(slice);
out.reserve(1 + 8 + payload);
write_typed_array_header_numeric(out, T::CLASS, T::BYTE_CODE, slice.len());
if slice.is_empty() {
return;
}
#[cfg(target_endian = "little")]
{
let start = out.len();
unsafe {
let dst = out.as_mut_ptr().add(start);
ptr::copy_nonoverlapping(slice.as_ptr() as *const u8, dst, payload);
out.set_len(start + payload);
}
}
#[cfg(not(target_endian = "little"))]
{
for v in slice {
T::write_one_le(v, out);
}
}
}
pub fn to_vec_typed_slice<T: BeveTypedSlice>(slice: &[T]) -> Vec<u8> {
let payload = slice.len() * T::ELEM_SIZE;
let mut out = Vec::with_capacity(1 + 8 + payload);
write_typed_slice(&mut out, slice);
out
}
pub fn to_writer_typed_slice<W: Write, T: BeveTypedSlice>(mut w: W, slice: &[T]) -> Result<()> {
w.write_all(&[make_header(TYPE_TYPED_ARRAY, T::CLASS, T::BYTE_CODE)])?;
write_size_to_writer(&mut w, slice.len() as u64)?;
if slice.is_empty() {
return Ok(());
}
#[cfg(target_endian = "little")]
{
let bytes = unsafe {
core::slice::from_raw_parts(slice.as_ptr() as *const u8, core::mem::size_of_val(slice))
};
w.write_all(bytes)?;
}
#[cfg(not(target_endian = "little"))]
{
let mut scratch = Vec::with_capacity(T::ELEM_SIZE);
for v in slice {
scratch.clear();
T::write_one_le(v, &mut scratch);
w.write_all(&scratch)?;
}
}
Ok(())
}
pub fn typed_slice_size<T: BeveTypedSlice>(slice: &[T]) -> u64 {
1 + size_encoded_len(slice.len() as u64) as u64 + core::mem::size_of_val(slice) as u64 }
pub fn read_typed_slice<T: BeveTypedSlice>(input: &[u8]) -> Result<Vec<T>> {
let mut pos = 0usize;
let header = *input.get(pos).ok_or(Error::Eof)?;
pos += 1;
if parse_type(header) != TYPE_TYPED_ARRAY {
return Err(Error::InvalidType("not a typed array"));
}
if parse_subtype(header) != T::CLASS || parse_byte_count_code(header) != T::BYTE_CODE {
return Err(Error::Mismatch("typed array element type does not match T"));
}
let len = read_size(input, &mut pos)? as usize;
let payload = len.checked_mul(T::ELEM_SIZE).ok_or(Error::InvalidSize)?;
if input.len() - pos < payload {
return Err(Error::Eof);
}
let src = &input[pos..pos + payload];
let mut out: Vec<T> = Vec::with_capacity(len);
if len != 0 {
#[cfg(target_endian = "little")]
{
unsafe {
ptr::copy_nonoverlapping(src.as_ptr(), out.as_mut_ptr() as *mut u8, payload);
out.set_len(len);
}
}
#[cfg(not(target_endian = "little"))]
{
unsafe {
let dst = out.as_mut_ptr() as *mut u8;
core::ptr::copy_nonoverlapping(src.as_ptr(), dst, payload);
let bytes = core::slice::from_raw_parts_mut(dst, payload);
for elem in bytes.chunks_exact_mut(T::ELEM_SIZE) {
elem.reverse();
}
out.set_len(len);
}
}
}
Ok(out)
}
pub fn write_bool_slice(out: &mut Vec<u8>, slice: &[bool]) {
write_typed_array_header_bool(out, slice.len());
let mut acc: u8 = 0;
let mut idx: u8 = 0; for &b in slice {
if b {
acc |= 1 << idx;
}
idx += 1;
if idx == 8 {
out.push(acc);
acc = 0;
idx = 0;
}
}
if idx != 0 {
out.push(acc);
}
}
pub fn to_vec_bool_slice(slice: &[bool]) -> Vec<u8> {
let mut out = Vec::new();
write_bool_slice(&mut out, slice);
out
}
pub fn write_str_slice(out: &mut Vec<u8>, slice: &[&str]) {
write_typed_array_header_string(out, slice.len());
for s in slice {
write_size(s.len() as u64, out);
out.extend_from_slice(s.as_bytes());
}
}
pub fn to_vec_str_slice(slice: &[&str]) -> Vec<u8> {
let mut out = Vec::new();
write_str_slice(&mut out, slice);
out
}
pub fn write_string_slice(out: &mut Vec<u8>, slice: &[String]) {
write_typed_array_header_string(out, slice.len());
for s in slice {
write_size(s.len() as u64, out);
out.extend_from_slice(s.as_bytes());
}
}
pub fn to_vec_string_slice(slice: &[String]) -> Vec<u8> {
let mut out = Vec::new();
write_string_slice(&mut out, slice);
out
}
#[inline]
fn write_complex_header(out: &mut Vec<u8>, is_array: bool, class: u8, byte_code: u8) {
out.push(((EXT_COMPLEX & 0x1f) << 3) | (TYPE_EXTENSION & 0b111));
let h = ((byte_code & 0b111) << 5) | ((class & 0b11) << 3) | (if is_array { 1 } else { 0 });
out.push(h);
}
pub fn write_complex_slice<T: BeveTypedSlice>(out: &mut Vec<u8>, slice: &[Complex<T>]) {
let payload = core::mem::size_of_val(slice);
out.reserve(2 + 8 + payload);
write_complex_header(out, true, T::CLASS, T::BYTE_CODE);
write_size(slice.len() as u64, out);
if !slice.is_empty() {
#[cfg(target_endian = "little")]
{
let start = out.len();
unsafe {
let dst = out.as_mut_ptr().add(start);
ptr::copy_nonoverlapping(slice.as_ptr() as *const u8, dst, payload);
out.set_len(start + payload);
}
}
#[cfg(not(target_endian = "little"))]
{
for c in slice {
T::write_one_le(&c.re, out);
T::write_one_le(&c.im, out);
}
}
}
}
pub fn to_vec_complex_slice<T: BeveTypedSlice>(slice: &[Complex<T>]) -> Vec<u8> {
let payload = core::mem::size_of_val(slice);
let mut out = Vec::with_capacity(2 + 8 + payload);
write_complex_slice(&mut out, slice);
out
}
pub fn to_writer_complex_slice<W: Write, T: BeveTypedSlice>(
mut w: W,
slice: &[Complex<T>],
) -> Result<()> {
w.write_all(&[
((EXT_COMPLEX & 0x1f) << 3) | (TYPE_EXTENSION & 0b111),
((T::BYTE_CODE & 0b111) << 5) | ((T::CLASS & 0b11) << 3) | 1,
])?;
write_size_to_writer(&mut w, slice.len() as u64)?;
if slice.is_empty() {
return Ok(());
}
#[cfg(target_endian = "little")]
{
let bytes = unsafe {
core::slice::from_raw_parts(slice.as_ptr() as *const u8, core::mem::size_of_val(slice))
};
w.write_all(bytes)?;
}
#[cfg(not(target_endian = "little"))]
{
let mut scratch = Vec::with_capacity(T::ELEM_SIZE);
for c in slice {
scratch.clear();
T::write_one_le(&c.re, &mut scratch);
T::write_one_le(&c.im, &mut scratch);
w.write_all(&scratch)?;
}
}
Ok(())
}
pub fn complex_slice_size<T: BeveTypedSlice>(slice: &[Complex<T>]) -> u64 {
2 + size_encoded_len(slice.len() as u64) as u64 + core::mem::size_of_val(slice) as u64 }
pub fn read_complex_slice<T: BeveTypedSlice>(input: &[u8]) -> Result<Vec<Complex<T>>> {
let mut pos = 0usize;
let header = *input.get(pos).ok_or(Error::Eof)?;
pos += 1;
if parse_type(header) != TYPE_EXTENSION || parse_extension_id(header) != EXT_COMPLEX {
return Err(Error::InvalidType("not a complex extension"));
}
let ch = *input.get(pos).ok_or(Error::Eof)?;
pos += 1;
let is_array = (ch & 0x01) != 0;
let class = (ch >> 3) & 0x03;
let byte_code = (ch >> 5) & 0x07;
if !is_array {
return Err(Error::InvalidType(
"expected a complex array, found a scalar complex",
));
}
if class != T::CLASS || byte_code != T::BYTE_CODE {
return Err(Error::Mismatch("complex element type does not match T"));
}
let len = read_size(input, &mut pos)? as usize;
let payload = len
.checked_mul(2 * T::ELEM_SIZE)
.ok_or(Error::InvalidSize)?;
if input.len() - pos < payload {
return Err(Error::Eof);
}
let src = &input[pos..pos + payload];
let mut out: Vec<Complex<T>> = Vec::with_capacity(len);
if len != 0 {
#[cfg(target_endian = "little")]
{
unsafe {
ptr::copy_nonoverlapping(src.as_ptr(), out.as_mut_ptr() as *mut u8, payload);
out.set_len(len);
}
}
#[cfg(not(target_endian = "little"))]
{
unsafe {
let dst = out.as_mut_ptr() as *mut u8;
core::ptr::copy_nonoverlapping(src.as_ptr(), dst, payload);
let bytes = core::slice::from_raw_parts_mut(dst, payload);
for elem in bytes.chunks_exact_mut(T::ELEM_SIZE) {
elem.reverse();
}
out.set_len(len);
}
}
}
Ok(out)
}
pub fn read_typed_slice_from_reader<T: BeveTypedSlice, R: Read>(mut reader: R) -> Result<Vec<T>> {
let mut hdr = [0u8; 1];
reader.read_exact(&mut hdr).map_err(|_| Error::Eof)?;
let header = hdr[0];
if parse_type(header) != TYPE_TYPED_ARRAY {
return Err(Error::InvalidType("not a typed array"));
}
if parse_subtype(header) != T::CLASS || parse_byte_count_code(header) != T::BYTE_CODE {
return Err(Error::Mismatch("typed array element type does not match T"));
}
let len = read_size_from_reader(&mut reader)? as usize;
read_le_payload::<T, R>(&mut reader, len, T::ELEM_SIZE)
}
pub fn read_complex_slice_from_reader<T: BeveTypedSlice, R: Read>(
mut reader: R,
) -> Result<Vec<Complex<T>>> {
let mut hdr = [0u8; 2];
reader.read_exact(&mut hdr).map_err(|_| Error::Eof)?;
if parse_type(hdr[0]) != TYPE_EXTENSION || parse_extension_id(hdr[0]) != EXT_COMPLEX {
return Err(Error::InvalidType("not a complex extension"));
}
let ch = hdr[1];
let is_array = (ch & 0x01) != 0;
let class = (ch >> 3) & 0x03;
let byte_code = (ch >> 5) & 0x07;
if !is_array {
return Err(Error::InvalidType(
"expected a complex array, found a scalar complex",
));
}
if class != T::CLASS || byte_code != T::BYTE_CODE {
return Err(Error::Mismatch("complex element type does not match T"));
}
let len = read_size_from_reader(&mut reader)? as usize;
read_le_payload::<Complex<T>, R>(&mut reader, len, T::ELEM_SIZE)
}
fn read_le_payload<E, R: Read>(reader: &mut R, len: usize, scalar_size: usize) -> Result<Vec<E>> {
let mut out: Vec<E> = Vec::new();
if len == 0 {
return Ok(out);
}
let elem = core::mem::size_of::<E>();
len.checked_mul(elem).ok_or(Error::InvalidSize)?;
const STEP_BYTES: usize = 8 * 1024 * 1024;
let step_elems = (STEP_BYTES / elem.max(1)).max(1);
let mut filled = 0usize;
while filled < len {
let take = step_elems.min(len - filled);
out.try_reserve(take)
.map_err(|_| Error::Message("allocation failed for slice"))?;
unsafe {
let start = out.as_mut_ptr().add(filled) as *mut u8;
let nbytes = take * elem;
core::ptr::write_bytes(start, 0, nbytes);
let dst = core::slice::from_raw_parts_mut(start, nbytes);
reader.read_exact(dst).map_err(Error::from)?;
out.set_len(filled + take);
}
filled += take;
}
#[cfg(not(target_endian = "little"))]
{
let payload = len * elem;
let bytes =
unsafe { core::slice::from_raw_parts_mut(out.as_mut_ptr() as *mut u8, payload) };
for s in bytes.chunks_exact_mut(scalar_size) {
s.reverse();
}
}
#[cfg(target_endian = "little")]
let _ = scalar_size;
Ok(out)
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MatrixLayoutFast {
Right,
Left,
}
pub fn to_vec_matrix_f64(layout: MatrixLayoutFast, extents: &[u64], data: &[f64]) -> Vec<u8> {
let mut out = Vec::new();
out.push(((EXT_MATRICES & 0x1f) << 3) | (TYPE_EXTENSION & 0b111));
let mh = match layout {
MatrixLayoutFast::Right => 0u8,
MatrixLayoutFast::Left => 1u8,
};
out.push(mh);
write_typed_array_header_numeric(&mut out, ARRAY_UNSIGNED, 3, extents.len());
for &e in extents {
out.extend_from_slice(&e.to_le_bytes());
}
write_typed_array_header_numeric(&mut out, ARRAY_FLOAT, 3, data.len());
for &v in data {
out.extend_from_slice(&v.to_le_bytes());
}
out
}