#[cfg(test)]
mod tests;
use core::prelude::v1::*;
use crate::io::prelude::*;
use core::cmp;
use core::fmt;
use crate::io::{
self, Error, ErrorKind, Initializer, IoSlice, IoSliceMut, SeekFrom, DEFAULT_BUF_SIZE,
};
use crate::io::memchr;
pub struct BufReader<R> {
inner: R,
buf: Box<[u8]>,
pos: usize,
cap: usize,
}
impl<R: Read> BufReader<R> {
pub fn new(inner: R) -> BufReader<R> {
BufReader::with_capacity(DEFAULT_BUF_SIZE, inner)
}
pub fn with_capacity(capacity: usize, inner: R) -> BufReader<R> {
unsafe {
let mut buffer = Vec::with_capacity(capacity);
buffer.set_len(capacity);
inner.initializer().initialize(&mut buffer);
BufReader { inner, buf: buffer.into_boxed_slice(), pos: 0, cap: 0 }
}
}
}
impl<R> BufReader<R> {
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 {
self.buf.len()
}
pub fn into_inner(self) -> R {
self.inner
}
#[inline]
fn discard_buffer(&mut self) {
self.pos = 0;
self.cap = 0;
}
}
impl<R: Seek> BufReader<R> {
pub fn seek_relative(&mut self, offset: i64) -> io::Result<()> {
let pos = self.pos as u64;
if offset < 0 {
if let Some(new_pos) = pos.checked_sub((-offset) as u64) {
self.pos = new_pos as usize;
return Ok(());
}
} else {
if let Some(new_pos) = pos.checked_add(offset as u64) {
if new_pos <= self.cap as u64 {
self.pos = new_pos as usize;
return Ok(());
}
}
}
self.seek(SeekFrom::Current(offset)).map(drop)
}
}
impl<R: Read> Read for BufReader<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if self.pos == self.cap && buf.len() >= self.buf.len() {
self.discard_buffer();
return self.inner.read(buf);
}
let nread = {
let mut rem = self.fill_buf()?;
rem.read(buf)?
};
self.consume(nread);
Ok(nread)
}
fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
let total_len = bufs.iter().map(|b| b.len()).sum::<usize>();
if self.pos == self.cap && total_len >= self.buf.len() {
self.discard_buffer();
return self.inner.read_vectored(bufs);
}
let nread = {
let mut rem = self.fill_buf()?;
rem.read_vectored(bufs)?
};
self.consume(nread);
Ok(nread)
}
fn is_read_vectored(&self) -> bool {
self.inner.is_read_vectored()
}
unsafe fn initializer(&self) -> Initializer {
self.inner.initializer()
}
}
impl<R: Read> BufRead for BufReader<R> {
fn fill_buf(&mut self) -> io::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> fmt::Debug for BufReader<R>
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, self.buf.len()))
.finish()
}
}
impl<R: Seek> Seek for BufReader<R> {
fn seek(&mut self, pos: SeekFrom) -> io::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> {
inner: Option<W>,
buf: Vec<u8>,
panicked: bool,
}
#[derive(Debug)]
pub struct IntoInnerError<W>(W, Error);
impl<W: Write> BufWriter<W> {
pub fn new(inner: W) -> BufWriter<W> {
BufWriter::with_capacity(DEFAULT_BUF_SIZE, inner)
}
pub fn with_capacity(capacity: usize, inner: W) -> BufWriter<W> {
BufWriter { inner: Some(inner), buf: Vec::with_capacity(capacity), panicked: false }
}
fn flush_buf(&mut self) -> io::Result<()> {
struct BufGuard<'a> {
buffer: &'a mut Vec<u8>,
written: usize,
}
impl<'a> BufGuard<'a> {
fn new(buffer: &'a mut Vec<u8>) -> 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 Drop for BufGuard<'_> {
fn drop(&mut self) {
if self.written > 0 {
self.buffer.drain(..self.written);
}
}
}
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() == io::ErrorKind::Interrupted => {}
Err(e) => return Err(e),
}
}
Ok(())
}
fn write_to_buf(&mut self, buf: &[u8]) -> usize {
let available = self.buf.capacity() - self.buf.len();
let amt_to_buffer = available.min(buf.len());
self.buf.extend_from_slice(&buf[..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 {
self.buf.capacity()
}
pub fn into_inner(mut self) -> Result<W, IntoInnerError<BufWriter<W>>> {
match self.flush_buf() {
Err(e) => Err(IntoInnerError(self, e)),
Ok(()) => Ok(self.inner.take().unwrap()),
}
}
}
impl<W: Write> Write for BufWriter<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if self.buf.len() + buf.len() > self.buf.capacity() {
self.flush_buf()?;
}
if buf.len() >= self.buf.capacity() {
self.panicked = true;
let r = self.get_mut().write(buf);
self.panicked = false;
r
} else {
self.buf.extend_from_slice(buf);
Ok(buf.len())
}
}
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
if self.buf.len() + buf.len() > self.buf.capacity() {
self.flush_buf()?;
}
if buf.len() >= self.buf.capacity() {
self.panicked = true;
let r = self.get_mut().write_all(buf);
self.panicked = false;
r
} else {
self.buf.extend_from_slice(buf);
Ok(())
}
}
fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
let total_len = bufs.iter().map(|b| b.len()).sum::<usize>();
if self.buf.len() + total_len > self.buf.capacity() {
self.flush_buf()?;
}
if total_len >= self.buf.capacity() {
self.panicked = true;
let r = self.get_mut().write_vectored(bufs);
self.panicked = false;
r
} else {
bufs.iter().for_each(|b| self.buf.extend_from_slice(b));
Ok(total_len)
}
}
fn is_write_vectored(&self) -> bool {
self.get_ref().is_write_vectored()
}
fn flush(&mut self) -> io::Result<()> {
self.flush_buf().and_then(|()| self.get_mut().flush())
}
}
impl<W: Write> fmt::Debug for BufWriter<W>
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(), self.buf.capacity()))
.finish()
}
}
impl<W: Write + Seek> Seek for BufWriter<W> {
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
self.flush_buf()?;
self.get_mut().seek(pos)
}
}
impl<W: Write> Drop for BufWriter<W> {
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> {
buffer: &'a mut BufWriter<W>,
}
impl<'a, W: Write> LineWriterShim<'a, W> {
pub fn new(buffer: &'a mut BufWriter<W>) -> 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) -> io::Result<()> {
match self.buffered().last().copied() {
Some(b'\n') => self.buffer.flush_buf(),
_ => Ok(()),
}
}
}
impl<'a, W: Write> Write for LineWriterShim<'a, W> {
fn write(&mut self, buf: &[u8]) -> io::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) -> io::Result<()> {
self.buffer.flush()
}
fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
if !self.is_write_vectored() {
return match bufs.iter().find(|buf| !buf.is_empty()) {
Some(buf) => self.write(buf),
None => Ok(0),
};
}
let last_newline_buf_idx = bufs
.iter()
.enumerate()
.rev()
.find_map(|(i, buf)| memchr::memchr(b'\n', buf).map(|_| i));
let last_newline_buf_idx = match last_newline_buf_idx {
None => {
self.flush_if_completed_line()?;
return self.buffer.write_vectored(bufs);
}
Some(i) => i,
};
self.buffer.flush_buf()?;
let (lines, tail) = bufs.split_at(last_newline_buf_idx + 1);
let flushed = self.inner_mut().write_vectored(lines)?;
if flushed == 0 {
return Ok(0);
}
let lines_len = lines.iter().map(|buf| buf.len()).sum();
if flushed < lines_len {
return Ok(flushed);
}
let buffered: usize = tail
.iter()
.filter(|buf| !buf.is_empty())
.map(|buf| self.buffer.write_to_buf(buf))
.take_while(|&n| n > 0)
.sum();
Ok(flushed + buffered)
}
fn is_write_vectored(&self) -> bool {
self.buffer.is_write_vectored()
}
fn write_all(&mut self, buf: &[u8]) -> io::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> {
inner: BufWriter<W>,
}
impl<W: Write> LineWriter<W> {
pub fn new(inner: W) -> LineWriter<W> {
LineWriter::with_capacity(1024, inner)
}
pub fn with_capacity(capacity: usize, inner: W) -> LineWriter<W> {
LineWriter { inner: BufWriter::with_capacity(capacity, 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, IntoInnerError<LineWriter<W>>> {
self.inner
.into_inner()
.map_err(|IntoInnerError(buf, e)| IntoInnerError(LineWriter { inner: buf }, e))
}
}
impl<W: Write> Write for LineWriter<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
LineWriterShim::new(&mut self.inner).write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.inner.flush()
}
fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
LineWriterShim::new(&mut self.inner).write_vectored(bufs)
}
fn is_write_vectored(&self) -> bool {
self.inner.is_write_vectored()
}
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
LineWriterShim::new(&mut self.inner).write_all(buf)
}
fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
LineWriterShim::new(&mut self.inner).write_all_vectored(bufs)
}
fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> io::Result<()> {
LineWriterShim::new(&mut self.inner).write_fmt(fmt)
}
}
impl<W: Write> fmt::Debug for LineWriter<W>
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.buf.len(), self.inner.buf.capacity()),
)
.finish()
}
}