use crate::errors::Result;
use crate::mmap::MemoryMappedFile;
use crate::utils::page_size;
use std::marker::PhantomData;
pub struct ChunkIterator<'a> {
mmap: &'a MemoryMappedFile,
chunk_size: usize,
current_offset: u64,
total_len: u64,
buffer: Vec<u8>,
}
impl<'a> ChunkIterator<'a> {
pub(crate) fn new(mmap: &'a MemoryMappedFile, chunk_size: usize) -> Result<Self> {
let total_len = mmap.current_len()?;
let buffer = Vec::with_capacity(chunk_size);
Ok(Self {
mmap,
chunk_size,
current_offset: 0,
total_len,
buffer,
})
}
}
impl<'a> Iterator for ChunkIterator<'a> {
type Item = Result<Vec<u8>>;
fn next(&mut self) -> Option<Self::Item> {
if self.current_offset >= self.total_len {
return None;
}
let remaining = self.total_len - self.current_offset;
let chunk_len = remaining.min(self.chunk_size as u64);
self.buffer.resize(chunk_len as usize, 0);
match self.mmap.read_into(self.current_offset, &mut self.buffer) {
Ok(()) => {
self.current_offset += chunk_len;
Some(Ok(self.buffer.clone()))
}
Err(e) => Some(Err(e)),
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let remaining = self.total_len.saturating_sub(self.current_offset);
let chunks = (remaining as usize).div_ceil(self.chunk_size);
(chunks, Some(chunks))
}
}
impl<'a> ExactSizeIterator for ChunkIterator<'a> {}
pub struct PageIterator<'a> {
inner: ChunkIterator<'a>,
}
impl<'a> PageIterator<'a> {
pub(crate) fn new(mmap: &'a MemoryMappedFile) -> Result<Self> {
let ps = page_size();
Ok(Self {
inner: ChunkIterator::new(mmap, ps)?,
})
}
}
impl<'a> Iterator for PageIterator<'a> {
type Item = Result<Vec<u8>>;
fn next(&mut self) -> Option<Self::Item> {
self.inner.next()
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<'a> ExactSizeIterator for PageIterator<'a> {}
pub struct ChunkIteratorMut<'a> {
mmap: &'a MemoryMappedFile,
chunk_size: usize,
current_offset: u64,
total_len: u64,
_phantom: PhantomData<&'a mut [u8]>,
}
impl<'a> ChunkIteratorMut<'a> {
pub(crate) fn new(mmap: &'a MemoryMappedFile, chunk_size: usize) -> Result<Self> {
let total_len = mmap.current_len()?;
Ok(Self {
mmap,
chunk_size,
current_offset: 0,
total_len,
_phantom: PhantomData,
})
}
pub fn for_each_mut<F, E>(mut self, mut f: F) -> Result<std::result::Result<(), E>>
where
F: FnMut(u64, &mut [u8]) -> std::result::Result<(), E>,
{
while self.current_offset < self.total_len {
let remaining = self.total_len - self.current_offset;
let chunk_len = remaining.min(self.chunk_size as u64);
let mut guard = self.mmap.as_slice_mut(self.current_offset, chunk_len)?;
let slice = guard.as_mut();
match f(self.current_offset, slice) {
Ok(()) => {}
Err(e) => return Ok(Err(e)),
}
self.current_offset += chunk_len;
}
Ok(Ok(()))
}
}
impl MemoryMappedFile {
#[cfg(feature = "iterator")]
pub fn chunks(&self, chunk_size: usize) -> ChunkIterator<'_> {
ChunkIterator::new(self, chunk_size).expect("chunk iterator creation should not fail")
}
#[cfg(feature = "iterator")]
pub fn pages(&self) -> PageIterator<'_> {
PageIterator::new(self).expect("page iterator creation should not fail")
}
#[cfg(feature = "iterator")]
pub fn chunks_mut(&self, chunk_size: usize) -> ChunkIteratorMut<'_> {
ChunkIteratorMut::new(self, chunk_size)
.expect("mutable chunk iterator creation should not fail")
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::create_mmap;
use std::fs;
use std::path::PathBuf;
fn tmp_path(name: &str) -> PathBuf {
let mut p = std::env::temp_dir();
p.push(format!(
"mmap_io_iterator_test_{}_{}",
name,
std::process::id()
));
p
}
#[test]
#[cfg(feature = "iterator")]
fn test_chunk_iterator() {
let path = tmp_path("chunk_iter");
let _ = fs::remove_file(&path);
let mmap = create_mmap(&path, 10240).expect("create");
for i in 0..10 {
let data = vec![i as u8; 1024];
mmap.update_region(i * 1024, &data).expect("write");
}
mmap.flush().expect("flush");
let chunks: Vec<_> = mmap
.chunks(1024)
.collect::<Result<Vec<_>>>()
.expect("collect chunks");
assert_eq!(chunks.len(), 10);
for (i, chunk) in chunks.iter().enumerate() {
assert_eq!(chunk.len(), 1024);
assert!(chunk.iter().all(|&b| b == i as u8));
}
let chunks: Vec<_> = mmap
.chunks(3000)
.collect::<Result<Vec<_>>>()
.expect("collect chunks");
assert_eq!(chunks.len(), 4); assert_eq!(chunks[3].len(), 1240);
fs::remove_file(&path).expect("cleanup");
}
#[test]
#[cfg(feature = "iterator")]
fn test_page_iterator() {
let path = tmp_path("page_iter");
let _ = fs::remove_file(&path);
let ps = page_size();
let file_size = ps * 3 + 100;
let mmap = create_mmap(&path, file_size as u64).expect("create");
let pages: Vec<_> = mmap
.pages()
.collect::<Result<Vec<_>>>()
.expect("collect pages");
assert_eq!(pages.len(), 4); assert_eq!(pages[0].len(), ps);
assert_eq!(pages[1].len(), ps);
assert_eq!(pages[2].len(), ps);
assert_eq!(pages[3].len(), 100);
fs::remove_file(&path).expect("cleanup");
}
#[test]
#[cfg(feature = "iterator")]
fn test_mutable_chunk_iterator() {
let path = tmp_path("mut_chunk_iter");
let _ = fs::remove_file(&path);
let mmap = create_mmap(&path, 4096).expect("create");
let result = mmap.chunks_mut(1024).for_each_mut(|offset, chunk| {
let value = (offset / 1024) as u8;
chunk.fill(value);
Ok::<(), std::io::Error>(())
});
assert!(result.is_ok());
assert!(result.unwrap().is_ok());
mmap.flush().expect("flush");
let mut buf = [0u8; 1024];
for i in 0..4 {
mmap.read_into(i * 1024, &mut buf).expect("read");
assert!(buf.iter().all(|&b| b == i as u8));
}
fs::remove_file(&path).expect("cleanup");
}
#[test]
#[cfg(feature = "iterator")]
fn test_iterator_size_hint() {
let path = tmp_path("size_hint");
let _ = fs::remove_file(&path);
let mmap = create_mmap(&path, 10000).expect("create");
let iter = mmap.chunks(1000);
assert_eq!(iter.size_hint(), (10, Some(10)));
let iter = mmap.chunks(3000);
assert_eq!(iter.size_hint(), (4, Some(4)));
fs::remove_file(&path).expect("cleanup");
}
#[test]
#[cfg(feature = "iterator")]
fn test_empty_file_iteration() {
let path = tmp_path("empty_iter");
let _ = fs::remove_file(&path);
let mmap = create_mmap(&path, 1).expect("create"); mmap.resize(1).expect("resize");
let chunks: Vec<_> = mmap
.chunks(1024)
.collect::<Result<Vec<_>>>()
.expect("collect");
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].len(), 1);
fs::remove_file(&path).expect("cleanup");
}
}