Skip to main content

trx_rs/
mmap_backing.rs

1use std::sync::Arc;
2
3use bytemuck::{cast_slice, cast_slice_mut, Pod};
4use memmap2::{Mmap, MmapMut};
5
6use crate::error::{Result, TrxError};
7
8/// Convert a `Vec<T: Pod>` into a `Vec<u8>` by copying the raw bytes.
9pub fn vec_to_bytes<T: Pod>(v: Vec<T>) -> Vec<u8> {
10    cast_slice::<T, u8>(&v).to_vec()
11}
12
13/// Owns the backing memory for a TRX data array.
14///
15/// May be a read-only mmap, a read-write mmap, an owned heap buffer, or a slice of a shared mmap.
16pub enum MmapBacking {
17    ReadOnly(Mmap),
18    ReadWrite(MmapMut),
19    SharedSlice {
20        mmap: Arc<Mmap>,
21        offset: usize,
22        len: usize,
23    },
24    Owned(Vec<u8>),
25    OwnedU64(Vec<u64>, usize),
26    OwnedU32(Vec<u32>, usize),
27}
28
29impl MmapBacking {
30    /// Raw bytes view.
31    pub fn as_bytes(&self) -> &[u8] {
32        match self {
33            MmapBacking::ReadOnly(m) => m,
34            MmapBacking::ReadWrite(m) => m,
35            MmapBacking::SharedSlice { mmap, offset, len } => &mmap[*offset..*offset + *len],
36            MmapBacking::Owned(v) => v,
37            MmapBacking::OwnedU64(v, len) => &bytemuck::cast_slice(v)[..*len],
38            MmapBacking::OwnedU32(v, len) => &bytemuck::cast_slice(v)[..*len],
39        }
40    }
41
42    /// Mutable raw bytes view (only for ReadWrite and Owned).
43    pub fn as_bytes_mut(&mut self) -> Result<&mut [u8]> {
44        match self {
45            MmapBacking::ReadOnly(_) => Err(TrxError::Argument(
46                "cannot mutably access read-only mmap".into(),
47            )),
48            MmapBacking::SharedSlice { .. } => Err(TrxError::Argument(
49                "cannot mutably access shared mmap slice".into(),
50            )),
51            MmapBacking::ReadWrite(m) => Ok(m.as_mut()),
52            MmapBacking::Owned(v) => Ok(v.as_mut_slice()),
53            MmapBacking::OwnedU64(v, len) => Ok(&mut bytemuck::cast_slice_mut(v)[..*len]),
54            MmapBacking::OwnedU32(v, len) => Ok(&mut bytemuck::cast_slice_mut(v)[..*len]),
55        }
56    }
57
58    /// Length in bytes.
59    pub fn len(&self) -> usize {
60        self.as_bytes().len()
61    }
62
63    /// Whether the backing is empty.
64    pub fn is_empty(&self) -> bool {
65        self.len() == 0
66    }
67
68    pub fn is_mapped(&self) -> bool {
69        matches!(
70            self,
71            MmapBacking::ReadOnly(_) | MmapBacking::ReadWrite(_) | MmapBacking::SharedSlice { .. }
72        )
73    }
74
75    /// Cast the raw bytes to a typed slice.
76    ///
77    /// Panics if the bytes are not aligned or the length is not a multiple
78    /// of `size_of::<T>()`.
79    pub fn cast_slice<T: Pod>(&self) -> &[T] {
80        cast_slice(self.as_bytes())
81    }
82
83    /// Cast the raw bytes to a mutable typed slice.
84    pub fn cast_slice_mut<T: Pod>(&mut self) -> Result<&mut [T]> {
85        let bytes = self.as_bytes_mut()?;
86        Ok(cast_slice_mut(bytes))
87    }
88}
89
90impl std::fmt::Debug for MmapBacking {
91    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92        match self {
93            MmapBacking::ReadOnly(m) => write!(f, "ReadOnly({} bytes)", m.len()),
94            MmapBacking::ReadWrite(m) => write!(f, "ReadWrite({} bytes)", m.len()),
95            MmapBacking::SharedSlice { offset, len, .. } => {
96                write!(f, "SharedSlice({} bytes at offset {})", len, offset)
97            }
98            MmapBacking::Owned(v) => write!(f, "Owned({} bytes)", v.len()),
99            MmapBacking::OwnedU64(_, len) => write!(f, "OwnedU64({} bytes)", len),
100            MmapBacking::OwnedU32(_, len) => write!(f, "OwnedU32({} bytes)", len),
101        }
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn shared_slice_behavior_and_mapping_status() {
111        let mut anon = memmap2::MmapMut::map_anon(16).unwrap();
112        anon[0..4].copy_from_slice(&1.0f32.to_ne_bytes());
113        anon[4..8].copy_from_slice(&2.0f32.to_ne_bytes());
114        let mmap = Arc::new(anon.make_read_only().unwrap());
115
116        let mut backing = MmapBacking::SharedSlice {
117            mmap: Arc::clone(&mmap),
118            offset: 0,
119            len: 8,
120        };
121
122        assert!(backing.is_mapped());
123        assert_eq!(backing.len(), 8);
124        assert_eq!(backing.as_bytes(), &mmap[0..8]);
125        assert_eq!(backing.cast_slice::<f32>(), &[1.0f32, 2.0f32]);
126        assert!(backing.as_bytes_mut().is_err());
127
128        let owned = MmapBacking::Owned(vec![1, 2, 3]);
129        assert!(!owned.is_mapped());
130    }
131}