openvm_circuit/system/memory/online/
memmap.rs

1use std::{
2    fmt::Debug,
3    mem::{align_of, size_of, size_of_val},
4};
5
6use memmap2::MmapMut;
7
8use super::{LinearMemory, PAGE_SIZE};
9
10pub const CELL_STRIDE: usize = 1;
11
12/// Mmap-backed linear memory. OS-memory pages are paged in on-demand and zero-initialized.
13#[derive(Debug)]
14pub struct MmapMemory {
15    mmap: MmapMut,
16    size: usize,
17}
18
19impl Clone for MmapMemory {
20    fn clone(&self) -> Self {
21        let mut new_mmap = MmapMut::map_anon(self.mmap.len()).unwrap();
22        new_mmap.copy_from_slice(&self.mmap);
23        Self {
24            mmap: new_mmap,
25            size: self.size,
26        }
27    }
28}
29
30impl MmapMemory {
31    #[inline(always)]
32    pub fn as_ptr(&self) -> *const u8 {
33        self.mmap.as_ptr()
34    }
35
36    #[inline(always)]
37    pub fn as_mut_ptr(&mut self) -> *mut u8 {
38        self.mmap.as_mut_ptr()
39    }
40
41    #[cfg(not(feature = "unprotected"))]
42    #[inline(always)]
43    fn check_bounds(&self, start: usize, size: usize) {
44        let memory_size = self.size();
45        if start > memory_size || size > memory_size - start {
46            panic_oob(start, size, memory_size);
47        }
48    }
49
50    #[cfg(feature = "unprotected")]
51    #[inline(always)]
52    fn check_bounds(&self, start: usize, size: usize) {
53        let memory_size = self.size();
54        debug_assert!(
55            start <= memory_size && size <= memory_size - start,
56            "Memory access out of bounds: start={} size={} memory_size={}",
57            start,
58            size,
59            memory_size
60        );
61    }
62}
63
64impl LinearMemory for MmapMemory {
65    /// Create a new MmapMemory with the given `size` in bytes.
66    /// We round `size` up to be a multiple of the mmap page size (4kb by default).
67    fn new(size: usize) -> Self {
68        let mmap_size = size.div_ceil(PAGE_SIZE) * PAGE_SIZE;
69        // anonymous mapping means pages are zero-initialized on first use
70        Self {
71            mmap: MmapMut::map_anon(mmap_size).unwrap(),
72            size,
73        }
74    }
75
76    fn size(&self) -> usize {
77        self.size
78    }
79
80    fn as_slice(&self) -> &[u8] {
81        &self.mmap[..self.size]
82    }
83
84    fn as_mut_slice(&mut self) -> &mut [u8] {
85        &mut self.mmap[..self.size]
86    }
87
88    #[cfg(target_os = "linux")]
89    fn fill_zero(&mut self) {
90        use libc::{madvise, MADV_DONTNEED};
91
92        let mmap = &mut self.mmap;
93        // SAFETY: our mmap is a memory-backed (not file-backed) anonymous private mapping.
94        // When we madvise MADV_DONTNEED, according to https://man7.org/linux/man-pages/man2/madvise.2.html
95        // > subsequent accesses of pages in the range will succeed, but
96        // > will result in either repopulating the memory contents from
97        // > the up-to-date contents of the underlying mapped file (for
98        // > shared file mappings, shared anonymous mappings, and shmem-
99        // > based techniques such as System V shared memory segments)
100        // > or zero-fill-on-demand pages for anonymous private
101        // > mappings.
102        unsafe {
103            let ret = madvise(
104                mmap.as_ptr() as *mut libc::c_void,
105                mmap.len(),
106                MADV_DONTNEED,
107            );
108            if ret != 0 {
109                // Fallback to write_bytes if madvise fails
110                std::ptr::write_bytes(mmap.as_mut_ptr(), 0, mmap.len());
111            }
112        }
113    }
114
115    #[inline(always)]
116    unsafe fn read<BLOCK: Copy>(&self, from: usize) -> BLOCK {
117        self.check_bounds(from, size_of::<BLOCK>());
118        let src = self.as_ptr().add(from) as *const BLOCK;
119        // SAFETY:
120        // - Bounds checked above (unless unprotected feature enabled)
121        // - We assume `src` is aligned to `BLOCK`
122        // - We assume `BLOCK` is "plain old data" so the underlying `src` bytes is valid to read as
123        //   an initialized value of `BLOCK`
124        core::ptr::read(src)
125    }
126
127    #[inline(always)]
128    unsafe fn read_unaligned<BLOCK: Copy>(&self, from: usize) -> BLOCK {
129        self.check_bounds(from, size_of::<BLOCK>());
130        let src = self.as_ptr().add(from) as *const BLOCK;
131        // SAFETY:
132        // - Bounds checked above (unless unprotected feature enabled)
133        // - We assume `BLOCK` is "plain old data" so the underlying `src` bytes is valid to read as
134        //   an initialized value of `BLOCK`
135        core::ptr::read_unaligned(src)
136    }
137
138    #[inline(always)]
139    unsafe fn write<BLOCK: Copy>(&mut self, start: usize, values: BLOCK) {
140        self.check_bounds(start, size_of::<BLOCK>());
141        let dst = self.as_mut_ptr().add(start) as *mut BLOCK;
142        // SAFETY:
143        // - Bounds checked above (unless unprotected feature enabled)
144        // - We assume `dst` is aligned to `BLOCK`
145        core::ptr::write(dst, values);
146    }
147
148    #[inline(always)]
149    unsafe fn write_unaligned<BLOCK: Copy>(&mut self, start: usize, values: BLOCK) {
150        self.check_bounds(start, size_of::<BLOCK>());
151        let dst = self.as_mut_ptr().add(start) as *mut BLOCK;
152        // SAFETY:
153        // - Bounds checked above (unless unprotected feature enabled)
154        core::ptr::write_unaligned(dst, values);
155    }
156
157    #[inline(always)]
158    unsafe fn swap<BLOCK: Copy>(&mut self, start: usize, values: &mut BLOCK) {
159        self.check_bounds(start, size_of::<BLOCK>());
160        // SAFETY:
161        // - Bounds checked above (unless unprotected feature enabled)
162        // - We assume `start` is aligned to `BLOCK`
163        core::ptr::swap(
164            self.as_mut_ptr().add(start) as *mut BLOCK,
165            values as *mut BLOCK,
166        );
167    }
168
169    #[inline(always)]
170    unsafe fn copy_nonoverlapping<T: Copy>(&mut self, to: usize, data: &[T]) {
171        self.check_bounds(to, size_of_val(data));
172        debug_assert_eq!(PAGE_SIZE % align_of::<T>(), 0);
173        let src = data.as_ptr();
174        let dst = self.as_mut_ptr().add(to) as *mut T;
175        // SAFETY:
176        // - Bounds checked above (unless unprotected feature enabled)
177        // - Assumes `to` is aligned to `T` and `self.as_mut_ptr()` is aligned to `T`, which implies
178        //   the same for `dst`.
179        core::ptr::copy_nonoverlapping::<T>(src, dst, data.len());
180    }
181
182    #[inline(always)]
183    unsafe fn get_aligned_slice<T: Copy>(&self, start: usize, len: usize) -> &[T] {
184        self.check_bounds(start, len * size_of::<T>());
185        let data = self.as_ptr().add(start) as *const T;
186        // SAFETY:
187        // - Bounds checked above (unless unprotected feature enabled)
188        // - Assumes `data` is aligned to `T`
189        // - `T` is "plain old data" (POD), so conversion from underlying bytes is properly
190        //   initialized
191        // - `self` will not be mutated while borrowed
192        core::slice::from_raw_parts(data, len)
193    }
194}
195
196#[cold]
197#[inline(never)]
198fn panic_oob(start: usize, size: usize, memory_size: usize) -> ! {
199    panic!("Memory access out of bounds: start={start} size={size} memory_size={memory_size}");
200}
openvm_circuit/system/memory/online/memmap.rs

openvm_circuit/system/memory/online/
memmap.rs
