gtfsort/
mmap.rs

1use std::marker::PhantomData;
2
3#[cfg(all(not(unix), not(windows)))]
4compile_error!(
5    "mmap is only supported on Unix and Windows platforms, please compile without the mmap feature"
6);
7
8#[cfg(windows)]
9macro_rules! high32 {
10    ($x:expr) => {
11        ($x >> 32) as u32
12    };
13}
14
15#[cfg(windows)]
16macro_rules! low32 {
17    ($x:expr) => {
18        $x as u32
19    };
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum Madvice {
24    Normal,
25    Random,
26    Sequential,
27    WillNeed,
28    DontNeed,
29    HugePage,
30}
31
32type CleanupFn<S> = Box<dyn FnOnce(&mut S) -> std::io::Result<()>>;
33
34pub struct MemoryMap<'a, T> {
35    ptr: *const T,
36    size: usize,
37    cleanup: Option<CleanupFn<Self>>,
38    _marker: PhantomData<&'a T>,
39}
40
41impl<'a, T> MemoryMap<'a, T> {
42    /// Creates a new MemoryMap instance from a pointer and size.
43    ///
44    /// # Safety
45    /// ptr must be a valid pointer to a memory-mapped region of size bytes.
46    pub unsafe fn new(ptr: *const T, size: usize) -> Self {
47        Self {
48            ptr,
49            size,
50            cleanup: None,
51            _marker: PhantomData,
52        }
53    }
54
55    /// Returns the mapped region size in bytes.
56    pub fn size_bytes(&self) -> usize {
57        self.size
58    }
59
60    /// Borrows the mapped region as an immutable slice.
61    pub fn as_slice(&self) -> &[T] {
62        if self.size == 0 {
63            return &[];
64        }
65        unsafe { std::slice::from_raw_parts(self.ptr, self.size / std::mem::size_of::<T>()) }
66    }
67
68    #[cfg(unix)]
69    /// Applies one or more access-pattern hints to the mapped region on Unix.
70    pub fn madvise(&self, advice: &[Madvice]) -> Result<(), std::io::Error> {
71        if self.ptr.is_null() {
72            return Ok(());
73        }
74
75        #[allow(unreachable_patterns)]
76        let advice = advice.iter().fold(0, |acc, &a| {
77            acc | match a {
78                Madvice::Normal => libc::MADV_NORMAL,
79                Madvice::Random => libc::MADV_RANDOM,
80                Madvice::Sequential => libc::MADV_SEQUENTIAL,
81                Madvice::WillNeed => libc::MADV_WILLNEED,
82                Madvice::DontNeed => libc::MADV_DONTNEED,
83                #[cfg(target_os = "linux")]
84                Madvice::HugePage => libc::MADV_HUGEPAGE,
85                _ => 0,
86            }
87        });
88
89        let ret = unsafe { libc::madvise(self.ptr as *mut _, self.size, advice) };
90
91        if ret == -1 {
92            return Err(std::io::Error::last_os_error());
93        }
94
95        Ok(())
96    }
97
98    #[cfg(not(unix))]
99    /// Accepts access-pattern hints as a no-op on non-Unix platforms.
100    pub fn madvise(&self, _advice: &[Madvice]) -> Result<(), std::io::Error> {
101        Ok(())
102    }
103
104    #[cfg(unix)]
105    /// Creates a new MemoryMap instance from a file descriptor and size.
106    ///
107    /// # Safety
108    /// fd must be a valid file descriptor.
109    /// The file descriptor must be open and readable.
110    /// Size must be a valid size for the file descriptor.
111    pub unsafe fn from_file<F>(fd: &'a F, size: usize) -> Result<Self, std::io::Error>
112    where
113        F: std::os::unix::io::AsRawFd,
114    {
115        if size == 0 {
116            return Ok(Self {
117                ptr: std::ptr::null(),
118                size,
119                cleanup: None,
120                _marker: PhantomData,
121            });
122        }
123
124        let ptr = libc::mmap(
125            std::ptr::null_mut(),
126            size,
127            libc::PROT_READ,
128            libc::MAP_SHARED,
129            fd.as_raw_fd(),
130            0,
131        );
132
133        if ptr == libc::MAP_FAILED {
134            return Err(std::io::Error::last_os_error());
135        }
136
137        Ok(Self {
138            ptr: ptr as *const T,
139            size,
140            cleanup: Some(Box::new(move |this| unsafe {
141                let ret = libc::munmap(this.ptr as *mut _, this.size);
142                if ret == -1 {
143                    let e = std::io::Error::last_os_error();
144                    log::warn!("munmap error: {}", e);
145                    Err(e)
146                } else {
147                    Ok(())
148                }
149            })),
150            _marker: PhantomData,
151        })
152    }
153
154    #[cfg(windows)]
155    /// Creates a new MemoryMap instance from a file handle and size.
156    ///
157    /// # Safety
158    /// handle must be a valid file handle.
159    /// The file handle must be open and readable.
160    /// Size must be a valid size for the file handle.
161    pub unsafe fn from_handle<F>(handle: &'a F, size: usize) -> Result<Self, std::io::Error>
162    where
163        F: std::os::windows::io::AsRawHandle,
164    {
165        use windows::{
166            core::*,
167            Win32::{
168                Foundation::{CloseHandle, HANDLE},
169                System::Memory::*,
170            },
171        };
172
173        if size == 0 {
174            return Ok(Self {
175                ptr: std::ptr::null(),
176                size,
177                cleanup: None,
178                _marker: PhantomData,
179            });
180        }
181
182        unsafe {
183            let handle = CreateFileMappingW(
184                HANDLE(handle.as_raw_handle()),
185                None,
186                PAGE_READONLY,
187                high32!(size),
188                low32!(size),
189                PCWSTR::null(),
190            )?;
191
192            if handle.0.is_null() {
193                return Err(std::io::Error::last_os_error());
194            }
195
196            let ptr = MapViewOfFile(handle, FILE_MAP_READ, 0, 0, size);
197
198            if ptr.Value.is_null() {
199                return Err(std::io::Error::last_os_error());
200            }
201
202            Ok(Self {
203                ptr: ptr.Value as *const T,
204                size,
205                cleanup: Some(Box::new(move |_this| {
206                    UnmapViewOfFile(ptr)?;
207                    CloseHandle(handle)?;
208                    Ok(())
209                })),
210                _marker: PhantomData,
211            })
212        }
213    }
214
215    /// Explicitly unmaps the region and reports cleanup errors.
216    pub fn close(mut self) -> Result<(), std::io::Error> {
217        if let Some(cleanup) = self.cleanup.take() {
218            cleanup(&mut self)?;
219        }
220        Ok(())
221    }
222}
223
224impl<T> Drop for MemoryMap<'_, T> {
225    /// Unmaps the region when the mapping goes out of scope.
226    fn drop(&mut self) {
227        if let Some(cleanup) = self.cleanup.take() {
228            cleanup(self).expect("failed to unmap memory, and error was ignored");
229        }
230    }
231}
232
233pub struct MemoryMapMut<'a, T> {
234    ptr: *mut T,
235    size: usize,
236    cleanup: Option<CleanupFn<Self>>,
237    _marker: PhantomData<&'a mut T>,
238}
239
240impl<'a, T> MemoryMapMut<'a, T> {
241    /// Creates a new MemoryMapMut instance from a pointer and size.
242    ///
243    /// # Safety
244    /// ptr must be a valid pointer to a memory-mapped region of size bytes.
245    pub unsafe fn new(ptr: *mut T, size: usize) -> Self {
246        Self {
247            ptr,
248            size,
249            cleanup: None,
250            _marker: PhantomData,
251        }
252    }
253
254    /// Borrows the mapped region as an immutable slice.
255    pub fn as_slice(&self) -> &[T] {
256        if self.size == 0 {
257            return &[];
258        }
259        unsafe { std::slice::from_raw_parts(self.ptr, self.size / std::mem::size_of::<T>()) }
260    }
261
262    /// Borrows the mapped region as a mutable slice.
263    pub fn as_mut_slice(&mut self) -> &mut [T] {
264        if self.size == 0 {
265            return &mut [];
266        }
267        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.size / std::mem::size_of::<T>()) }
268    }
269
270    #[cfg(unix)]
271    /// Applies one or more access-pattern hints to the mutable mapping on Unix.
272    pub fn madvise(&self, advice: &[Madvice]) -> Result<(), std::io::Error> {
273        if self.ptr.is_null() {
274            return Ok(());
275        }
276        #[allow(unreachable_patterns)]
277        let advice = advice.iter().fold(0, |acc, &a| {
278            acc | match a {
279                Madvice::Normal => libc::MADV_NORMAL,
280                Madvice::Random => libc::MADV_RANDOM,
281                Madvice::Sequential => libc::MADV_SEQUENTIAL,
282                Madvice::WillNeed => libc::MADV_WILLNEED,
283                Madvice::DontNeed => libc::MADV_DONTNEED,
284                #[cfg(target_os = "linux")]
285                Madvice::HugePage => libc::MADV_HUGEPAGE,
286                _ => 0,
287            }
288        });
289
290        let ret = unsafe { libc::madvise(self.ptr as *mut _, self.size, advice) };
291
292        if ret == -1 {
293            return Err(std::io::Error::last_os_error());
294        }
295
296        Ok(())
297    }
298
299    #[cfg(not(unix))]
300    /// Accepts access-pattern hints as a no-op on non-Unix platforms.
301    pub fn madvise(&self, _advice: &[Madvice]) -> Result<(), std::io::Error> {
302        Ok(())
303    }
304
305    #[cfg(unix)]
306    /// Creates a new MemoryMapMut instance from a file descriptor and size.
307    ///
308    /// # Safety
309    /// fd must be a valid file descriptor.
310    /// The file descriptor must be open and readable.
311    /// Size must be a valid size for the file descriptor.
312    pub unsafe fn from_file<F>(fd: &'a F, size: usize) -> Result<Self, std::io::Error>
313    where
314        F: std::os::unix::io::AsRawFd,
315    {
316        if size == 0 {
317            return Ok(Self {
318                ptr: std::ptr::null_mut(),
319                size,
320                cleanup: None,
321                _marker: PhantomData,
322            });
323        }
324
325        let ptr = libc::mmap(
326            std::ptr::null_mut(),
327            size,
328            libc::PROT_READ | libc::PROT_WRITE,
329            libc::MAP_SHARED,
330            fd.as_raw_fd(),
331            0,
332        );
333
334        if ptr == libc::MAP_FAILED {
335            return Err(std::io::Error::last_os_error());
336        }
337
338        Ok(Self {
339            ptr: ptr as *mut T,
340            size,
341            cleanup: Some(Box::new(move |this| unsafe {
342                let ret = libc::munmap(this.ptr as *mut _, this.size);
343                if ret == -1 {
344                    let e = std::io::Error::last_os_error();
345                    eprintln!("munmap failed: {}", e);
346                    Err(e)
347                } else {
348                    Ok(())
349                }
350            })),
351            _marker: PhantomData,
352        })
353    }
354
355    #[cfg(windows)]
356    /// Creates a new MemoryMapMut instance from a file handle and size.
357    /// # Safety
358    /// handle must be a valid file handle.
359    /// The file handle must be open and writable.
360    /// Size must be a valid size for the file handle.
361    pub unsafe fn from_handle<F>(handle: &'a F, size: usize) -> Result<Self, std::io::Error>
362    where
363        F: std::os::windows::io::AsRawHandle,
364    {
365        use windows::{
366            core::*,
367            Win32::{
368                Foundation::{CloseHandle, HANDLE},
369                System::Memory::*,
370            },
371        };
372
373        if size == 0 {
374            return Ok(Self {
375                ptr: std::ptr::null_mut(),
376                size,
377                cleanup: None,
378                _marker: PhantomData,
379            });
380        }
381
382        unsafe {
383            let handle = CreateFileMappingW(
384                HANDLE(handle.as_raw_handle()),
385                None,
386                PAGE_READWRITE,
387                high32!(size),
388                low32!(size),
389                PCWSTR::null(),
390            )?;
391
392            if handle.0.is_null() {
393                return Err(std::io::Error::last_os_error());
394            }
395
396            let ptr = MapViewOfFile(handle, FILE_MAP_WRITE, 0, 0, size);
397
398            if ptr.Value.is_null() {
399                return Err(std::io::Error::last_os_error());
400            }
401
402            Ok(Self {
403                ptr: ptr.Value as *mut T,
404                size,
405                cleanup: Some(Box::new(move |_this| {
406                    UnmapViewOfFile(ptr)?;
407                    CloseHandle(handle)?;
408                    Ok(())
409                })),
410                _marker: PhantomData,
411            })
412        }
413    }
414
415    /// Explicitly unmaps the mutable region and reports cleanup errors.
416    pub fn close(mut self) -> Result<(), std::io::Error> {
417        if let Some(cleanup) = self.cleanup.take() {
418            cleanup(&mut self)?;
419        }
420        Ok(())
421    }
422}
423
424impl<T> Drop for MemoryMapMut<'_, T> {
425    /// Unmaps the mutable region when the mapping goes out of scope.
426    fn drop(&mut self) {
427        if let Some(cleanup) = self.cleanup.take() {
428            cleanup(self).expect("failed to unmap memory, and error was ignored");
429        }
430    }
431}
432
433#[cfg(test)]
434mod tests {
435    use super::*;
436    use std::{fs::OpenOptions, io::Write, path::PathBuf, sync::atomic::AtomicU64};
437
438    static COUNTER: AtomicU64 = AtomicU64::new(0);
439
440    /// Creates a uniquely named read-only mmap test fixture.
441    fn tempfile_ro(data: &[u8]) -> (PathBuf, std::fs::File) {
442        let path = std::env::temp_dir().join(format!(
443            "gtfsort_mmap_test_{}",
444            COUNTER.fetch_add(1, std::sync::atomic::Ordering::AcqRel)
445        ));
446
447        let mut file = OpenOptions::new()
448            .read(true)
449            .write(true)
450            .create(true)
451            .truncate(true)
452            .open(&path)
453            .unwrap();
454
455        file.write_all(data).unwrap();
456
457        file.flush().unwrap();
458        drop(file);
459
460        (
461            path.clone(),
462            OpenOptions::new().read(true).open(&path).unwrap(),
463        )
464    }
465
466    /// Creates a uniquely named read-write mmap test fixture.
467    fn tempfile_rw(data: &[u8]) -> (PathBuf, std::fs::File) {
468        let path = std::env::temp_dir().join(format!(
469            "gtfsort_mmap_test_{}",
470            COUNTER.fetch_add(1, std::sync::atomic::Ordering::AcqRel)
471        ));
472
473        let mut file = OpenOptions::new()
474            .read(true)
475            .write(true)
476            .create(true)
477            .truncate(true)
478            .open(&path)
479            .unwrap();
480
481        file.write_all(data).unwrap();
482
483        file.flush().unwrap();
484
485        (path, file)
486    }
487
488    #[test]
489    /// Verifies immutable file mapping and explicit cleanup.
490    fn test_mmap() {
491        let (path, file) = tempfile_ro(b"hello world");
492
493        #[cfg(unix)]
494        let mmap = unsafe { MemoryMap::<u8>::from_file(&file, 5).unwrap() };
495        #[cfg(windows)]
496        let mmap = unsafe { MemoryMap::<u8>::from_handle(&file, 5).unwrap() };
497
498        assert_eq!(mmap.as_slice(), b"hello");
499
500        mmap.close().unwrap();
501
502        drop(file);
503
504        assert_eq!(std::fs::read_to_string(path).unwrap(), "hello world");
505    }
506
507    #[test]
508    /// Verifies that mutable mappings persist changes to the file.
509    fn test_mmap_mut() {
510        let (path, file) = tempfile_rw(b"hello world");
511
512        #[cfg(unix)]
513        let mut mmap = unsafe { MemoryMapMut::<u8>::from_file(&file, 11).unwrap() };
514        #[cfg(windows)]
515        let mut mmap = unsafe { MemoryMapMut::<u8>::from_handle(&file, 11).unwrap() };
516
517        assert_eq!(mmap.as_slice(), b"hello world");
518
519        mmap.as_mut_slice()["hello ".len()..].copy_from_slice(b"WORLD");
520
521        assert_eq!(mmap.as_slice(), b"hello WORLD");
522
523        mmap.close().unwrap();
524
525        drop(file);
526
527        assert_eq!(std::fs::read_to_string(path).unwrap(), "hello WORLD");
528    }
529
530    #[test]
531    /// Verifies immutable mapping of an empty file.
532    fn test_mmap_zero_size() {
533        let (path, file) = tempfile_ro(b"");
534
535        #[cfg(unix)]
536        let mmap = unsafe { MemoryMap::<u8>::from_file(&file, 0).unwrap() };
537        #[cfg(windows)]
538        let mmap = unsafe { MemoryMap::<u8>::from_handle(&file, 0).unwrap() };
539
540        assert_eq!(mmap.as_slice(), b"");
541
542        mmap.close().unwrap();
543
544        drop(file);
545
546        assert_eq!(std::fs::read_to_string(path).unwrap(), "");
547    }
548
549    #[test]
550    /// Verifies mutable mapping of an empty file.
551    fn test_mmap_mut_zero_size() {
552        let (path, file) = tempfile_rw(b"");
553
554        #[cfg(unix)]
555        let mmap = unsafe { MemoryMapMut::<u8>::from_file(&file, 0).unwrap() };
556        #[cfg(windows)]
557        let mmap = unsafe { MemoryMapMut::<u8>::from_handle(&file, 0).unwrap() };
558
559        assert_eq!(mmap.as_slice(), b"");
560
561        mmap.close().unwrap();
562
563        drop(file);
564
565        assert_eq!(std::fs::read_to_string(path).unwrap(), "");
566    }
567
568    #[test]
569    /// Verifies that access-pattern advice can be applied.
570    fn test_mmap_madvise() {
571        let (path, file) = tempfile_ro(b"hello world");
572
573        #[cfg(unix)]
574        let mmap = unsafe { MemoryMap::<u8>::from_file(&file, 5).unwrap() };
575        #[cfg(windows)]
576        let mmap = unsafe { MemoryMap::<u8>::from_handle(&file, 5).unwrap() };
577
578        mmap.madvise(&[Madvice::Random]).unwrap();
579
580        mmap.close().unwrap();
581
582        drop(file);
583
584        assert_eq!(std::fs::read_to_string(path).unwrap(), "hello world");
585    }
586}
gtfsort/mmap.rs

gtfsort/
mmap.rs
