swh_graph/utils/
mmap.rs

1// Copyright (C) 2023-2026  The Software Heritage developers
2// See the AUTHORS file at the top-level directory of this distribution
3// License: GNU General Public License version 3, or any later version
4// See top-level LICENSE file for more information
5
6use std::fs::File;
7use std::marker::PhantomData;
8use std::path::Path;
9
10use anyhow::{bail, Context, Result};
11use byteorder::ByteOrder;
12use mmap_rs::{Mmap, MmapFlags};
13
14/// Newtype for [`Mmap`] used to store arrays of any integers
15///
16/// instead of slices of u8
17pub struct NumberMmap<E: ByteOrder, N: common_traits::AsBytes, B> {
18    data: B,
19    len: usize,
20    offset: usize,
21    _number: PhantomData<N>,
22    _endianness: PhantomData<E>,
23}
24
25impl<E: ByteOrder, N: common_traits::AsBytes> NumberMmap<E, N, Mmap> {
26    pub fn new<P: AsRef<Path>>(path: P, len: usize) -> Result<NumberMmap<E, N, Mmap>> {
27        let path = path.as_ref();
28        let file_len = path
29            .metadata()
30            .with_context(|| format!("Could not stat {}", path.display()))?
31            .len();
32        if file_len < (len * N::BYTES) as u64 {
33            // We have to allow length > num_nodes because graphs compressed
34            // with the Java implementation used zero padding at the end
35            bail!(
36                "{} is too short: expected at least {} bytes ({} items), got {}",
37                path.display(),
38                len * N::BYTES,
39                len,
40                file_len,
41            );
42        }
43        let file =
44            File::open(path).with_context(|| format!("Could not open {}", path.display()))?;
45        Self::with_file_and_offset(path, len, file, 0)
46    }
47
48    pub fn with_file_and_offset<P: AsRef<Path>>(
49        path: P,
50        len: usize,
51        file: File,
52        offset: usize,
53    ) -> Result<NumberMmap<E, N, Mmap>> {
54        let path = path.as_ref();
55        let file_len = len * N::BYTES;
56        let data = unsafe {
57            mmap_rs::MmapOptions::new(file_len as _)
58                .with_context(|| format!("Could not initialize mmap of size {file_len}"))?
59                .with_flags(MmapFlags::TRANSPARENT_HUGE_PAGES | MmapFlags::RANDOM_ACCESS)
60                .with_file(&file, 0)
61                .map()
62                .with_context(|| format!("Could not mmap {}", path.display()))?
63        };
64
65        if data.len() % N::BYTES != 0 {
66            bail!(
67                "Cannot interpret mmap of size {} as array of {}",
68                data.len(),
69                std::any::type_name::<N>()
70            );
71        }
72        Ok(NumberMmap {
73            data,
74            len,
75            offset,
76            _number: PhantomData,
77            _endianness: PhantomData,
78        })
79    }
80
81    #[allow(clippy::len_without_is_empty)]
82    #[inline(always)]
83    pub fn len(&self) -> usize {
84        self.len
85    }
86}
87
88impl<E: ByteOrder, N: common_traits::AsBytes> NumberMmap<E, N, Mmap> {
89    fn get_slice(&self, index: usize) -> Option<&[u8]> {
90        let start = (index * N::BYTES) + self.offset;
91        self.data.get(start..(start + N::BYTES))
92    }
93
94    unsafe fn get_slice_unchecked(&self, index: usize) -> &[u8] {
95        let start = (index * N::BYTES) + self.offset;
96        self.data.get_unchecked(start..(start + N::BYTES))
97    }
98}
99
100macro_rules! impl_number_mmap {
101    ($ty:ty, $fn:ident) => {
102        impl<E: ByteOrder> crate::utils::GetIndex for &NumberMmap<E, $ty, Mmap> {
103            type Output = $ty;
104
105            #[inline(always)]
106            fn len(&self) -> usize {
107                NumberMmap::len(self)
108            }
109
110            /// Returns an item
111            #[inline(always)]
112            fn get(&self, index: usize) -> Option<$ty> {
113                self.get_slice(index).map(E::$fn)
114            }
115
116            /// Returns an item
117            ///
118            /// # Safety
119            ///
120            /// Undefined behavior if `index >= len()`
121            #[inline(always)]
122            unsafe fn get_unchecked(&self, index: usize) -> $ty {
123                E::$fn(self.get_slice_unchecked(index))
124            }
125        }
126    };
127}
128
129impl_number_mmap!(i16, read_i16);
130impl_number_mmap!(u16, read_u16);
131impl_number_mmap!(i32, read_i32);
132impl_number_mmap!(u32, read_u32);
133impl_number_mmap!(f32, read_f32);
134impl_number_mmap!(i64, read_i64);
135impl_number_mmap!(u64, read_u64);
136impl_number_mmap!(f64, read_f64);
swh_graph/utils/mmap.rs

swh_graph/utils/
mmap.rs
