extern crate alloc;
use alloc::ffi::CString;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use crate::{ErrorKind, OrtResult, Read, common::buf_read, libc, ort_error};
pub fn read<R: Read, const MAX_CHUNK_SIZE: usize>(
r: buf_read::OrtBufReader<R>,
) -> ChunkedIterator<R, MAX_CHUNK_SIZE> {
ChunkedIterator::new(r)
}
pub struct ChunkedIterator<R: Read, const MAX_CHUNK_SIZE: usize> {
r: buf_read::OrtBufReader<R>,
size_buf: String,
data_buf: Vec<u8>,
}
impl<R: Read, const MAX_CHUNK_SIZE: usize> ChunkedIterator<R, MAX_CHUNK_SIZE> {
fn new(r: buf_read::OrtBufReader<R>) -> ChunkedIterator<R, MAX_CHUNK_SIZE> {
ChunkedIterator {
r,
size_buf: String::with_capacity(16),
data_buf: Vec::with_capacity(MAX_CHUNK_SIZE),
}
}
pub fn next_chunk(&mut self) -> Option<OrtResult<&str>> {
let mut bytes_read = 0;
loop {
self.size_buf.clear();
match self.r.read_line(&mut self.size_buf) {
Ok(0) => {
return Some(Err(ort_error(ErrorKind::ChunkedEofInSize, "")));
}
Ok(_) => {}
Err(err) => {
err.debug_print();
return Some(Err(ort_error(ErrorKind::ChunkedSizeReadError, "")));
}
}
let size_str = self.size_buf.trim();
if size_str.is_empty() {
continue;
}
let size = match usize::from_str_radix(size_str, 16) {
Ok(n) => n,
Err(_err) => {
let c_s = CString::new("ERROR invalid chunked size: ".to_string() + size_str)
.unwrap();
unsafe {
libc::write(2, c_s.as_ptr().cast(), c_s.count_bytes());
}
return Some(Err(ort_error(ErrorKind::ChunkedInvalidSize, "")));
}
};
if size == 0 {
return None;
}
if bytes_read == 0 {
self.data_buf.clear();
}
self.data_buf.reserve_exact(size);
unsafe { self.data_buf.set_len(size + bytes_read) };
if let Err(_err) = self.r.read_exact(&mut self.data_buf[bytes_read..]) {
return Some(Err(ort_error(ErrorKind::ChunkedDataReadError, "")));
};
bytes_read += size;
let last_byte = self.data_buf[self.data_buf.len() - 1];
if (last_byte & 0b1000_0000) != 0 {
continue;
}
break;
}
Some(Ok(unsafe { str::from_utf8_unchecked(&self.data_buf) }))
}
}