rsbag/src/reader/io.rs

92 lines
2.7 KiB
Rust

use std::io;
use bytes::{Buf, BytesMut};
use super::BagReader;
use crate::{parse, Error, Result};
const READ_SIZE: usize = 4096;
pub struct IoReader<R> {
read: R,
buffer: BytesMut,
consumed: usize,
}
impl<R: io::Read + io::Seek> IoReader<R> {
pub fn new(read: R) -> Self {
Self {
read,
buffer: BytesMut::with_capacity(READ_SIZE),
consumed: 0,
}
}
fn read_more(&mut self, n: usize) -> Result<()> {
if n == 0 {
return Ok(());
}
let old_size = self.buffer.len();
self.buffer.resize(old_size + n, 0);
let read = self.read.read(&mut self.buffer[old_size..])?;
self.buffer.truncate(old_size + read);
if read == 0 {
return Err(Error::Eof);
}
Ok(())
}
}
impl<R: io::Read + io::Seek> BagReader for IoReader<R> {
fn read_parser<'a, O: 'a, P>(&'a mut self, mut parser: P) -> Result<O>
where
P: nom::Parser<&'a [u8], O, parse::Error<&'a [u8]>>,
{
self.buffer.advance(self.consumed);
self.consumed = 0;
let this = self as *mut Self;
loop {
match parser.parse(&self.buffer) {
Ok((rest, output)) => {
self.consumed += self.buffer.len() - rest.len();
return Ok(output);
}
Err(nom::Err::Incomplete(needed)) => {
let needed = match needed {
nom::Needed::Unknown => 0,
nom::Needed::Size(n) => n.get(),
};
// Safety: this.buffer is only borrowed in the Ok case above, which
// immediately returns.
unsafe { &mut *this }.read_more(needed.max(READ_SIZE))?;
}
Err(nom::Err::Error(e) | nom::Err::Failure(e)) => {
return Err(e.into());
}
}
}
}
fn seek(&mut self, mut pos: io::SeekFrom) -> Result<()> {
if let io::SeekFrom::Current(pos) = &mut pos {
// If seeking relative to current position, subtract data
// read from the file but not yet consumed.
let remaining = (self.buffer.len() - self.consumed) as i64;
let new_pos = *pos - remaining;
if *pos >= 0 && new_pos < 0 {
// The new position is within the already read data, just consume more data
self.consumed += *pos as usize;
return Ok(());
}
*pos = new_pos;
}
self.buffer.clear();
self.consumed = 0;
self.read.seek(pos)?;
Ok(())
}
}