kevy_map/
alloc.rs

1//! Allocation lifecycle for [`crate::KevyMap`] — `alloc_table` (the only
2//! growing constructor) and the matching `Drop` impl. Split out so
3//! [`crate::map`] stays under the 500-LOC house rule. Both halves dispatch
4//! between the global allocator and a 2 MiB-aligned `mmap` path (E13)
5//! based on the per-instance `mmap_backed` flag.
6
7use core::marker::PhantomData;
8use core::mem::MaybeUninit;
9use core::ptr;
10use core::ptr::NonNull;
11use std::alloc::{Layout, alloc, dealloc, handle_alloc_error};
12
13use crate::map::{EMPTY, GROUP_WIDTH, KevyMap, MIN_CAP, table_layout};
14
15/// Tables smaller than this stay on the global allocator. mmap's
16/// over-allocate-and-trim alignment trick costs one extra HP of address
17/// space + 2 munmap syscalls — only worth it once the payload itself
18/// approaches HP scale.
19const THP_BACKED_THRESHOLD: usize = 1024 * 1024; // 1 MiB
20
21impl<K, V> KevyMap<K, V> {
22    /// Allocate a freshly-zeroed table sized for `cap` slots. `cap` must be
23    /// a power of two and ≥ `MIN_CAP`. Used by [`crate::KevyMap::with_capacity`]
24    /// and by the growth path on rehash.
25    ///
26    /// E13: when the combined layout is ≥ [`THP_BACKED_THRESHOLD`], allocate
27    /// directly via 2 MiB-aligned `mmap` so the kernel's `khugepaged` can
28    /// actually promote the region to 2 MiB pages. With the global
29    /// allocator (jemalloc-like chunk placement) the base pointer is only
30    /// 4 KiB-aligned, so `khugepaged` cannot find a candidate even when
31    /// `MADV_HUGEPAGE` is set — observed as `AnonHugePages: 0 kB` in
32    /// `/proc/PID/smaps` despite the hint.
33    pub(crate) fn alloc_table(cap: usize) -> Self {
34        debug_assert!(cap.is_power_of_two());
35        debug_assert!(cap >= MIN_CAP);
36
37        let (layout, meta_offset) = table_layout::<(K, V)>(cap);
38        let (base, mmap_backed) = if layout.size() >= THP_BACKED_THRESHOLD {
39            if let Some(p) = kevy_madvise::mmap_anon_aligned_2mb(layout.size()) {
40                (p.as_ptr(), true)
41            } else {
42                (fallback_alloc(layout), false)
43            }
44        } else {
45            (fallback_alloc(layout), false)
46        };
47        // Initialise the metadata range (real + mirror tail) to EMPTY in a
48        // single memset. The slot array is left uninitialised — slots
49        // become initialised only when their metadata byte transitions
50        // out of the high-bit-set state (EMPTY/DELETED).
51        let meta_byte_ptr = unsafe { base.add(meta_offset) };
52        unsafe { ptr::write_bytes(meta_byte_ptr, EMPTY, cap + GROUP_WIDTH) };
53
54        let slots_ptr = base.cast::<MaybeUninit<(K, V)>>();
55        let metadata_ptr = meta_byte_ptr;
56
57        // Re-hint THP on the buffer. On the `mmap_backed = true` path the
58        // 2 MiB-aligned mmap already called MADV_HUGEPAGE inside
59        // `mmap_anon_aligned_2mb`; this second call is redundant there but
60        // harmless. On the global-alloc fallback path it's the only hint.
61        if !mmap_backed {
62            kevy_madvise::advise_hugepage(base.cast_const(), layout.size());
63        }
64
65        Self {
66            // SAFETY: alloc returned non-null; raw pointers are derived
67            // within the same allocation.
68            slots_ptr: unsafe { NonNull::new_unchecked(slots_ptr) },
69            metadata_ptr: unsafe { NonNull::new_unchecked(metadata_ptr) },
70            cap,
71            mask: cap - 1,
72            occupied: 0,
73            deleted: 0,
74            mmap_backed,
75            _marker: PhantomData,
76        }
77    }
78}
79
80/// Global-allocator path that aborts on OOM. Cohesive helper so both
81/// branches of `alloc_table` can call it.
82fn fallback_alloc(layout: Layout) -> *mut u8 {
83    // SAFETY: layout has non-zero size (metadata alone is ≥ MIN_CAP +
84    // GROUP_WIDTH - 1 ≥ 31 bytes). alloc returns either a valid
85    // allocation of `layout` or null.
86    let p = unsafe { alloc(layout) };
87    if p.is_null() {
88        handle_alloc_error(layout);
89    }
90    p
91}
92
93impl<K, V> Drop for KevyMap<K, V> {
94    fn drop(&mut self) {
95        if self.cap == 0 {
96            return;
97        }
98        if std::mem::needs_drop::<(K, V)>() {
99            for i in 0..self.cap {
100                // SAFETY: i < cap ⇒ in-bounds.
101                let meta = unsafe { *self.metadata_ptr.as_ptr().add(i) };
102                if meta & 0x80 == 0 {
103                    // SAFETY: full slot ⇒ initialised.
104                    unsafe {
105                        ptr::drop_in_place(self.slots_ptr.as_ptr().add(i).cast::<(K, V)>());
106                    }
107                }
108            }
109        }
110        // Free the single combined allocation. `slots_ptr` IS the base of
111        // the allocation (see `alloc_table`'s layout computation: slots are
112        // at offset 0; metadata sits at meta_offset).
113        let (layout, _) = table_layout::<(K, V)>(self.cap);
114        // SAFETY: cap > 0 ⇒ slots_ptr is non-null and was returned by either
115        // `alloc` or `mmap_anon_aligned_2mb` with the same `layout.size()`;
116        // `mmap_backed` records which path was used so dealloc matches.
117        if self.mmap_backed {
118            // SAFETY: slots_ptr came from mmap_anon_aligned_2mb with
119            // layout.size(); munmap_2mb rounds the len back up internally.
120            unsafe {
121                kevy_madvise::munmap_2mb(self.slots_ptr.cast(), layout.size());
122            }
123        } else {
124            // SAFETY: slots_ptr came from `alloc` with this layout.
125            unsafe {
126                dealloc(self.slots_ptr.as_ptr().cast::<u8>(), layout);
127            }
128        }
129    }
130}
kevy_map/alloc.rs

kevy_map/
alloc.rs
