quick_file_transfer/
mmap_reader.rs

1use anyhow::{bail, Result};
2use memmap2::Mmap;
3use std::{fs::File, ops::Range, path::Path};
4
5pub struct MemoryMapWrapper {
6    mmap: Mmap, // Hold this or segfault
7}
8
9impl MemoryMapWrapper {
10    pub fn new(path: &Path) -> Result<Self> {
11        let file = File::open(path)?;
12        let mmap = unsafe { Mmap::map(&file)? };
13
14        #[cfg(target_os = "linux")]
15        mmap.advise(memmap2::Advice::PopulateRead)?;
16
17        Ok(Self { mmap })
18    }
19
20    pub fn borrow_full(&self) -> &[u8] {
21        &self.mmap[..]
22    }
23
24    pub fn borrow_slice(&self, range: Range<usize>) -> Result<&[u8]> {
25        if range.start > range.end || range.end > self.mmap.len() {
26            bail!("Invalid slice range");
27        }
28        Ok(&self.mmap[range])
29    }
30
31    pub fn flen(&self) -> usize {
32        self.mmap.len()
33    }
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39    use std::fs;
40    use temp_dir::TempDir;
41    use testresult::TestResult;
42
43    #[test]
44    fn test_borrow_slice() -> TestResult {
45        let d = TempDir::new()?;
46        let c = d.child("file");
47
48        let content = b"Hello, world!";
49        fs::write(c.as_path(), content)?;
50
51        let reader = MemoryMapWrapper::new(c.as_path())?;
52        let slice = reader.borrow_slice(0..reader.flen())?;
53
54        assert_eq!(slice, content);
55        Ok(())
56    }
57
58    #[test]
59    fn test_borrow_slice_at() -> TestResult {
60        let d = TempDir::new()?;
61        let c = d.child("file");
62        let content = b"Hello, world!";
63        fs::write(c.as_path(), content)?;
64
65        let path = c.as_path();
66        let reader = MemoryMapWrapper::new(path)?;
67
68        let slice1 = reader.borrow_slice(0..5)?;
69        let over_lapping_slice = reader.borrow_slice(4..9)?;
70        let slice2 = reader.borrow_slice(7..12)?;
71        assert_eq!(slice2, b"world");
72        assert_eq!(slice1, b"Hello");
73        assert_eq!(over_lapping_slice.len(), 5);
74
75        // Test out of bounds
76        assert!(reader.borrow_slice(5..14).is_err());
77
78        Ok(())
79    }
80}