use std::ops::{Deref, DerefMut};
use std::path::Path;
use std::sync::Arc;
use std::{fmt, mem, slice};
use memmap2::MmapMut;
use super::advice::{Advice, AdviceSetting, Madviseable};
use super::ops;
use crate::common::bitvec::BitSlice;
type Result<T> = std::result::Result<T, Error>;
pub type MmapFlusher = Box<dyn FnOnce() -> Result<()> + Send>;
pub struct MmapType<T>
where
T: ?Sized + 'static,
{
r#type: &'static mut T,
mmap: Arc<MmapMut>,
}
impl<T: ?Sized> fmt::Debug for MmapType<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("MmapType")
.field("mmap", &self.mmap)
.finish_non_exhaustive()
}
}
impl<T> MmapType<T>
where
T: Sized + 'static,
{
pub unsafe fn from(mmap_with_type: MmapMut) -> Self {
unsafe { Self::try_from(mmap_with_type).unwrap() }
}
pub unsafe fn try_from(mut mmap_with_type: MmapMut) -> Result<Self> {
let r#type = unsafe { mmap_prefix_to_type_unbounded(&mut mmap_with_type)? };
let mmap = Arc::new(mmap_with_type);
Ok(Self { r#type, mmap })
}
}
impl<T> MmapType<[T]>
where
T: 'static,
{
pub unsafe fn try_slice_from(mut mmap_with_slice: MmapMut) -> Result<Self> {
let r#type = unsafe { mmap_to_slice_unbounded(&mut mmap_with_slice, 0)? };
let mmap = Arc::new(mmap_with_slice);
Ok(Self { r#type, mmap })
}
}
impl<T> MmapType<T>
where
T: ?Sized + 'static,
{
pub fn flusher(&self) -> MmapFlusher {
Box::new({
let mmap = self.mmap.clone();
move || {
if !mmap.is_empty() {
mmap.flush()?;
}
Ok(())
}
})
}
#[cfg(unix)]
pub unsafe fn unchecked_advise(&self, advice: memmap2::UncheckedAdvice) -> std::io::Result<()> {
unsafe { self.mmap.unchecked_advise(advice) }
}
pub fn populate(&self) -> std::io::Result<()> {
self.mmap.populate();
Ok(())
}
pub fn clear_cache(&self) -> std::io::Result<()> {
let Self { r#type: _, mmap } = self;
mmap.clear_cache();
Ok(())
}
}
impl<T> Deref for MmapType<T>
where
T: ?Sized + 'static,
{
type Target = T;
#[allow(clippy::needless_lifetimes)]
fn deref<'bounded>(&'bounded self) -> &'bounded Self::Target {
self.r#type
}
}
impl<T> DerefMut for MmapType<T>
where
T: ?Sized + 'static,
{
#[allow(clippy::needless_lifetimes)]
fn deref_mut<'bounded>(&'bounded mut self) -> &'bounded mut Self::Target {
self.r#type
}
}
pub struct MmapSlice<T>
where
T: Sized + 'static,
{
mmap: MmapType<[T]>,
}
impl<T> fmt::Debug for MmapSlice<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("MmapSlice")
.field("mmap", &self.mmap)
.finish_non_exhaustive()
}
}
impl<T> MmapSlice<T> {
pub unsafe fn from(mmap_with_slice: MmapMut) -> Self {
unsafe { Self::try_from(mmap_with_slice).unwrap() }
}
pub unsafe fn try_from(mmap_with_slice: MmapMut) -> Result<Self> {
let r#type = unsafe { MmapType::try_slice_from(mmap_with_slice) };
r#type.map(|mmap| Self { mmap })
}
pub fn flusher(&self) -> MmapFlusher {
self.mmap.flusher()
}
pub fn create(path: &Path, mut iter: impl ExactSizeIterator<Item = T>) -> Result<()> {
let file_len = iter.len() * mem::size_of::<T>();
let _file = ops::create_and_ensure_length(path, file_len)?;
let mmap = ops::open_write_mmap(
path,
AdviceSetting::Advice(Advice::Normal), false,
)?;
let mut mmap_slice = unsafe { Self::try_from(mmap)? };
mmap_slice.fill_with(|| iter.next().expect("iterator size mismatch"));
mmap_slice.flusher()()?;
Ok(())
}
pub fn populate(&self) -> std::io::Result<()> {
self.mmap.populate()?;
Ok(())
}
pub fn clear_cache(&self) -> std::io::Result<()> {
let Self { mmap } = self;
mmap.clear_cache()?;
Ok(())
}
}
impl<T> Deref for MmapSlice<T> {
type Target = MmapType<[T]>;
fn deref(&self) -> &Self::Target {
&self.mmap
}
}
impl<T> DerefMut for MmapSlice<T> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.mmap
}
}
impl<T: 'static> AsRef<[T]> for MmapSlice<T> {
fn as_ref(&self) -> &[T] {
&self.mmap
}
}
#[derive(Debug)]
pub struct MmapBitSlice {
mmap: MmapType<BitSlice>,
}
impl MmapBitSlice {
const MIN_FILE_SIZE: usize = mem::size_of::<usize>();
pub fn from(mmap: MmapMut, header_size: usize) -> Self {
Self::try_from(mmap, header_size).unwrap()
}
pub fn try_from(mut mmap: MmapMut, header_size: usize) -> Result<Self> {
let data = unsafe { mmap_to_slice_unbounded(&mut mmap, header_size)? };
let bitslice = BitSlice::from_slice_mut(data);
let mmap = Arc::new(mmap);
Ok(Self {
mmap: MmapType {
r#type: bitslice,
mmap,
},
})
}
pub fn flusher(&self) -> MmapFlusher {
self.mmap.flusher()
}
pub fn create(path: &Path, bitslice: &BitSlice) -> Result<()> {
let bits_count = bitslice.len();
let bytes_count = bits_count
.div_ceil(u8::BITS as usize)
.next_multiple_of(Self::MIN_FILE_SIZE);
let _file = ops::create_and_ensure_length(path, bytes_count)?;
let mmap = ops::open_write_mmap(
path,
AdviceSetting::Advice(Advice::Normal), false,
)?;
let mut mmap_bitslice = MmapBitSlice::try_from(mmap, 0)?;
mmap_bitslice.fill_with(|idx| {
bitslice
.get(idx)
.map(|bitref| bitref.as_ref().to_owned())
.unwrap_or(false)
});
mmap_bitslice.flusher()()?;
Ok(())
}
pub fn populate(&self) -> std::io::Result<()> {
self.mmap.populate()?;
Ok(())
}
pub fn clear_cache(&self) -> std::io::Result<()> {
let Self { mmap } = self;
mmap.clear_cache()?;
Ok(())
}
}
impl Deref for MmapBitSlice {
type Target = BitSlice;
fn deref(&self) -> &BitSlice {
&self.mmap
}
}
impl DerefMut for MmapBitSlice {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.mmap
}
}
#[derive(thiserror::Error, Debug)]
pub enum Error {
#[error("Mmap length must be {0} to match the size of type, but it is {1}")]
SizeExact(usize, usize),
#[error("Mmap length must be at least {0} to match the size of type, but it is {1}")]
SizeLess(usize, usize),
#[error("Mmap length must be multiple of {0} to match the size of type, but it is {1}")]
SizeMultiple(usize, usize),
#[error("{0}")]
Io(#[from] std::io::Error),
#[error("File not found: {0}")]
MissingFile(String),
}
unsafe fn mmap_prefix_to_type_unbounded<'unbnd, T>(mmap: &mut MmapMut) -> Result<&'unbnd mut T>
where
T: Sized,
{
let size_t = mem::size_of::<T>();
if mmap.len() < size_t {
return Err(Error::SizeLess(size_t, mmap.len()));
}
let bytes: &'unbnd mut [u8] = unsafe {
let slice = mmap.deref_mut();
slice::from_raw_parts_mut(slice.as_mut_ptr(), size_t)
};
assert_alignment::<_, T>(bytes);
#[cfg(debug_assertions)]
if mmap.len() != size_t {
log::warn!(
"Mmap length {} is not equal to size of type {}",
mmap.len(),
size_t,
);
}
#[cfg(debug_assertions)]
if bytes.len() != mem::size_of::<T>() {
return Err(Error::SizeExact(mem::size_of::<T>(), bytes.len()));
}
let ptr = bytes.as_mut_ptr().cast::<T>();
Ok(unsafe { &mut *ptr })
}
unsafe fn mmap_to_slice_unbounded<'unbnd, T>(
mmap: &mut MmapMut,
header_size: usize,
) -> Result<&'unbnd mut [T]>
where
T: Sized,
{
let size_t = mem::size_of::<T>();
if size_t == 0 {
debug_assert_eq!(
mmap.len().saturating_sub(header_size),
0,
"mmap data must be zero-sized, because size T is zero",
);
} else {
debug_assert_eq!(header_size % size_t, 0, "header not multiple of size T");
if !mmap.len().is_multiple_of(size_t) {
return Err(Error::SizeMultiple(size_t, mmap.len()));
}
}
let bytes: &'unbnd mut [u8] = unsafe {
let slice = mmap.deref_mut();
&mut slice::from_raw_parts_mut(slice.as_mut_ptr(), slice.len())[header_size..]
};
assert_alignment::<_, T>(bytes);
debug_assert_eq!(bytes.len() + header_size, mmap.len());
unsafe {
Ok(slice::from_raw_parts_mut(
bytes.as_mut_ptr().cast::<T>(),
bytes.len().checked_div(size_t).unwrap_or(0),
))
}
}
fn assert_alignment<S, T>(bytes: &[S]) {
assert_eq!(
bytes.as_ptr().align_offset(mem::align_of::<T>()),
0,
"type must be aligned",
);
}
#[cfg(test)]
mod tests {
use std::fmt::Debug;
use std::iter;
use rand::rngs::{SmallRng, StdRng};
use rand::{RngExt, SeedableRng};
use tempfile::{Builder, NamedTempFile};
use super::*;
use crate::common::mmap::AdviceSetting;
fn create_temp_mmap_file(len: usize) -> NamedTempFile {
let tempfile = Builder::new()
.prefix("test.")
.suffix(".mmap")
.tempfile()
.unwrap();
#[allow(clippy::disallowed_methods, reason = "test code")]
tempfile.as_file().set_len(len as u64).unwrap();
tempfile
}
#[test]
fn test_open_zero_type() {
check_open_zero_type::<()>(());
check_open_zero_type::<u8>(0);
check_open_zero_type::<usize>(0);
check_open_zero_type::<f32>(0.0);
}
fn check_open_zero_type<T: Sized + PartialEq + Debug + 'static>(zero: T) {
let bytes = mem::size_of::<T>();
let tempfile = create_temp_mmap_file(bytes);
let mmap = ops::open_write_mmap(tempfile.path(), AdviceSetting::Global, false).unwrap();
let mmap_type: MmapType<T> = unsafe { MmapType::from(mmap) };
assert_eq!(mmap_type.deref(), &zero);
}
#[test]
fn test_open_zero_slice() {
check_open_zero_slice::<()>(0, ());
check_open_zero_slice::<u8>(0, 0);
check_open_zero_slice::<u8>(1, 0);
check_open_zero_slice::<u8>(131, 0);
check_open_zero_slice::<usize>(0, 0);
check_open_zero_slice::<usize>(1, 0);
check_open_zero_slice::<usize>(131, 0);
check_open_zero_slice::<f32>(0, 0.0);
check_open_zero_slice::<f32>(1, 0.0);
check_open_zero_slice::<f32>(131, 0.0);
}
#[test]
#[should_panic]
fn test_open_zero_slice_infinite_length() {
check_open_zero_slice::<()>(1, ());
}
fn check_open_zero_slice<T: Sized + PartialEq + Debug + 'static>(len: usize, zero: T) {
let bytes = mem::size_of::<T>() * len;
let tempfile = create_temp_mmap_file(bytes);
let mmap = ops::open_write_mmap(tempfile.path(), AdviceSetting::Global, false).unwrap();
let mmap_slice: MmapSlice<T> = unsafe { MmapSlice::from(mmap) };
assert_eq!(mmap_slice.len(), len);
assert!(mmap_slice.iter().all(|i| i == &zero));
}
#[test]
fn test_reopen_random() {
let mut rng = SmallRng::seed_from_u64(42);
check_reopen_random::<(), _>(0, || rng.random());
check_reopen_random::<u8, _>(0, || rng.random());
check_reopen_random::<u8, _>(1, || rng.random());
check_reopen_random::<u8, _>(131, || rng.random());
check_reopen_random::<u64, _>(0, || rng.random());
check_reopen_random::<u64, _>(1, || rng.random());
check_reopen_random::<u64, _>(131, || rng.random());
check_reopen_random::<f32, _>(0, || rng.random());
check_reopen_random::<f32, _>(1, || rng.random());
check_reopen_random::<f32, _>(131, || rng.random());
}
fn check_reopen_random<T, R>(len: usize, rng: R)
where
T: Sized + Copy + PartialEq + Debug + 'static,
R: FnMut() -> T,
{
let bytes = mem::size_of::<T>() * len;
let tempfile = create_temp_mmap_file(bytes);
let template: Vec<T> = iter::repeat_with(rng).take(len).collect();
{
let mmap = ops::open_write_mmap(tempfile.path(), AdviceSetting::Global, false).unwrap();
let mut mmap_slice: MmapSlice<T> = unsafe { MmapSlice::from(mmap) };
assert_eq!(mmap_slice.len(), len);
mmap_slice.copy_from_slice(&template);
}
{
let mmap = ops::open_write_mmap(tempfile.path(), AdviceSetting::Global, false).unwrap();
let mmap_slice: MmapSlice<T> = unsafe { MmapSlice::from(mmap) };
assert_eq!(mmap_slice.as_ref(), template);
}
}
#[test]
fn test_bitslice() {
check_bitslice_with_header(0, 0);
check_bitslice_with_header(0, 128);
check_bitslice_with_header(512, 0);
check_bitslice_with_header(512, 256);
check_bitslice_with_header(11721 * 8, 256);
}
fn check_bitslice_with_header(bits: usize, header_size: usize) {
let bytes = (mem::size_of::<usize>() * bits / 8) + header_size;
let tempfile = create_temp_mmap_file(bytes);
{
let mut rng = StdRng::seed_from_u64(42);
let mmap = ops::open_write_mmap(tempfile.path(), AdviceSetting::Global, false).unwrap();
let mut mmap_bitslice = MmapBitSlice::from(mmap, header_size);
(0..bits).for_each(|i| mmap_bitslice.set(i, rng.random()));
}
{
let mut rng = StdRng::seed_from_u64(42);
let mmap = ops::open_write_mmap(tempfile.path(), AdviceSetting::Global, false).unwrap();
let mmap_bitslice = MmapBitSlice::from(mmap, header_size);
(0..bits).for_each(|i| assert_eq!(mmap_bitslice[i], rng.random::<bool>()));
}
}
#[test]
fn test_zero_sized_type() {
{
let tempfile = create_temp_mmap_file(0);
let mmap = ops::open_write_mmap(tempfile.path(), AdviceSetting::Global, false).unwrap();
let result = unsafe { MmapType::<()>::try_from(mmap).unwrap() };
assert_eq!(result.deref(), &());
}
{
let tempfile = create_temp_mmap_file(0);
let mmap = ops::open_write_mmap(tempfile.path(), AdviceSetting::Global, false).unwrap();
let result = unsafe { MmapSlice::<()>::try_from(mmap).unwrap() };
assert_eq!(result.as_ref(), &[]);
assert_alignment::<_, ()>(result.as_ref());
}
}
#[test]
fn test_double_read_mmap() {
let tempfile = create_temp_mmap_file(1024);
let mut mmap_write =
ops::open_write_mmap(tempfile.path(), AdviceSetting::Global, false).unwrap();
let mmap_read = ops::open_read_mmap(
tempfile.path(),
AdviceSetting::Advice(Advice::Sequential),
false,
)
.unwrap();
mmap_write[333] = 42;
assert_eq!(mmap_read[333], 42);
}
}