use std::io::ErrorKind;
use std::io::Read;
use std::io::Result;
use std::io::Write;
use std::io::copy as copy_all;
use crate::Streams;
use crate::std_io::ext::internal::read_ext_impl;
use crate::util::SliceRange;
use crate::util::UncheckedSlice;
use crate::util::allocation_error;
use crate::util::try_reserve_string;
const DISCARD_BUFFER_SIZE: usize = 8 * 1024;
pub trait ReadExt: Read {
unsafe fn read_unchecked(&mut self, buffer: &mut [u8], start_index: usize, count: usize) -> Result<usize>;
unsafe fn read_exact_unchecked(&mut self, buffer: &mut [u8], start_index: usize, count: usize) -> Result<()>;
unsafe fn read_exact_or_eof_unchecked(
&mut self,
buffer: &mut [u8],
start_index: usize,
count: usize,
) -> Result<usize>;
fn read_exact_or_eof(&mut self, buffer: &mut [u8]) -> Result<usize>;
fn read_exact_array<const N: usize>(&mut self) -> Result<[u8; N]>;
fn read_exact_vec_limited(&mut self, len: usize, max_len: usize) -> Result<Vec<u8>>;
fn read_exact_vec_limited_into(&mut self, output: &mut Vec<u8>, len: usize, max_len: usize) -> Result<()>;
fn discard_exact_or_eof(&mut self, bytes: u64) -> Result<u64>;
fn copy_to(&mut self, writer: &mut dyn Write) -> Result<u64>;
fn copy_to_at_most(&mut self, writer: &mut dyn Write, max_bytes: u64) -> Result<u64>;
fn copy_to_end_limited(&mut self, writer: &mut dyn Write, max_bytes: u64) -> Result<u64>;
fn read_to_end_limited(&mut self, max_len: usize) -> Result<Vec<u8>>;
fn read_to_end_limited_into(&mut self, output: &mut Vec<u8>, max_len: usize) -> Result<usize>;
fn read_to_string_limited(&mut self, max_len: usize) -> Result<String>;
fn read_to_string_limited_into(&mut self, output: &mut String, max_len: usize) -> Result<usize>;
}
unsafe fn read_exact_or_eof_unchecked_impl(
reader: &mut dyn Read,
buffer: &mut [u8],
start_index: usize,
count: usize,
) -> Result<usize> {
debug_assert!(
SliceRange::range_fits(buffer.len(), start_index, count),
"unchecked read range exceeds buffer"
);
let mut total = 0;
while total < count {
let target = unsafe { UncheckedSlice::subslice_mut(buffer, start_index + total, count - total) };
match reader.read(target) {
Ok(0) => break,
Ok(read) => total += read,
Err(error) => {
if error.kind() == ErrorKind::Interrupted {
continue;
}
return Err(error);
}
}
}
Ok(total)
}
#[inline]
fn read_exact_array_impl<const N: usize>(reader: &mut dyn Read) -> Result<[u8; N]> {
let mut buffer = [0; N];
reader.read_exact(&mut buffer)?;
Ok(buffer)
}
#[inline]
fn read_exact_vec_limited_impl(reader: &mut dyn Read, len: usize, max_len: usize) -> Result<Vec<u8>> {
let mut output = Vec::new();
read_ext_impl::read_exact_vec_limited_into(reader, &mut output, len, max_len)?;
Ok(output)
}
fn discard_exact_or_eof_impl(reader: &mut dyn Read, bytes: u64) -> Result<u64> {
let mut buffer = [0; DISCARD_BUFFER_SIZE];
let mut remaining = bytes;
let mut discarded = 0;
while remaining > 0 {
let requested = remaining.min(DISCARD_BUFFER_SIZE as u64) as usize;
match reader.read(&mut buffer[..requested]) {
Ok(0) => break,
Ok(count) => {
let count = count as u64;
remaining -= count;
discarded += count;
}
Err(error) => {
if error.kind() == ErrorKind::Interrupted {
continue;
}
return Err(error);
}
}
}
Ok(discarded)
}
#[inline]
fn read_to_string_limited_impl(reader: &mut dyn Read, max_len: usize) -> Result<String> {
let bytes = read_ext_impl::read_to_end_limited(reader, max_len)?;
String::from_utf8(bytes).map_err(read_ext_impl::invalid_utf8_error)
}
#[inline]
fn read_to_string_limited_into_impl(reader: &mut dyn Read, output: &mut String, max_len: usize) -> Result<usize> {
let bytes = read_ext_impl::read_to_end_limited(reader, max_len)?;
let text = String::from_utf8(bytes).map_err(read_ext_impl::invalid_utf8_error)?;
let count = text.len();
try_reserve_string(output, count).map_err(allocation_error)?;
output.push_str(&text);
Ok(count)
}
impl<T> ReadExt for T
where
T: Read + ?Sized,
{
#[inline(always)]
unsafe fn read_unchecked(&mut self, buffer: &mut [u8], start_index: usize, count: usize) -> Result<usize> {
let target = unsafe { UncheckedSlice::subslice_mut(buffer, start_index, count) };
self.read(target)
}
#[inline(always)]
unsafe fn read_exact_or_eof_unchecked(
&mut self,
buffer: &mut [u8],
start_index: usize,
count: usize,
) -> Result<usize> {
let mut reader = self;
unsafe { read_exact_or_eof_unchecked_impl(&mut reader, buffer, start_index, count) }
}
#[inline(always)]
unsafe fn read_exact_unchecked(&mut self, buffer: &mut [u8], start_index: usize, count: usize) -> Result<()> {
let target = unsafe { UncheckedSlice::subslice_mut(buffer, start_index, count) };
self.read_exact(target)
}
#[inline(always)]
fn read_exact_or_eof(&mut self, buffer: &mut [u8]) -> Result<usize> {
let mut reader = self;
read_ext_impl::read_exact_or_eof(&mut reader, buffer)
}
#[inline(always)]
fn read_exact_array<const N: usize>(&mut self) -> Result<[u8; N]> {
let mut reader = self;
read_exact_array_impl(&mut reader)
}
#[inline(always)]
fn read_exact_vec_limited(&mut self, len: usize, max_len: usize) -> Result<Vec<u8>> {
let mut reader = self;
read_exact_vec_limited_impl(&mut reader, len, max_len)
}
#[inline(always)]
fn read_exact_vec_limited_into(&mut self, output: &mut Vec<u8>, len: usize, max_len: usize) -> Result<()> {
let mut reader = self;
read_ext_impl::read_exact_vec_limited_into(&mut reader, output, len, max_len)
}
#[inline(always)]
fn discard_exact_or_eof(&mut self, bytes: u64) -> Result<u64> {
let mut reader = self;
discard_exact_or_eof_impl(&mut reader, bytes)
}
#[inline(always)]
fn copy_to(&mut self, writer: &mut dyn Write) -> Result<u64> {
copy_all(self, writer)
}
#[inline(always)]
fn copy_to_at_most(&mut self, writer: &mut dyn Write, max_bytes: u64) -> Result<u64> {
let mut reader = self;
Streams::copy_at_most(&mut reader, writer, max_bytes)
}
#[inline(always)]
fn copy_to_end_limited(&mut self, writer: &mut dyn Write, max_bytes: u64) -> Result<u64> {
let mut reader = self;
Streams::copy_to_end_limited(&mut reader, writer, max_bytes)
}
#[inline(always)]
fn read_to_end_limited(&mut self, max_len: usize) -> Result<Vec<u8>> {
let mut reader = self;
read_ext_impl::read_to_end_limited(&mut reader, max_len)
}
#[inline(always)]
fn read_to_end_limited_into(&mut self, output: &mut Vec<u8>, max_len: usize) -> Result<usize> {
let mut reader = self;
read_ext_impl::read_to_end_limited_into(&mut reader, output, max_len)
}
#[inline(always)]
fn read_to_string_limited(&mut self, max_len: usize) -> Result<String> {
let mut reader = self;
read_to_string_limited_impl(&mut reader, max_len)
}
#[inline(always)]
fn read_to_string_limited_into(&mut self, output: &mut String, max_len: usize) -> Result<usize> {
let mut reader = self;
read_to_string_limited_into_impl(&mut reader, output, max_len)
}
}