101 lines
2.6 KiB
Rust
101 lines
2.6 KiB
Rust
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
use std::io;
|
|
use std::io::{BufRead, Read};
|
|
|
|
pub enum ReadInfo {
|
|
NormalChunk(usize),
|
|
LastChunk(usize),
|
|
EmptyChunk,
|
|
}
|
|
|
|
pub struct AheadReader {
|
|
inner: Box<dyn BufRead>,
|
|
buf: Vec<u8>,
|
|
bufsz: usize,
|
|
capacity: usize,
|
|
}
|
|
|
|
impl AheadReader {
|
|
pub fn from(reader: Box<dyn BufRead>, capacity: usize) -> Self {
|
|
Self {
|
|
inner: reader,
|
|
buf: vec![0; capacity],
|
|
bufsz: 0,
|
|
capacity,
|
|
}
|
|
}
|
|
|
|
fn read_until_full(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
|
|
let mut total = 0;
|
|
loop {
|
|
match self.inner.read(buf) {
|
|
Ok(0) => break,
|
|
Ok(n) => {
|
|
total += n;
|
|
let tmp = buf;
|
|
buf = &mut tmp[n..];
|
|
}
|
|
Err(e) => match e.kind() {
|
|
io::ErrorKind::Interrupted => continue,
|
|
_ => return Err(e),
|
|
},
|
|
}
|
|
}
|
|
|
|
Ok(total)
|
|
}
|
|
|
|
pub fn read_ahead(&mut self, userbuf: &mut [u8]) -> io::Result<ReadInfo> {
|
|
// 1st read
|
|
if self.bufsz == 0 {
|
|
eprintln!("[reader] first read");
|
|
return self.first_read(userbuf);
|
|
}
|
|
|
|
eprintln!("[reader] normal read");
|
|
// normal read (not the 1st one)
|
|
self.normal_read(userbuf)
|
|
}
|
|
|
|
pub fn read_exact(&mut self, userbuf: &mut [u8]) -> io::Result<()> {
|
|
self.inner.read_exact(userbuf)
|
|
}
|
|
|
|
fn first_read(&mut self, userbuf: &mut [u8]) -> io::Result<ReadInfo> {
|
|
// 1st read directly to userbuf (we have no cached data yet)
|
|
let n = self.read_until_full(userbuf)?;
|
|
if n == 0 {
|
|
return Ok(ReadInfo::EmptyChunk);
|
|
}
|
|
|
|
// 2nd read directly into our internal buf
|
|
let mut tmp = vec![0u8; self.capacity];
|
|
let n2 = self.read_until_full(&mut tmp)?;
|
|
self.buf = tmp;
|
|
self.bufsz = n2;
|
|
if n2 == 0 {
|
|
return Ok(ReadInfo::LastChunk(n));
|
|
}
|
|
|
|
Ok(ReadInfo::NormalChunk(n))
|
|
}
|
|
|
|
fn normal_read(&mut self, userbuf: &mut [u8]) -> io::Result<ReadInfo> {
|
|
// copy internal buf to userbuf
|
|
userbuf.copy_from_slice(&self.buf);
|
|
let userbuf_sz = self.bufsz;
|
|
|
|
// 2nd read directly into our internal buf
|
|
let mut tmp = vec![0u8; self.capacity];
|
|
let n2 = self.read_until_full(&mut tmp)?;
|
|
self.buf = tmp;
|
|
self.bufsz = n2;
|
|
if n2 == 0 {
|
|
return Ok(ReadInfo::LastChunk(userbuf_sz));
|
|
}
|
|
|
|
Ok(ReadInfo::NormalChunk(userbuf_sz))
|
|
}
|
|
}
|