sparse_mmap/
alloc.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Memory functionality that needs a refactor.
5
6use std::ops::Deref;
7use std::ops::DerefMut;
8use std::slice;
9use std::sync::atomic::AtomicU8;
10#[cfg(unix)]
11use unix as sys;
12#[cfg(windows)]
13use windows as sys;
14
15#[derive(Debug)]
16pub struct Allocation {
17    ptr: *mut u8,
18    size: usize,
19    _dummy: std::marker::PhantomData<[u8]>,
20}
21
22unsafe impl Send for Allocation {}
23unsafe impl Sync for Allocation {}
24
25impl Allocation {
26    pub fn new(size: usize) -> Result<Self, std::io::Error> {
27        let ptr = sys::alloc(size)?;
28        Ok(Allocation {
29            ptr,
30            size,
31            _dummy: std::marker::PhantomData,
32        })
33    }
34}
35
36impl DerefMut for Allocation {
37    fn deref_mut(&mut self) -> &mut [u8] {
38        unsafe { slice::from_raw_parts_mut(self.ptr, self.size) }
39    }
40}
41
42impl Deref for Allocation {
43    type Target = [u8];
44
45    fn deref(&self) -> &[u8] {
46        unsafe { slice::from_raw_parts(self.ptr, self.size) }
47    }
48}
49
50impl Drop for Allocation {
51    fn drop(&mut self) {
52        unsafe {
53            sys::free(self.ptr, self.size);
54        }
55    }
56}
57
58#[derive(Debug)]
59pub struct SharedMem {
60    alloc: Allocation,
61}
62
63impl SharedMem {
64    pub fn new(alloc: Allocation) -> Self {
65        SharedMem { alloc }
66    }
67}
68
69impl Deref for SharedMem {
70    type Target = [AtomicU8];
71
72    fn deref(&self) -> &Self::Target {
73        unsafe { slice::from_raw_parts(self.alloc.ptr as *const AtomicU8, self.alloc.size) }
74    }
75}
76
77#[cfg(windows)]
78mod windows {
79    use std::ptr;
80    use windows_sys::Win32::System::Memory::MEM_COMMIT;
81    use windows_sys::Win32::System::Memory::MEM_RELEASE;
82    use windows_sys::Win32::System::Memory::MEM_RESERVE;
83    use windows_sys::Win32::System::Memory::PAGE_READWRITE;
84    use windows_sys::Win32::System::Memory::VirtualAlloc;
85    use windows_sys::Win32::System::Memory::VirtualFree;
86
87    pub fn alloc(size: usize) -> std::io::Result<*mut u8> {
88        let ptr = unsafe {
89            VirtualAlloc(
90                ptr::null_mut(),
91                size,
92                MEM_RESERVE | MEM_COMMIT,
93                PAGE_READWRITE,
94            )
95        };
96        if ptr.is_null() {
97            return Err(std::io::Error::last_os_error());
98        }
99        Ok(ptr.cast::<u8>())
100    }
101
102    pub unsafe fn free(ptr: *mut u8, _size: usize) {
103        let ret = unsafe { VirtualFree(ptr.cast(), 0, MEM_RELEASE) };
104        assert!(ret != 0);
105    }
106}
107
108#[cfg(unix)]
109mod unix {
110    use std::ptr;
111
112    pub fn alloc(size: usize) -> std::io::Result<*mut u8> {
113        let ptr = unsafe {
114            libc::mmap(
115                ptr::null_mut(),
116                size,
117                libc::PROT_READ | libc::PROT_WRITE,
118                libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
119                -1,
120                0,
121            )
122        };
123        if ptr == libc::MAP_FAILED {
124            return Err(std::io::Error::last_os_error());
125        }
126        Ok(ptr.cast::<u8>())
127    }
128
129    pub unsafe fn free(ptr: *mut u8, size: usize) {
130        let ret = unsafe { libc::munmap(ptr.cast::<libc::c_void>(), size) };
131        assert!(ret == 0);
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[test]
140    fn test_alloc_free() -> Result<(), Box<dyn std::error::Error>> {
141        unsafe {
142            let x = sys::alloc(4096)?;
143            sys::free(x, 4096);
144            Ok(())
145        }
146    }
147}