use std::io::{Read, Result};
pub trait ReadExt {
fn read_byte(&mut self) -> Result<u8>;
fn read_n<const N: usize>(&mut self) -> Result<[u8; N]>;
fn read_string(&mut self) -> Result<String>;
}
pub trait BufReadExt {
fn read_until<const N: usize>(
&mut self,
delimeter: &[u8; N],
buf: &mut Vec<u8>,
) -> Result<usize>;
}
impl<R: Read> ReadExt for R {
fn read_byte(&mut self) -> Result<u8> {
let [c] = self.read_n::<1>()?;
Ok(c)
}
fn read_n<const N: usize>(&mut self) -> Result<[u8; N]> {
let mut buf = [0; N];
self.read_exact(&mut buf)?;
Ok(buf)
}
fn read_string(&mut self) -> Result<String> {
let mut body = String::new();
self.read_to_string(&mut body)?;
Ok(body)
}
}
mod read_until {
use bstr::ByteSlice;
use std::io::BufRead;
use std::io::Result;
use super::BufReadExt;
fn step<const N: usize>(
excess: &[u8],
delimeter: &[u8; N],
buf: &mut Vec<u8>,
) -> (bool, usize) {
if let Some(i) = excess.find(&delimeter) {
buf.extend_from_slice(&excess[..i + N]);
return (true, i + N);
}
for n in 1..N {
if buf.ends_with(&delimeter[..n]) && excess.starts_with(&delimeter[n..N]) {
let used = N - n;
buf.extend_from_slice(&excess[..used]);
return (true, used);
}
}
buf.extend_from_slice(excess);
(false, excess.len())
}
impl<R: BufRead> BufReadExt for R {
fn read_until<const N: usize>(
&mut self,
delimiter: &[u8; N],
buf: &mut Vec<u8>,
) -> Result<usize> {
let mut read = 0;
loop {
let (done, used) = step(self.fill_buf()?, delimiter, buf);
self.consume(used);
read += used;
if done || used == 0 {
return Ok(read);
}
}
}
}
#[cfg(test)]
mod test {
use super::BufReadExt;
use std::io::BufReader;
#[test]
fn read1() {
let mut buf0: &[u8] = b"ebfdefghijkl";
let mut buf1 = Vec::new();
let exp = b"ebf";
let mut bufread = BufReader::new(&mut buf0);
bufread.read_until(b"f", &mut buf1).unwrap();
assert_eq!(buf1, exp);
}
#[test]
fn read2() {
let mut buf0: &[u8] = b"ebfdefghijkl";
let mut buf1 = Vec::new();
let exp = b"ebfdef";
let mut bufread = BufReader::new(&mut buf0);
bufread.read_until(b"ef", &mut buf1).unwrap();
assert_eq!(buf1, exp);
}
#[test]
fn read3() {
let mut buf0: &[u8] = b"ebfdefghijkl";
let mut buf1 = Vec::new();
let exp = b"ebfdef";
let mut bufread = BufReader::new(&mut buf0);
bufread.read_until(b"def", &mut buf1).unwrap();
assert_eq!(buf1, exp);
}
#[test]
fn read4() {
let mut buf0: &[u8] = b"ebfdefghijkl";
let mut buf1 = Vec::new();
let exp = b"ebfdef";
let mut bufread = BufReader::new(&mut buf0);
bufread.read_until(b"fdef", &mut buf1).unwrap();
assert_eq!(buf1, exp);
}
#[test]
fn read5() {
let mut buf0: &[u8] = b"ebfdefghijkl";
let mut buf1 = Vec::new();
let exp = b"ebfdef";
let mut bufread = BufReader::new(&mut buf0);
bufread.read_until(b"bfdef", &mut buf1).unwrap();
assert_eq!(buf1, exp);
}
#[test]
fn read6() {
let mut buf0: &[u8] = b"ebfdefghijkl";
let mut buf1 = Vec::new();
let exp = b"ebfdef";
let mut bufread = BufReader::new(&mut buf0);
bufread.read_until(b"ebfdef", &mut buf1).unwrap();
assert_eq!(buf1, exp);
}
}
}