use crate::{
ByteSearch, IoBufRead, IoError, IoErrorKind, IoRead, IoResult, IoSeek, IoSeekFrom, IoWrite,
};
use ::core::{cmp, fmt};
#[doc = crate::_tags!(io)]
pub struct IoBufReader<R, const S: usize> {
inner: R,
buf: [u8; S],
pos: usize,
cap: usize,
}
impl<R: IoRead, const S: usize> IoBufReader<R, S> {
pub const fn new(inner: R) -> IoBufReader<R, S> {
IoBufReader { inner, buf: [0; S], pos: 0, cap: 0 }
}
}
#[rustfmt::skip]
impl<R, const S: usize> IoBufReader<R, S> {
pub const fn get_ref(&self) -> &R { &self.inner }
pub const fn get_mut(&mut self) -> &mut R { &mut self.inner }
pub const fn buffer(&self) -> &[u8] { crate::Slice::range(&self.buf, self.pos, self.cap) }
pub const fn capacity(&self) -> usize { S }
pub fn into_inner(self) -> R { self.inner }
const fn discard_buffer(&mut self) { self.pos = 0; self.cap = 0; }
}
impl<R: IoRead, const S: usize> IoRead for IoBufReader<R, S> {
fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
if self.pos == self.cap && buf.len() >= S {
self.discard_buffer();
return self.inner.read(buf);
}
let nread = {
let mut rem = self.fill_buf()?;
rem.read(buf)?
};
self.consume(nread);
Ok(nread)
}
}
impl<R: IoRead, const S: usize> IoBufRead for IoBufReader<R, S> {
fn fill_buf(&mut self) -> IoResult<&[u8]> {
if self.pos >= self.cap {
debug_assert!(self.pos == self.cap);
self.cap = self.inner.read(&mut self.buf)?;
self.pos = 0;
}
Ok(&self.buf[self.pos..self.cap])
}
fn consume(&mut self, amt: usize) {
self.pos = cmp::min(self.pos + amt, self.cap);
}
}
impl<R: fmt::Debug, const S: usize> fmt::Debug for IoBufReader<R, S> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("IoBufReader")
.field("reader", &self.inner)
.field("buffer", &format_args!("{}/{}", self.cap - self.pos, S))
.finish()
}
}
impl<R: IoSeek, const S: usize> IoSeek for IoBufReader<R, S> {
fn seek(&mut self, pos: IoSeekFrom) -> IoResult<u64> {
let result: u64;
if let IoSeekFrom::Current(n) = pos {
let remainder = (self.cap - self.pos) as i64;
if let Some(offset) = n.checked_sub(remainder) {
result = self.inner.seek(IoSeekFrom::Current(offset))?;
} else {
self.inner.seek(IoSeekFrom::Current(-remainder))?;
self.discard_buffer();
result = self.inner.seek(IoSeekFrom::Current(n))?;
}
} else {
result = self.inner.seek(pos)?;
}
self.discard_buffer();
Ok(result)
}
}
#[doc = crate::_tags!(io)]
pub struct IoBufWriter<W: IoWrite, const S: usize> {
inner: Option<W>,
buf: [u8; S],
len: usize,
panicked: bool,
}
#[doc = crate::_tags!(io error)]
#[derive(Debug)]
pub struct IoIntoInnerError<W>(W, IoError);
impl<W: IoWrite, const S: usize> IoBufWriter<W, S> {
pub fn new(inner: W) -> IoBufWriter<W, S> {
IoBufWriter {
inner: Some(inner),
buf: [0; S],
len: 0,
panicked: false,
}
}
fn flush_buf(&mut self) -> IoResult<()> {
struct BufGuard<'a, const S: usize> {
buffer: &'a mut [u8; S],
written: usize,
}
impl<'a, const S: usize> BufGuard<'a, S> {
fn new(buffer: &'a mut [u8; S]) -> Self {
Self { buffer, written: 0 }
}
fn remaining(&self) -> &[u8] {
&self.buffer[self.written..]
}
fn consume(&mut self, amt: usize) {
self.written += amt;
}
fn done(&self) -> bool {
self.written >= self.buffer.len()
}
}
impl<const S: usize> Drop for BufGuard<'_, S> {
fn drop(&mut self) {
if self.written > 0 {
let mut new_buf = [0; S];
new_buf.copy_from_slice(&self.buffer[self.written..]);
*self.buffer = new_buf;
}
}
}
let mut guard = BufGuard::new(&mut self.buf);
let inner = self.inner.as_mut().unwrap();
while !guard.done() {
self.panicked = true;
let r = inner.write(guard.remaining());
self.panicked = false;
match r {
Ok(0) => {
return Err(IoError::new(
IoErrorKind::WriteZero,
"failed to write the buffered data",
));
}
Ok(n) => guard.consume(n),
Err(ref e) if e.kind() == IoErrorKind::Interrupted => {}
Err(e) => return Err(e),
}
}
Ok(())
}
fn write_to_buf(&mut self, buf: &[u8]) -> usize {
let available = S - self.len;
let amt_to_buffer = available.min(buf.len());
self.buf[available..].copy_from_slice(&buf[..amt_to_buffer]);
self.len += amt_to_buffer;
amt_to_buffer
}
pub fn get_ref(&self) -> &W {
self.inner.as_ref().unwrap()
}
pub fn get_mut(&mut self) -> &mut W {
self.inner.as_mut().unwrap()
}
pub fn buffer(&self) -> &[u8] {
&self.buf
}
pub fn capacity(&self) -> usize {
S
}
pub fn into_inner(mut self) -> crate::Result<W, IoIntoInnerError<IoBufWriter<W, S>>> {
match self.flush_buf() {
Err(e) => Err(IoIntoInnerError(self, e)),
Ok(()) => Ok(self.inner.take().unwrap()),
}
}
}
impl<W: IoWrite, const S: usize> IoWrite for IoBufWriter<W, S> {
fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
if self.len + buf.len() > S {
self.flush_buf()?;
}
if buf.len() >= S {
self.panicked = true;
let r = self.get_mut().write(buf);
self.panicked = false;
r
} else {
self.buf.copy_from_slice(buf);
Ok(buf.len())
}
}
fn write_all(&mut self, buf: &[u8]) -> IoResult<()> {
if self.len + buf.len() > S {
self.flush_buf()?;
}
if buf.len() >= S {
self.panicked = true;
let r = self.get_mut().write_all(buf);
self.panicked = false;
r
} else {
self.buf.copy_from_slice(buf);
Ok(())
}
}
fn flush(&mut self) -> IoResult<()> {
self.flush_buf().and_then(|()| self.get_mut().flush())
}
}
impl<W: IoWrite + fmt::Debug, const S: usize> fmt::Debug for IoBufWriter<W, S> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("IoBufWriter")
.field("writer", &self.inner.as_ref().unwrap())
.field("buffer", &format_args!("{}/{}", self.buf.len(), S))
.finish()
}
}
impl<W: IoWrite + IoSeek, const S: usize> IoSeek for IoBufWriter<W, S> {
fn seek(&mut self, pos: IoSeekFrom) -> IoResult<u64> {
self.flush_buf()?;
self.get_mut().seek(pos)
}
}
impl<W: IoWrite, const S: usize> Drop for IoBufWriter<W, S> {
fn drop(&mut self) {
if self.inner.is_some() && !self.panicked {
let _r = self.flush_buf();
}
}
}
impl<W> IoIntoInnerError<W> {
pub fn error(&self) -> &IoError {
&self.1
}
pub fn into_inner(self) -> W {
self.0
}
}
impl<W> From<IoIntoInnerError<W>> for IoError {
fn from(iie: IoIntoInnerError<W>) -> IoError {
iie.1
}
}
impl<W> fmt::Display for IoIntoInnerError<W> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.error().fmt(f)
}
}
#[doc = crate::_tags!(io)]
#[derive(Debug)]
pub(super) struct LineWriterShim<'a, W: IoWrite, const S: usize> {
buffer: &'a mut IoBufWriter<W, S>,
}
impl<'a, W: IoWrite, const S: usize> LineWriterShim<'a, W, S> {
pub fn new(buffer: &'a mut IoBufWriter<W, S>) -> Self {
Self { buffer }
}
fn inner_mut(&mut self) -> &mut W {
self.buffer.get_mut()
}
fn buffered(&self) -> &[u8] {
self.buffer.buffer()
}
fn flush_if_completed_line(&mut self) -> IoResult<()> {
match self.buffered().last().copied() {
Some(b'\n') => self.buffer.flush_buf(),
_ => Ok(()),
}
}
}
impl<W: IoWrite, const S: usize> IoWrite for LineWriterShim<'_, W, S> {
fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
let newline_idx = match ByteSearch::last1(b'\n', buf) {
None => {
self.flush_if_completed_line()?;
return self.buffer.write(buf);
}
Some(newline_idx) => newline_idx + 1,
};
self.buffer.flush_buf()?;
let lines = &buf[..newline_idx];
let flushed = self.inner_mut().write(lines)?;
if flushed == 0 {
return Ok(0);
}
let tail = if flushed >= newline_idx {
&buf[flushed..]
} else if newline_idx - flushed <= self.buffer.capacity() {
&buf[flushed..newline_idx]
} else {
let scan_area = &buf[flushed..];
let scan_area = &scan_area[..self.buffer.capacity()];
match ByteSearch::last1(b'\n', scan_area) {
Some(newline_idx) => &scan_area[..newline_idx + 1],
None => scan_area,
}
};
let buffered = self.buffer.write_to_buf(tail);
Ok(flushed + buffered)
}
fn flush(&mut self) -> IoResult<()> {
self.buffer.flush()
}
fn write_all(&mut self, buf: &[u8]) -> IoResult<()> {
match ByteSearch::last1(b'\n', buf) {
None => {
self.flush_if_completed_line()?;
self.buffer.write_all(buf)
}
Some(newline_idx) => {
let (lines, tail) = buf.split_at(newline_idx + 1);
if self.buffered().is_empty() {
self.inner_mut().write_all(lines)?;
} else {
self.buffer.write_all(lines)?;
self.buffer.flush_buf()?;
}
self.buffer.write_all(tail)
}
}
}
}
#[doc = crate::_tags!(io)]
pub struct IoLineWriter<W: IoWrite, const S: usize> {
inner: IoBufWriter<W, S>,
}
impl<W: IoWrite, const S: usize> IoLineWriter<W, S> {
pub fn new(inner: W) -> IoLineWriter<W, S> {
IoLineWriter { inner: IoBufWriter::new(inner) }
}
pub fn get_ref(&self) -> &W {
self.inner.get_ref()
}
pub fn get_mut(&mut self) -> &mut W {
self.inner.get_mut()
}
pub fn into_inner(self) -> Result<W, IoIntoInnerError<IoLineWriter<W, S>>> {
self.inner
.into_inner()
.map_err(|IoIntoInnerError(buf, e)| IoIntoInnerError(IoLineWriter { inner: buf }, e))
}
}
impl<W: IoWrite, const S: usize> IoWrite for IoLineWriter<W, S> {
fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
LineWriterShim::new(&mut self.inner).write(buf)
}
fn flush(&mut self) -> IoResult<()> {
self.inner.flush()
}
fn write_all(&mut self, buf: &[u8]) -> IoResult<()> {
LineWriterShim::new(&mut self.inner).write_all(buf)
}
fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> IoResult<()> {
LineWriterShim::new(&mut self.inner).write_fmt(fmt)
}
}
impl<W: IoWrite + fmt::Debug, const S: usize> fmt::Debug for IoLineWriter<W, S> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("IoLineWriter")
.field("writer", &self.inner.inner)
.field("buffer", &format_args!("{}/{}", self.inner.len, S))
.finish()
}
}