use std::fs::File;
use std::io::{self, BufReader, BufWriter, Read, Write};
use std::path::Path;
pub const DEFAULT_BUFFER_SIZE: usize = 64 * 1024;
pub const LARGE_BUFFER_SIZE: usize = 1024 * 1024;
pub const SMALL_BUFFER_SIZE: usize = 4 * 1024;
pub fn buffered_reader<P: AsRef<Path>>(path: P) -> io::Result<BufReader<File>> {
let file = File::open(path)?;
Ok(BufReader::with_capacity(DEFAULT_BUFFER_SIZE, file))
}
pub fn buffered_reader_with_capacity<P: AsRef<Path>>(
path: P,
capacity: usize,
) -> io::Result<BufReader<File>> {
let file = File::open(path)?;
Ok(BufReader::with_capacity(capacity, file))
}
pub fn buffered_writer<P: AsRef<Path>>(path: P) -> io::Result<BufWriter<File>> {
let file = File::create(path)?;
Ok(BufWriter::with_capacity(DEFAULT_BUFFER_SIZE, file))
}
pub fn buffered_writer_with_capacity<P: AsRef<Path>>(
path: P,
capacity: usize,
) -> io::Result<BufWriter<File>> {
let file = File::create(path)?;
Ok(BufWriter::with_capacity(capacity, file))
}
pub fn read_chunks<P, F>(path: P, chunk_size: usize, mut callback: F) -> io::Result<()>
where
P: AsRef<Path>,
F: FnMut(&[u8]) -> io::Result<()>,
{
let file = File::open(path)?;
let mut reader = BufReader::with_capacity(chunk_size.max(4096), file);
let mut buffer = vec![0u8; chunk_size];
loop {
let n = reader.read(&mut buffer)?;
if n == 0 {
break;
}
callback(&buffer[..n])?;
}
Ok(())
}
pub fn write_chunks<P, I, D>(path: P, chunks: I) -> io::Result<()>
where
P: AsRef<Path>,
I: IntoIterator<Item = D>,
D: AsRef<[u8]>,
{
let file = File::create(path)?;
let mut writer = BufWriter::with_capacity(DEFAULT_BUFFER_SIZE, file);
for chunk in chunks {
writer.write_all(chunk.as_ref())?;
}
writer.flush()?;
Ok(())
}
pub fn copy_buffered<R: Read, W: Write>(
reader: &mut R,
writer: &mut W,
buffer_size: usize,
) -> io::Result<u64> {
let mut buffer = vec![0u8; buffer_size];
let mut total = 0u64;
loop {
let n = reader.read(&mut buffer)?;
if n == 0 {
break;
}
writer.write_all(&buffer[..n])?;
total += n as u64;
}
Ok(total)
}
pub struct ChunkStream<R> {
reader: BufReader<R>,
chunk_size: usize,
}
impl<R: Read> ChunkStream<R> {
pub fn new(reader: R) -> Self {
Self::with_chunk_size(reader, DEFAULT_BUFFER_SIZE)
}
pub fn with_chunk_size(reader: R, chunk_size: usize) -> Self {
Self {
reader: BufReader::with_capacity(chunk_size.max(4096), reader),
chunk_size,
}
}
pub fn next_chunk(&mut self) -> io::Result<Option<Vec<u8>>> {
let mut buffer = vec![0u8; self.chunk_size];
let n = self.reader.read(&mut buffer)?;
if n == 0 {
return Ok(None);
}
buffer.truncate(n);
Ok(Some(buffer))
}
pub fn process_all<F>(&mut self, mut callback: F) -> io::Result<()>
where
F: FnMut(&[u8]) -> io::Result<()>,
{
loop {
let mut buffer = vec![0u8; self.chunk_size];
let n = self.reader.read(&mut buffer)?;
if n == 0 {
break;
}
callback(&buffer[..n])?;
}
Ok(())
}
}
#[cfg(feature = "async")]
pub mod async_buffer {
use std::io;
use std::path::Path;
use tokio::fs::File;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader, BufWriter};
use super::DEFAULT_BUFFER_SIZE;
pub async fn buffered_reader<P: AsRef<Path>>(path: P) -> io::Result<BufReader<File>> {
let file = File::open(path).await?;
Ok(BufReader::with_capacity(DEFAULT_BUFFER_SIZE, file))
}
pub async fn buffered_writer<P: AsRef<Path>>(path: P) -> io::Result<BufWriter<File>> {
let file = File::create(path).await?;
Ok(BufWriter::with_capacity(DEFAULT_BUFFER_SIZE, file))
}
pub async fn read_chunks<P, F, Fut>(
path: P,
chunk_size: usize,
mut callback: F,
) -> io::Result<()>
where
P: AsRef<Path>,
F: FnMut(Vec<u8>) -> Fut,
Fut: std::future::Future<Output = io::Result<()>>,
{
let file = File::open(path).await?;
let mut reader = BufReader::with_capacity(chunk_size.max(4096), file);
let mut buffer = vec![0u8; chunk_size];
loop {
let n = reader.read(&mut buffer).await?;
if n == 0 {
break;
}
callback(buffer[..n].to_vec()).await?;
}
Ok(())
}
pub async fn copy_buffered<R, W>(
reader: &mut R,
writer: &mut W,
buffer_size: usize,
) -> io::Result<u64>
where
R: AsyncReadExt + Unpin,
W: AsyncWriteExt + Unpin,
{
let mut buffer = vec![0u8; buffer_size];
let mut total = 0u64;
loop {
let n = reader.read(&mut buffer).await?;
if n == 0 {
break;
}
writer.write_all(&buffer[..n]).await?;
total += n as u64;
}
Ok(total)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
#[test]
fn test_chunk_stream() {
let data = b"Hello, world! This is a test.";
let cursor = Cursor::new(data);
let mut stream = ChunkStream::with_chunk_size(cursor, 10);
let mut chunks = Vec::new();
while let Some(chunk) = stream.next_chunk().unwrap() {
chunks.push(chunk);
}
let reconstructed: Vec<u8> = chunks.into_iter().flatten().collect();
assert_eq!(reconstructed, data);
}
#[test]
fn test_copy_buffered() {
let data = b"Test data for copying";
let mut reader = Cursor::new(data);
let mut writer = Vec::new();
let copied = copy_buffered(&mut reader, &mut writer, 8).unwrap();
assert_eq!(copied, data.len() as u64);
assert_eq!(writer, data);
}
#[test]
fn test_process_all() {
let data = b"Process all chunks";
let cursor = Cursor::new(data);
let mut stream = ChunkStream::with_chunk_size(cursor, 5);
let mut total_bytes = 0;
stream
.process_all(|chunk| {
total_bytes += chunk.len();
Ok(())
})
.unwrap();
assert_eq!(total_bytes, data.len());
}
}