data_pile/
shared_mmap.rs

1use memmap2::Mmap;
2use std::{
3    ops::{Bound::*, RangeBounds},
4    slice,
5    sync::Arc,
6};
7
8/// A structure that implements a view into memory mapping.
9#[derive(Debug, Clone)]
10pub struct SharedMmap {
11    mmap: Arc<Mmap>,
12    len: usize,
13    slice: *const u8,
14}
15
16impl SharedMmap {
17    pub(crate) fn new(mmap: Mmap) -> SharedMmap {
18        let len = mmap.len();
19        let slice = mmap.as_ptr();
20        SharedMmap {
21            mmap: Arc::new(mmap),
22            len,
23            slice,
24        }
25    }
26
27    pub fn len(&self) -> usize {
28        self.len
29    }
30
31    pub fn is_empty(&self) -> bool {
32        self.len == 0
33    }
34
35    /// Get a sub-view. It will point to the same memory mapping as the parent
36    /// mapping.
37    pub fn slice(&self, bounds: impl RangeBounds<usize>) -> SharedMmap {
38        if self.len == 0 {
39            return SharedMmap {
40                len: 0,
41                ..self.clone()
42            };
43        }
44        let start = match bounds.start_bound() {
45            Included(start) => *start,
46            Excluded(start) => start + 1,
47            Unbounded => 0,
48        };
49
50        if start >= self.len {
51            return SharedMmap {
52                len: 0,
53                ..self.clone()
54            };
55        }
56
57        let end = match bounds.end_bound() {
58            Included(end) => *end,
59            Excluded(end) if *end == 0 => {
60                return SharedMmap {
61                    len: 0,
62                    ..self.clone()
63                };
64            }
65            Excluded(end) => end - 1,
66            Unbounded => self.len - 1,
67        };
68        let end = std::cmp::min(end, self.len - 1);
69
70        let len = if start <= end { end - start + 1 } else { 0 };
71        let slice = unsafe { self.slice.add(start) };
72
73        SharedMmap {
74            mmap: self.mmap.clone(),
75            len,
76            slice,
77        }
78    }
79
80    fn get_ref(&self) -> &[u8] {
81        unsafe { slice::from_raw_parts(self.slice, self.len) }
82    }
83}
84
85// Those are safe to implement because the underlying `*const u8` is never
86// modified.
87unsafe impl Send for SharedMmap {}
88unsafe impl Sync for SharedMmap {}
89
90impl AsRef<[u8]> for SharedMmap {
91    fn as_ref(&self) -> &[u8] {
92        self.get_ref()
93    }
94}