forest/utils/io/
mmap.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use std::{fs, io, path::Path};
5
6use memmap2::MmapAsRawDesc;
7use positioned_io::{RandomAccessFile, ReadAt, Size};
8
9/// Wrapper type of [`memmap2::Mmap`] that implements [`ReadAt`] and [`Size`]
10pub struct Mmap(memmap2::Mmap);
11
12impl Mmap {
13    pub fn map(file: impl MmapAsRawDesc) -> io::Result<Self> {
14        Ok(Self(unsafe { memmap2::Mmap::map(file)? }))
15    }
16
17    pub fn map_path(path: impl AsRef<Path>) -> io::Result<Self> {
18        Self::map(&fs::File::open(path.as_ref())?)
19    }
20}
21
22impl ReadAt for Mmap {
23    #[allow(clippy::indexing_slicing)]
24    fn read_at(&self, pos: u64, buf: &mut [u8]) -> io::Result<usize> {
25        let start = pos as usize;
26        if start >= self.0.len() {
27            // This matches the behaviour for seeking past the end of a file
28            return Ok(0);
29        }
30        let end = start + buf.len();
31        if end <= self.0.len() {
32            buf.copy_from_slice(&self.0[start..end]);
33            Ok(buf.len())
34        } else {
35            let len = self.0.len() - start;
36            buf[..len].copy_from_slice(&self.0[start..]);
37            Ok(len)
38        }
39    }
40}
41
42impl Size for Mmap {
43    fn size(&self) -> io::Result<Option<u64>> {
44        Ok(Some(self.0.len() as _))
45    }
46}
47
48pub enum EitherMmapOrRandomAccessFile {
49    Mmap(Mmap),
50    RandomAccessFile(RandomAccessFile),
51}
52
53impl EitherMmapOrRandomAccessFile {
54    pub fn open(path: impl AsRef<Path>) -> io::Result<Self> {
55        Ok(if should_use_file_io() {
56            Self::RandomAccessFile(RandomAccessFile::open(path)?)
57        } else {
58            Self::Mmap(Mmap::map_path(path)?)
59        })
60    }
61}
62
63impl ReadAt for EitherMmapOrRandomAccessFile {
64    fn read_at(&self, pos: u64, buf: &mut [u8]) -> io::Result<usize> {
65        use EitherMmapOrRandomAccessFile::*;
66        match self {
67            Mmap(mmap) => mmap.read_at(pos, buf),
68            RandomAccessFile(file) => file.read_at(pos, buf),
69        }
70    }
71}
72
73impl Size for EitherMmapOrRandomAccessFile {
74    fn size(&self) -> io::Result<Option<u64>> {
75        use EitherMmapOrRandomAccessFile::*;
76        match self {
77            Mmap(mmap) => mmap.size(),
78            RandomAccessFile(file) => file.size(),
79        }
80    }
81}
82
83// Use mmap by default, switch to file-io when `FOREST_CAR_LOADER_FILE_IO` is set to `1` or `true`
84crate::def_is_env_truthy!(should_use_file_io, "FOREST_CAR_LOADER_FILE_IO");
85
86#[cfg(test)]
87mod tests {
88    use std::fs;
89
90    use super::*;
91    use quickcheck_macros::quickcheck;
92
93    #[quickcheck]
94    fn test_mmap_read_at_and_size(bytes: Vec<u8>) -> anyhow::Result<()> {
95        let tmp = tempfile::Builder::new().tempfile()?.into_temp_path();
96        fs::write(&tmp, &bytes)?;
97        let mmap = Mmap::map(&fs::File::open(&tmp)?)?;
98
99        assert_eq!(mmap.size()?.unwrap_or_default() as usize, bytes.len());
100
101        let mut buffer = [0; 128];
102        for pos in 0..bytes.len() {
103            let size = mmap.read_at(pos as _, &mut buffer)?;
104            assert_eq!(&bytes[pos..(pos + size)], &buffer[..size]);
105        }
106
107        Ok(())
108    }
109
110    #[test]
111    fn test_out_of_band_mmap_read() {
112        let temp_file = tempfile::Builder::new()
113            .tempfile()
114            .unwrap()
115            .into_temp_path();
116        let mmap = Mmap::map(&fs::File::open(&temp_file).unwrap()).unwrap();
117
118        let mut buffer = [];
119        // This matches the behaviour for seeking past the end of a file
120        assert_eq!(mmap.read_at(u64::MAX, &mut buffer).unwrap(), 0);
121    }
122}
forest/utils/io/mmap.rs

forest/utils/io/
mmap.rs
