use crate::io::{BufRead, Error, ErrorKind, Read, Result, Seek, SeekFrom, Write};
use core::{cmp, fmt};
pub struct BufReader<R, const S: usize> {
inner: R,
buf: [u8; S],
pos: usize,
cap: usize,
}
impl<R: Read, const S: usize> BufReader<R, S> {
pub fn new(inner: R) -> BufReader<R, S> {
BufReader {
inner,
buf: [0; S],
pos: 0,
cap: 0,
}
}
}
impl<R, const S: usize> BufReader<R, S> {
pub fn get_ref(&self) -> &R {
&self.inner
}
pub fn get_mut(&mut self) -> &mut R {
&mut self.inner
}
pub fn buffer(&self) -> &[u8] {
&self.buf[self.pos..self.cap]
}
pub fn capacity(&self) -> usize {
S
}
pub fn into_inner(self) -> R {
self.inner
}
#[inline]
fn discard_buffer(&mut self) {
self.pos = 0;
self.cap = 0;
}
}
impl<R: Read, const S: usize> Read for BufReader<R, S> {
fn read(&mut self, buf: &mut [u8]) -> Result<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: Read, const S: usize> BufRead for BufReader<R, S> {
fn fill_buf(&mut self) -> Result<&[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, const S: usize> fmt::Debug for BufReader<R, S>
where
R: fmt::Debug,
{
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("BufReader")
.field("reader", &self.inner)
.field("buffer", &format_args!("{}/{}", self.cap - self.pos, S))
.finish()
}
}
impl<R: Seek, const S: usize> Seek for BufReader<R, S> {
fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
let result: u64;
if let SeekFrom::Current(n) = pos {
let remainder = (self.cap - self.pos) as i64;
if let Some(offset) = n.checked_sub(remainder) {
result = self.inner.seek(SeekFrom::Current(offset))?;
} else {
self.inner.seek(SeekFrom::Current(-remainder))?;
self.discard_buffer();
result = self.inner.seek(SeekFrom::Current(n))?;
}
} else {
result = self.inner.seek(pos)?;
}
self.discard_buffer();
Ok(result)
}
}
pub struct BufWriter<W: Write, const S: usize> {
inner: Option<W>,
buf: [u8; S],
len: usize,
panicked: bool,
}
#[derive(Debug)]
pub struct IntoInnerError<W>(W, Error);
impl<W, const S: usize> BufWriter<W, S>
where
W: Write,
{
pub fn new(inner: W) -> BufWriter<W, S> {
BufWriter {
inner: Some(inner),
buf: [0; S],
len: 0,
panicked: false,
}
}
fn flush_buf(&mut self) -> Result<()> {
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(Error::new(
ErrorKind::WriteZero,
"failed to write the buffered data",
));
}
Ok(n) => guard.consume(n),
Err(ref e) if e.kind() == ErrorKind::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());
(&mut 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) -> core::result::Result<W, IntoInnerError<BufWriter<W, S>>> {
match self.flush_buf() {
Err(e) => Err(IntoInnerError(self, e)),
Ok(()) => Ok(self.inner.take().unwrap()),
}
}
}
impl<W: Write, const S: usize> Write for BufWriter<W, S> {
fn write(&mut self, buf: &[u8]) -> Result<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]) -> Result<()> {
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) -> Result<()> {
self.flush_buf().and_then(|()| self.get_mut().flush())
}
}
impl<W: Write, const S: usize> fmt::Debug for BufWriter<W, S>
where
W: fmt::Debug,
{
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("BufWriter")
.field("writer", &self.inner.as_ref().unwrap())
.field("buffer", &format_args!("{}/{}", self.buf.len(), S))
.finish()
}
}
impl<W: Write + Seek, const S: usize> Seek for BufWriter<W, S> {
fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
self.flush_buf()?;
self.get_mut().seek(pos)
}
}
impl<W: Write, const S: usize> Drop for BufWriter<W, S> {
fn drop(&mut self) {
if self.inner.is_some() && !self.panicked {
let _r = self.flush_buf();
}
}
}
impl<W> IntoInnerError<W> {
pub fn error(&self) -> &Error {
&self.1
}
pub fn into_inner(self) -> W {
self.0
}
}
impl<W> From<IntoInnerError<W>> for Error {
fn from(iie: IntoInnerError<W>) -> Error {
iie.1
}
}
impl<W> fmt::Display for IntoInnerError<W> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.error().fmt(f)
}
}
#[derive(Debug)]
pub(super) struct LineWriterShim<'a, W: Write, const S: usize> {
buffer: &'a mut BufWriter<W, S>,
}
impl<'a, W: Write, const S: usize> LineWriterShim<'a, W, S> {
pub fn new(buffer: &'a mut BufWriter<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) -> Result<()> {
match self.buffered().last().copied() {
Some(b'\n') => self.buffer.flush_buf(),
_ => Ok(()),
}
}
}
impl<'a, W: Write, const S: usize> Write for LineWriterShim<'a, W, S> {
fn write(&mut self, buf: &[u8]) -> Result<usize> {
let newline_idx = match memchr::memrchr(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 memchr::memrchr(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) -> Result<()> {
self.buffer.flush()
}
fn write_all(&mut self, buf: &[u8]) -> Result<()> {
match memchr::memrchr(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)
}
}
}
}
pub struct LineWriter<W: Write, const S: usize> {
inner: BufWriter<W, S>,
}
impl<W: Write, const S: usize> LineWriter<W, S> {
pub fn new(inner: W) -> LineWriter<W, S> {
LineWriter {
inner: BufWriter::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) -> core::result::Result<W, IntoInnerError<LineWriter<W, S>>> {
self.inner
.into_inner()
.map_err(|IntoInnerError(buf, e)| IntoInnerError(LineWriter { inner: buf }, e))
}
}
impl<W: Write, const S: usize> Write for LineWriter<W, S> {
fn write(&mut self, buf: &[u8]) -> Result<usize> {
LineWriterShim::new(&mut self.inner).write(buf)
}
fn flush(&mut self) -> Result<()> {
self.inner.flush()
}
fn write_all(&mut self, buf: &[u8]) -> Result<()> {
LineWriterShim::new(&mut self.inner).write_all(buf)
}
fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> Result<()> {
LineWriterShim::new(&mut self.inner).write_fmt(fmt)
}
}
impl<W: Write, const S: usize> fmt::Debug for LineWriter<W, S>
where
W: fmt::Debug,
{
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("LineWriter")
.field("writer", &self.inner.inner)
.field("buffer", &format_args!("{}/{}", self.inner.len, S))
.finish()
}
}