use std::fs::File;
use std::ops::Deref;
use std::path::Path;
use std::sync::Arc;
use memmap2::Mmap;
use crate::error::{Error, Result};
#[derive(Clone)]
enum StorageBacking {
Bytes(Arc<[u8]>),
Mmap(Arc<Mmap>),
}
#[derive(Clone)]
pub struct StorageBuffer {
backing: StorageBacking,
start: usize,
len: usize,
}
impl StorageBuffer {
pub fn from_vec(bytes: Vec<u8>) -> Self {
let len = bytes.len();
Self {
backing: StorageBacking::Bytes(Arc::<[u8]>::from(bytes)),
start: 0,
len,
}
}
pub(crate) fn from_arc_bytes(bytes: Arc<[u8]>, start: usize, len: usize) -> Self {
Self {
backing: StorageBacking::Bytes(bytes),
start,
len,
}
}
pub(crate) fn from_arc_mmap(mmap: Arc<Mmap>, start: usize, len: usize) -> Self {
Self {
backing: StorageBacking::Mmap(mmap),
start,
len,
}
}
pub fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
}
impl AsRef<[u8]> for StorageBuffer {
fn as_ref(&self) -> &[u8] {
self
}
}
impl Deref for StorageBuffer {
type Target = [u8];
fn deref(&self) -> &Self::Target {
match &self.backing {
StorageBacking::Bytes(bytes) => &bytes[self.start..self.start + self.len],
StorageBacking::Mmap(mmap) => &mmap[self.start..self.start + self.len],
}
}
}
pub trait Storage: Send + Sync {
fn len(&self) -> u64;
fn is_empty(&self) -> bool {
self.len() == 0
}
fn read_range(&self, offset: u64, len: usize) -> Result<StorageBuffer>;
}
pub type DynStorage = Arc<dyn Storage>;
pub struct BytesStorage {
data: Arc<[u8]>,
}
impl BytesStorage {
pub fn new(data: Vec<u8>) -> Self {
Self {
data: Arc::<[u8]>::from(data),
}
}
}
impl Storage for BytesStorage {
fn len(&self) -> u64 {
self.data.len() as u64
}
fn read_range(&self, offset: u64, len: usize) -> Result<StorageBuffer> {
let start = usize::try_from(offset).map_err(|_| Error::OffsetOutOfBounds(offset))?;
let end = start
.checked_add(len)
.ok_or(Error::OffsetOutOfBounds(offset))?;
if end > self.data.len() {
return Err(Error::UnexpectedEof {
offset,
needed: len as u64,
available: self.len().saturating_sub(offset),
});
}
Ok(StorageBuffer::from_arc_bytes(self.data.clone(), start, len))
}
}
pub struct MmapStorage {
mmap: Arc<Mmap>,
}
impl MmapStorage {
pub fn new(mmap: Mmap) -> Self {
Self {
mmap: Arc::new(mmap),
}
}
}
impl Storage for MmapStorage {
fn len(&self) -> u64 {
self.mmap.len() as u64
}
fn read_range(&self, offset: u64, len: usize) -> Result<StorageBuffer> {
let start = usize::try_from(offset).map_err(|_| Error::OffsetOutOfBounds(offset))?;
let end = start
.checked_add(len)
.ok_or(Error::OffsetOutOfBounds(offset))?;
if end > self.mmap.len() {
return Err(Error::UnexpectedEof {
offset,
needed: len as u64,
available: self.len().saturating_sub(offset),
});
}
Ok(StorageBuffer::from_arc_mmap(self.mmap.clone(), start, len))
}
}
pub struct FileStorage {
file: Arc<File>,
len: u64,
}
impl FileStorage {
pub fn open(path: impl AsRef<Path>) -> Result<Self> {
let file = File::open(path)?;
let len = file.metadata()?.len();
Ok(Self {
file: Arc::new(file),
len,
})
}
}
impl Storage for FileStorage {
fn len(&self) -> u64 {
self.len
}
fn read_range(&self, offset: u64, len: usize) -> Result<StorageBuffer> {
let needed = u64::try_from(len).map_err(|_| Error::OffsetOutOfBounds(offset))?;
let end = offset
.checked_add(needed)
.ok_or(Error::OffsetOutOfBounds(offset))?;
if end > self.len {
return Err(Error::UnexpectedEof {
offset,
needed,
available: self.len.saturating_sub(offset),
});
}
let mut buf = vec![0u8; len];
read_exact_at(self.file.as_ref(), &mut buf, offset)?;
Ok(StorageBuffer::from_vec(buf))
}
}
#[cfg(unix)]
fn read_exact_at(file: &File, mut buf: &mut [u8], mut offset: u64) -> std::io::Result<()> {
use std::os::unix::fs::FileExt;
while !buf.is_empty() {
let n = file.read_at(buf, offset)?;
if n == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"failed to fill whole buffer",
));
}
offset += n as u64;
buf = &mut buf[n..];
}
Ok(())
}
#[cfg(windows)]
fn read_exact_at(file: &File, mut buf: &mut [u8], mut offset: u64) -> std::io::Result<()> {
use std::os::windows::fs::FileExt;
while !buf.is_empty() {
let n = file.seek_read(buf, offset)?;
if n == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"failed to fill whole buffer",
));
}
offset += n as u64;
buf = &mut buf[n..];
}
Ok(())
}