openvm_circuit/system/memory/online/
paged_vec.rs

1use std::fmt::Debug;
2
3use openvm_stark_backend::p3_maybe_rayon::prelude::*;
4
5#[derive(Debug, Clone)]
6pub struct PagedVec<T, const PAGE_SIZE: usize> {
7    pages: Vec<Option<Box<[T; PAGE_SIZE]>>>,
8}
9
10unsafe impl<T: Send, const PAGE_SIZE: usize> Send for PagedVec<T, PAGE_SIZE> {}
11unsafe impl<T: Sync, const PAGE_SIZE: usize> Sync for PagedVec<T, PAGE_SIZE> {}
12
13impl<T: Copy + Default, const PAGE_SIZE: usize> PagedVec<T, PAGE_SIZE> {
14    #[inline]
15    /// `total_size` is the capacity of elements of type `T`.
16    pub fn new(total_size: usize) -> Self {
17        let num_pages = total_size.div_ceil(PAGE_SIZE);
18        Self {
19            pages: vec![None; num_pages],
20        }
21    }
22
23    #[cold]
24    #[inline(never)]
25    fn create_zeroed_page() -> Box<[T; PAGE_SIZE]> {
26        // SAFETY:
27        // - layout is valid since PAGE_SIZE is non-zero
28        // - alloc_zeroed returns properly aligned memory for T
29        // - Box::from_raw takes ownership of the allocated memory
30        unsafe {
31            let layout = std::alloc::Layout::array::<T>(PAGE_SIZE).unwrap();
32            let ptr = std::alloc::alloc_zeroed(layout) as *mut [T; PAGE_SIZE];
33            Box::from_raw(ptr)
34        }
35    }
36
37    /// Get value at index without allocating new pages.
38    /// Panics if index is out of bounds. Returns default value if page doesn't exist.
39    #[inline]
40    pub fn get(&self, index: usize) -> T {
41        let page_idx = index / PAGE_SIZE;
42        let offset = index % PAGE_SIZE;
43
44        // SAFETY:
45        // - offset < PAGE_SIZE by construction (from modulo operation)
46        // - page exists when as_ref() returns Some
47        self.pages[page_idx]
48            .as_ref()
49            .map(|page| unsafe { *page.get_unchecked(offset) })
50            .unwrap_or_default()
51    }
52
53    /// Panics if the index is out of bounds. Creates new page before write when necessary.
54    #[inline]
55    pub fn set(&mut self, index: usize, value: T) {
56        let page_idx = index / PAGE_SIZE;
57        let offset = index % PAGE_SIZE;
58
59        let page = self.pages[page_idx].get_or_insert_with(Self::create_zeroed_page);
60
61        // SAFETY: offset < PAGE_SIZE by construction
62        unsafe {
63            *page.get_unchecked_mut(offset) = value;
64        }
65    }
66
67    pub fn par_iter(&self) -> impl ParallelIterator<Item = (usize, T)> + '_
68    where
69        T: Send + Sync,
70    {
71        self.pages
72            .par_iter()
73            .enumerate()
74            .filter_map(move |(page_idx, page)| {
75                page.as_ref().map(move |p| {
76                    p.par_iter()
77                        .enumerate()
78                        .map(move |(offset, &value)| (page_idx * PAGE_SIZE + offset, value))
79                })
80            })
81            .flatten()
82    }
83
84    pub fn iter(&self) -> impl Iterator<Item = (usize, T)> + '_
85    where
86        T: Send + Sync,
87    {
88        self.pages
89            .iter()
90            .enumerate()
91            .filter_map(move |(page_idx, page)| {
92                page.as_ref().map(move |p| {
93                    p.iter()
94                        .enumerate()
95                        .map(move |(offset, &value)| (page_idx * PAGE_SIZE + offset, value))
96                })
97            })
98            .flatten()
99    }
100}
openvm_circuit/system/memory/online/paged_vec.rs

openvm_circuit/system/memory/online/
paged_vec.rs
