use crate::file::tempfile;
use std::fs::File;
use std::io::{self, Cursor, Read, Seek, SeekFrom, Write};
#[derive(Debug)]
pub enum SpooledData {
InMemory(Cursor<Vec<u8>>),
OnDisk(File),
}
#[derive(Debug)]
pub struct SpooledTempFile {
max_size: usize,
inner: SpooledData,
}
#[inline]
pub fn spooled_tempfile(max_size: usize) -> SpooledTempFile {
SpooledTempFile::new(max_size)
}
impl SpooledTempFile {
#[must_use]
pub fn new(max_size: usize) -> SpooledTempFile {
SpooledTempFile {
max_size,
inner: SpooledData::InMemory(Cursor::new(Vec::new())),
}
}
#[must_use]
pub fn is_rolled(&self) -> bool {
match self.inner {
SpooledData::InMemory(_) => false,
SpooledData::OnDisk(_) => true,
}
}
pub fn roll(&mut self) -> io::Result<()> {
if !self.is_rolled() {
let mut file = tempfile()?;
if let SpooledData::InMemory(cursor) = &mut self.inner {
file.write_all(cursor.get_ref())?;
file.seek(SeekFrom::Start(cursor.position()))?;
}
self.inner = SpooledData::OnDisk(file);
}
Ok(())
}
pub fn set_len(&mut self, size: u64) -> Result<(), io::Error> {
if size > self.max_size as u64 {
self.roll()?; }
match &mut self.inner {
SpooledData::InMemory(cursor) => {
cursor.get_mut().resize(size as usize, 0);
Ok(())
}
SpooledData::OnDisk(file) => file.set_len(size),
}
}
#[must_use]
pub fn into_inner(self) -> SpooledData {
self.inner
}
}
impl Read for SpooledTempFile {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.read(buf),
SpooledData::OnDisk(file) => file.read(buf),
}
}
fn read_vectored(&mut self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.read_vectored(bufs),
SpooledData::OnDisk(file) => file.read_vectored(bufs),
}
}
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.read_to_end(buf),
SpooledData::OnDisk(file) => file.read_to_end(buf),
}
}
fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.read_to_string(buf),
SpooledData::OnDisk(file) => file.read_to_string(buf),
}
}
fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.read_exact(buf),
SpooledData::OnDisk(file) => file.read_exact(buf),
}
}
}
impl Write for SpooledTempFile {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if matches! {
&self.inner, SpooledData::InMemory(cursor)
if cursor.position().saturating_add(buf.len() as u64) > self.max_size as u64
} {
self.roll()?;
}
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.write(buf),
SpooledData::OnDisk(file) => file.write(buf),
}
}
fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
if matches! {
&self.inner, SpooledData::InMemory(cursor)
if bufs
.iter()
.fold(cursor.position(), |a, b| a.saturating_add(b.len() as u64))
> self.max_size as u64
} {
self.roll()?;
}
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.write_vectored(bufs),
SpooledData::OnDisk(file) => file.write_vectored(bufs),
}
}
#[inline]
fn flush(&mut self) -> io::Result<()> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.flush(),
SpooledData::OnDisk(file) => file.flush(),
}
}
}
impl Seek for SpooledTempFile {
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.seek(pos),
SpooledData::OnDisk(file) => file.seek(pos),
}
}
}