Skip to main content

plugmem_arena/
arena.rs

1//! Sharded sorted arena over one flat byte pool.
2//!
3//! Full v2 mechanics of: each shard owns a **chain of
4//! range-partitioned pages** — every page holds a sorted run of keys, pages
5//! in a chain follow each other in ascending key ranges (a B+-tree leaf
6//! level without interior nodes; chains stay short because shard hashing
7//! spreads the load). A full page **splits** in half instead of failing, and
8//! a page emptied by removals is unlinked into a **free-list** for reuse, so
9//! capacity is bounded only by [`ArenaCfg::max_bytes`].
10//!
11//! The structure is tuned for **wasm environments first**: `no_std + alloc`,
12//! a single linear byte pool (snapshot = memcpy), no threads, and the one
13//! measured `unsafe` (uninitialized page allocation) whose entire benefit is
14//! on the wasm allocation path — see [`Arena::insert`] internals and the
15//! crate-level documentation for the numbers.
16//!
17//! # Overlay opens
18//!
19//! [`Arena::load_overlay`] opens an arena whose existing pages are **borrowed**
20//! from a longer-lived buffer — typically a memory-mapped file — while the
21//! arena stays fully mutable. Because an arena mutates pages *in place* (slot
22//! shifts, page splits), it cannot use the append-only tail of
23//! [`BlobHeap`](crate::BlobHeap): instead the first write to a borrowed page
24//! copies just *that* page into an owned overlay, and pages grown after open
25//! live in an owned tail (per-page copy-on-write). The borrowed base is never
26//! mutated and never cloned as a whole, so a multi-gigabyte arena can be
27//! *written to* while resident only in the pages it actually touches. Staying
28//! true to the crate's flat philosophy, the overlay is two flat `Vec`s (a
29//! dense redirect and one contiguous copy pool) — no `Box`, no map. See
30//! `examples/overlay.rs`.
31
32use alloc::vec::Vec;
33use core::fmt;
34use core::marker::PhantomData;
35
36#[cfg(feature = "counters")]
37use core::cell::Cell;
38
39use crate::error::Error;
40use crate::paged::Paged;
41use crate::slot::Slot;
42
43/// Size of one arena page in bytes.
44///
45/// Fixed, not per-slot-count: whatever the slot size, a page is one
46/// L1-friendly unit of work, and every operation touches O(1) pages.
47pub const PAGE_BYTES: usize = 4096;
48
49/// Sentinel for "no page": an empty chain head, the end of a chain, or an
50/// empty free-list.
51const NONE: u32 = u32::MAX;
52
53/// Fibonacci hashing multiplier (2^64 / phi), used by [`ShardMode::Uniform`].
54const FIB: u64 = 0x9E37_79B9_7F4A_7C15;
55
56/// Fixed prefix of the serialized metadata section (see
57/// [`Arena::dump_meta`]).
58const IMAGE_HEADER: usize = 24;
59
60/// Serialized width of one page index entry (`heads`, `next` — little-endian
61/// `u32`), used to size and offset those metadata arrays.
62const PAGE_IDX_BYTES: usize = core::mem::size_of::<u32>();
63
64/// Serialized width of one page's slot count (`counts` — little-endian `u16`).
65const COUNT_BYTES: usize = core::mem::size_of::<u16>();
66
67/// Serialized bytes of per-page metadata: one `next` index plus one `counts`
68/// entry.
69const PAGE_META_BYTES: usize = PAGE_IDX_BYTES + COUNT_BYTES;
70
71/// `true` when the `counters` feature is enabled; lets hot loops keep their
72/// counting code in one expression that constant-folds away otherwise.
73const COUNT: bool = cfg!(feature = "counters");
74
75/// Adds `$n` to counter field `$field` when the `counters` feature is on;
76/// expands to nothing otherwise (zero cost, and the counter expression is
77/// not even evaluated).
78macro_rules! bump {
79    ($self:ident, $field:ident, $n:expr) => {
80        #[cfg(feature = "counters")]
81        {
82            let mut c = $self.counters.get();
83            c.$field += $n as u64;
84            $self.counters.set(c);
85        }
86    };
87}
88
89/// How keys are mapped to shards.
90#[derive(Clone, Copy, Debug, PartialEq, Eq)]
91#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
92pub enum ShardMode {
93    /// Shard = Fibonacci hash of the key's leading bytes. Spreads any key
94    /// distribution (including sequential ids) evenly across shards. Global
95    /// iteration order is *not* the key order. Use for lookup tables.
96    Uniform,
97    /// Shard = top bits of the key's leading bytes. Shard index order equals
98    /// key order, so [`Arena::iter`] yields globally ascending keys and
99    /// [`Arena::range`] scans work. Keys arriving in a narrow value range
100    /// will concentrate in few shards — their chains simply grow longer;
101    /// that is the trade-off for ordering.
102    Ordered,
103}
104
105/// Arena configuration.
106///
107/// Stored verbatim in snapshots later, so everything that
108/// affects byte interpretation lives here.
109#[derive(Clone, Copy, Debug, PartialEq, Eq)]
110#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
111pub struct ArenaCfg {
112    /// Number of shards; must be a non-zero power of two.
113    pub shards: usize,
114    /// Hard ceiling for the byte pool; exceeding it fails with
115    /// [`Error::CapacityExceeded`] instead of growing.
116    pub max_bytes: usize,
117    /// Key-to-shard mapping mode.
118    pub mode: ShardMode,
119}
120
121impl ArenaCfg {
122    /// Creates a config with the given shard count and mode, and no byte
123    /// limit (`max_bytes = usize::MAX`).
124    pub const fn new(shards: usize, mode: ShardMode) -> Self {
125        Self {
126            shards,
127            max_bytes: usize::MAX,
128            mode,
129        }
130    }
131
132    /// Returns the config with `max_bytes` replaced.
133    pub const fn with_max_bytes(mut self, max_bytes: usize) -> Self {
134        self.max_bytes = max_bytes;
135        self
136    }
137}
138
139impl Default for ArenaCfg {
140    /// 1024 shards, [`ShardMode::Uniform`], unlimited bytes.
141    fn default() -> Self {
142        Self::new(1024, ShardMode::Uniform)
143    }
144}
145
146/// Deterministic work counters (feature `counters`).
147///
148/// These are the basis of the project's CI performance gates: unlike
149/// wall-clock time they are identical on every machine, so a complexity
150/// regression fails the same way everywhere.
151#[cfg(feature = "counters")]
152#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
153pub struct Counters {
154    /// Key comparisons performed by binary searches.
155    pub cmp_ops: u64,
156    /// Bytes moved by insert/remove shifts and page splits.
157    pub bytes_shifted: u64,
158    /// Pages taken from the pool or the free-list.
159    pub pages_allocated: u64,
160    /// Steps taken walking page chains (each step peeks one page's first
161    /// key — one cache line).
162    pub chain_steps: u64,
163    /// Page splits performed by inserts into full pages.
164    pub splits: u64,
165}
166
167/// A sharded collection of fixed-size records in one contiguous byte pool,
168/// sorted by key prefix within each shard.
169///
170/// See the [crate-level documentation](crate) for the philosophy and the
171/// module documentation for the chain/split mechanics. Highlights:
172///
173/// - all state is `pool` + three small arrays — persisting the arena is a
174///   `memcpy` of defined bytes;
175/// - every operation touches O(1) pages: a short chain walk (one first-key
176///   peek per step), then binary search and shifts inside one 4 KiB page;
177/// - capacity is bounded only by [`ArenaCfg::max_bytes`] — full pages split,
178///   emptied pages are recycled through a free-list;
179/// - `Arena` deliberately implements neither `Clone` nor `PartialEq`: pages
180///   are allocated **uninitialized** (the measured wasm optimization), and a
181///   byte-wise clone/compare would read the uninitialized tails.
182pub struct Arena<'a, T: Slot> {
183    /// Page pool; length is always a multiple of [`PAGE_BYTES`]. Bytes of a
184    /// page beyond `counts[page] * T::SIZE` are uninitialized and must
185    /// never be read.
186    ///
187    /// A [`Paged`] backing so the arena can either own the pool (`new`/`load` —
188    /// the writable default, and the only shape on wasm32) or borrow a base
189    /// from a longer-lived buffer such as a memory-mapped snapshot
190    /// (`load_borrowed`/`load_overlay` —). The overlay open borrows
191    /// the base read-only and copies just the touched pages up on first write
192    /// (per-page copy-on-write), so a mapped database is mutated without
193    /// cloning it. The backing is `alloc`-only (no `std`/`Arc`), so the crate
194    /// stays `no_std`, and the lifetime `'a` lets the compiler prove the arena
195    /// never outlives its base. The uninitialized-tail invariant is untouched:
196    /// base pages are a fully-initialized dumped image, overlay copies are
197    /// taken from them, and freshly grown pages keep the same
198    /// read-only-up-to-`count` contract as before.
199    pool: Paged<'a, PAGE_BYTES>,
200    /// Shard -> first page of its chain (`NONE` = empty shard).
201    heads: Vec<u32>,
202    /// Page -> successor: the next page in a shard chain, or the next free
203    /// page when the page sits in the free-list.
204    next: Vec<u32>,
205    /// Page -> number of occupied slots.
206    counts: Vec<u16>,
207    /// Head of the free-list of recycled pages (linked through `next`).
208    free_head: u32,
209    /// Total records across all shards.
210    total: usize,
211    /// Reusable serialization buffer for [`Arena::insert`] (`T::SIZE` bytes,
212    /// allocated once) — inserts do not allocate after the first call except
213    /// when growing the pool itself.
214    scratch: Vec<u8>,
215    /// Reusable buffer for the cross-page slot moves in [`Arena::split`]. A
216    /// split reads a run of slots out of one page and writes it into another;
217    /// under the overlay backing those may be distinct owned pages, so the
218    /// bytes route through this buffer instead of a single in-pool
219    /// `copy_within`. Grown once; splits are amortized and off the recall path.
220    split_buf: Vec<u8>,
221    cfg: ArenaCfg,
222    #[cfg(feature = "counters")]
223    counters: Cell<Counters>,
224    /// `fn() -> T` keeps the arena `Send`/`Sync`-neutral and avoids bounding
225    /// auto-traits on `T` itself.
226    _marker: PhantomData<fn() -> T>,
227}
228
229/// Where a key's record lives (or would live).
230struct Target {
231    /// Page preceding `page` in its chain, `NONE` when `page` is the head.
232    prev: u32,
233    /// The chain page whose key range covers the key.
234    page: u32,
235}
236
237impl<'a, T: Slot> Arena<'a, T> {
238    /// Creates an empty arena.
239    ///
240    /// # Errors
241    ///
242    /// - [`Error::BadSlot`] unless `1 <= T::KEY_LEN <= T::SIZE <= PAGE_BYTES`;
243    /// - [`Error::BadShardCount`] unless `cfg.shards` is a non-zero power of
244    ///   two.
245    pub fn new(cfg: ArenaCfg) -> Result<Self, Error> {
246        if T::SIZE == 0 || T::SIZE > PAGE_BYTES || T::KEY_LEN == 0 || T::KEY_LEN > T::SIZE {
247            return Err(Error::BadSlot {
248                size: T::SIZE,
249                key_len: T::KEY_LEN,
250            });
251        }
252        if cfg.shards == 0 || !cfg.shards.is_power_of_two() {
253            return Err(Error::BadShardCount { got: cfg.shards });
254        }
255        Ok(Self {
256            pool: Paged::owned_empty(),
257            heads: alloc::vec![NONE; cfg.shards],
258            next: Vec::new(),
259            counts: Vec::new(),
260            free_head: NONE,
261            total: 0,
262            scratch: Vec::new(),
263            split_buf: Vec::new(),
264            cfg,
265            #[cfg(feature = "counters")]
266            counters: Cell::new(Counters::default()),
267            _marker: PhantomData,
268        })
269    }
270
271    /// Number of slots one page holds for this slot size (at least 1).
272    pub const fn slots_per_page() -> usize {
273        PAGE_BYTES / T::SIZE
274    }
275
276    /// Total number of records.
277    pub fn len(&self) -> usize {
278        self.total
279    }
280
281    /// `true` when the arena holds no records.
282    pub fn is_empty(&self) -> bool {
283        self.total == 0
284    }
285
286    /// Bytes currently allocated in the page pool (including recycled
287    /// free-list pages — the pool never shrinks; emptied pages are reused).
288    pub fn pool_bytes(&self) -> usize {
289        self.pool.len()
290    }
291
292    /// The configuration this arena was created with.
293    pub fn cfg(&self) -> &ArenaCfg {
294        &self.cfg
295    }
296
297    /// Inserts a record, keeping its shard's chain sorted by key prefix.
298    ///
299    /// Returns `Ok(false)` if a record with the same key prefix already
300    /// exists (the arena is a map keyed by prefix; existing payload is left
301    /// untouched — use [`Arena::payload_mut`] to update in place). A full
302    /// target page splits in half — inserts never fail on page capacity.
303    ///
304    /// # Errors
305    ///
306    /// [`Error::CapacityExceeded`] if a needed page allocation would cross
307    /// [`ArenaCfg::max_bytes`].
308    pub fn insert(&mut self, value: &T) -> Result<bool, Error> {
309        // Take the scratch buffer out of `self` to sidestep aliasing between
310        // `&self.scratch` and `&mut self.pool` below.
311        let mut buf = core::mem::take(&mut self.scratch);
312        buf.resize(T::SIZE, 0);
313        value.write(&mut buf);
314        let result = self.insert_bytes(&buf);
315        self.scratch = buf;
316        result
317    }
318
319    /// Insert path operating on the serialized slot.
320    fn insert_bytes(&mut self, slot: &[u8]) -> Result<bool, Error> {
321        let key = &slot[..T::KEY_LEN];
322        let shard = self.shard_of(key);
323
324        let target = match self.find_page(shard, key) {
325            Some(t) => t,
326            None => {
327                // Empty shard: allocate its first chain page.
328                let page = self.alloc_page()?;
329                self.heads[shard] = page;
330                Target { prev: NONE, page }
331            }
332        };
333
334        let mut page = target.page;
335        let mut count = self.counts[page as usize] as usize;
336        let mut pos = {
337            let mut cmps = 0u64;
338            let found = self.search_in(page, count, key, &mut cmps);
339            bump!(self, cmp_ops, cmps);
340            match found {
341                Ok(_) => return Ok(false),
342                Err(pos) => pos,
343            }
344        };
345
346        if count == Self::slots_per_page() {
347            // Split the full page; afterwards (page, pos, count) address the
348            // half that must receive the new record.
349            (page, pos, count) = self.split(page, pos)?;
350        }
351
352        let slot_start = pos * T::SIZE;
353        let used_end = count * T::SIZE;
354        let shifted = used_end - slot_start;
355        let bytes = self.pool.page_mut(page);
356        if shifted > 0 {
357            // Shift the sorted tail right by one slot; stays within the page.
358            bytes.copy_within(slot_start..used_end, slot_start + T::SIZE);
359        }
360        bytes[slot_start..slot_start + T::SIZE].copy_from_slice(slot);
361        bump!(self, bytes_shifted, shifted);
362
363        self.counts[page as usize] += 1;
364        self.total += 1;
365        Ok(true)
366    }
367
368    /// Splits full `page`, given the insert position `pos` inside it.
369    /// Returns `(target_page, target_pos, target_count)` for the record that
370    /// triggered the split.
371    fn split(&mut self, page: u32, pos: usize) -> Result<(u32, usize, usize), Error> {
372        let spp = Self::slots_per_page();
373        let fresh = self.alloc_page()?;
374        bump!(self, splits, 1);
375
376        // Link the fresh page right after the split one.
377        self.next[fresh as usize] = self.next[page as usize];
378        self.next[page as usize] = fresh;
379
380        if spp == 1 {
381            // Degenerate single-slot pages (T::SIZE > PAGE_BYTES / 2): the
382            // "upper half" is the whole record when the new key precedes it,
383            // otherwise the fresh page simply receives the new record.
384            return Ok(if pos == 0 {
385                self.move_slots(page, 0, fresh, T::SIZE);
386                bump!(self, bytes_shifted, T::SIZE);
387                self.counts[page as usize] = 0;
388                self.counts[fresh as usize] = 1;
389                (page, 0, 0)
390            } else {
391                (fresh, 0, 0)
392            });
393        }
394
395        // Move the upper half [half..spp) into the fresh page.
396        let half = spp / 2;
397        let moved = spp - half;
398        self.move_slots(page, half * T::SIZE, fresh, moved * T::SIZE);
399        bump!(self, bytes_shifted, moved * T::SIZE);
400        self.counts[page as usize] = half as u16;
401        self.counts[fresh as usize] = moved as u16;
402
403        // `pos` was computed against the full page; place the new record in
404        // whichever half now owns that position (pos == half belongs at the
405        // end of the lower page: the key sorts before the fresh page's first).
406        Ok(if pos <= half {
407            (page, pos, half)
408        } else {
409            (fresh, pos - half, moved)
410        })
411    }
412
413    /// Copies `len` bytes from `src_page` (starting at in-page offset
414    /// `src_off`) to the start of `dst_page`. The two pages may be distinct
415    /// owned pages under the overlay backing, so the bytes route through the
416    /// reusable [`Arena::split_buf`] rather than a single in-pool
417    /// `copy_within` — one immutable read of the source, then one write of the
418    /// destination (which copies-up if it is still a borrowed base page).
419    fn move_slots(&mut self, src_page: u32, src_off: usize, dst_page: u32, len: usize) {
420        let mut buf = core::mem::take(&mut self.split_buf);
421        buf.clear();
422        buf.extend_from_slice(&self.pool.page(src_page)[src_off..src_off + len]);
423        self.pool.page_mut(dst_page)[..len].copy_from_slice(&buf);
424        self.split_buf = buf;
425    }
426
427    /// Returns the record with the given key prefix, if present.
428    ///
429    /// # Panics
430    ///
431    /// Panics if `key.len() != T::KEY_LEN` (a caller bug, not a data
432    /// condition — mirrors slice indexing).
433    pub fn get(&self, key: &[u8]) -> Option<T> {
434        self.locate(key).map(|off| {
435            let (page, rel) = (off / PAGE_BYTES, off % PAGE_BYTES);
436            T::read(&self.pool.page(page as u32)[rel..rel + T::SIZE])
437        })
438    }
439
440    /// `true` if a record with the given key prefix exists.
441    ///
442    /// # Panics
443    ///
444    /// Panics if `key.len() != T::KEY_LEN`.
445    pub fn contains(&self, key: &[u8]) -> bool {
446        self.locate(key).is_some()
447    }
448
449    /// Returns the raw bytes of the record with the given key prefix.
450    ///
451    /// # Panics
452    ///
453    /// Panics if `key.len() != T::KEY_LEN`.
454    pub fn get_slot(&self, key: &[u8]) -> Option<&[u8]> {
455        self.locate(key).map(|off| {
456            let (page, rel) = (off / PAGE_BYTES, off % PAGE_BYTES);
457            &self.pool.page(page as u32)[rel..rel + T::SIZE]
458        })
459    }
460
461    /// Returns a mutable view of the record's **payload** (the bytes after
462    /// the key prefix), for in-place updates.
463    ///
464    /// The key prefix itself is not reachable through this method — sorted
465    /// order cannot be corrupted by construction, which is why this is safe
466    /// to expose at all.
467    ///
468    /// # Panics
469    ///
470    /// Panics if `key.len() != T::KEY_LEN`.
471    pub fn payload_mut(&mut self, key: &[u8]) -> Option<&mut [u8]> {
472        let off = self.locate(key)?;
473        let (page, rel) = (off / PAGE_BYTES, off % PAGE_BYTES);
474        Some(&mut self.pool.page_mut(page as u32)[rel + T::KEY_LEN..rel + T::SIZE])
475    }
476
477    /// Removes the record with the given key prefix. Returns `true` if it
478    /// existed. A page emptied by the removal is unlinked from its chain and
479    /// recycled through the free-list.
480    ///
481    /// # Panics
482    ///
483    /// Panics if `key.len() != T::KEY_LEN`.
484    pub fn remove(&mut self, key: &[u8]) -> bool {
485        assert_eq!(key.len(), T::KEY_LEN, "key length must equal Slot::KEY_LEN");
486        let shard = self.shard_of(key);
487        let Some(Target { prev, page }) = self.find_page(shard, key) else {
488            return false;
489        };
490        let count = self.counts[page as usize] as usize;
491        let pos = {
492            let mut cmps = 0u64;
493            let found = self.search_in(page, count, key, &mut cmps);
494            bump!(self, cmp_ops, cmps);
495            match found {
496                Ok(pos) => pos,
497                Err(_) => return false,
498            }
499        };
500
501        let slot_start = pos * T::SIZE;
502        let used_end = count * T::SIZE;
503        let tail = used_end - (slot_start + T::SIZE);
504        if tail > 0 {
505            // Shift the tail left over the removed slot; stays within the page.
506            self.pool
507                .page_mut(page)
508                .copy_within(slot_start + T::SIZE..used_end, slot_start);
509        }
510        bump!(self, bytes_shifted, tail);
511        self.counts[page as usize] -= 1;
512        self.total -= 1;
513
514        if self.counts[page as usize] == 0 {
515            // Unlink the emptied page and recycle it.
516            let successor = self.next[page as usize];
517            if prev == NONE {
518                self.heads[shard] = successor;
519            } else {
520                self.next[prev as usize] = successor;
521            }
522            self.next[page as usize] = self.free_head;
523            self.free_head = page;
524        }
525        true
526    }
527
528    /// Iterates all records shard by shard, following each shard's chain.
529    ///
530    /// In [`ShardMode::Ordered`] the shard index order equals key order and
531    /// chains are range-partitioned, so this yields **globally ascending
532    /// keys**. In [`ShardMode::Uniform`] the global order is unspecified
533    /// (each shard is still internally sorted).
534    pub fn iter(&self) -> Iter<'_, T> {
535        Iter {
536            arena: self,
537            shard: 0,
538            page: NONE,
539            idx: 0,
540            remaining: self.total,
541        }
542    }
543
544    /// Iterates records whose key prefix lies in `[from, to)` — `from`
545    /// inclusive, `to` exclusive — in ascending key order.
546    ///
547    /// Only meaningful when shard order equals key order, hence restricted
548    /// to [`ShardMode::Ordered`].
549    ///
550    /// # Panics
551    ///
552    /// Panics if the arena is in [`ShardMode::Uniform`], or if either bound's
553    /// length differs from `T::KEY_LEN`.
554    pub fn range<'s>(&'s self, from: &[u8], to: &'s [u8]) -> Range<'s, T> {
555        assert_eq!(
556            self.cfg.mode,
557            ShardMode::Ordered,
558            "range scans require ShardMode::Ordered"
559        );
560        assert_eq!(
561            from.len(),
562            T::KEY_LEN,
563            "key length must equal Slot::KEY_LEN"
564        );
565        assert_eq!(to.len(), T::KEY_LEN, "key length must equal Slot::KEY_LEN");
566
567        // Position on the first record with key >= `from`.
568        let shard = self.shard_of(from);
569        let mut page = NONE;
570        let mut idx = 0usize;
571        if let Some(t) = self.find_page(shard, from) {
572            let count = self.counts[t.page as usize] as usize;
573            let mut cmps = 0u64;
574            let pos = match self.search_in(t.page, count, from, &mut cmps) {
575                Ok(p) | Err(p) => p,
576            };
577            bump!(self, cmp_ops, cmps);
578            if pos < count {
579                page = t.page;
580                idx = pos;
581            } else {
582                // Past the last record of the covering page: continue with
583                // the next page in this chain, or fall through to the next
584                // shard.
585                page = self.next[t.page as usize];
586            }
587        }
588        Range {
589            arena: self,
590            // The shard the iterator moves to once the current chain (if
591            // any) is exhausted.
592            shard: shard + 1,
593            page,
594            idx,
595            to,
596        }
597    }
598
599    /// A snapshot of the work counters (feature `counters`).
600    #[cfg(feature = "counters")]
601    pub fn counters(&self) -> Counters {
602        self.counters.get()
603    }
604
605    /// Resets all work counters to zero (feature `counters`).
606    #[cfg(feature = "counters")]
607    pub fn reset_counters(&self) {
608        self.counters.set(Counters::default());
609    }
610
611    /// Appends the arena's metadata section to `out`.
612    ///
613    /// Layout (all little-endian): `[shards u32][pages u32][free_head u32]
614    /// [total u64][mode u8][reserved 3]`, then `heads` (`shards × u32`),
615    /// `next` (`pages × u32`), `counts` (`pages × u16`). Together with
616    /// [`Arena::dump_pool`] this is the complete state; both dumps are
617    /// canonical — dump → [`Arena::load`] → dump reproduces identical
618    /// bytes.
619    pub fn dump_meta(&self, out: &mut Vec<u8>) {
620        let pages = self.pool.len() / PAGE_BYTES;
621        debug_assert!(self.cfg.shards <= u32::MAX as usize && pages < u32::MAX as usize);
622        out.reserve(IMAGE_HEADER + self.heads.len() * PAGE_IDX_BYTES + pages * PAGE_META_BYTES);
623        out.extend_from_slice(&(self.cfg.shards as u32).to_le_bytes());
624        out.extend_from_slice(&(pages as u32).to_le_bytes());
625        out.extend_from_slice(&self.free_head.to_le_bytes());
626        out.extend_from_slice(&(self.total as u64).to_le_bytes());
627        out.push(match self.cfg.mode {
628            ShardMode::Uniform => 0,
629            ShardMode::Ordered => 1,
630        });
631        out.extend_from_slice(&[0u8; 3]);
632        for &head in &self.heads {
633            out.extend_from_slice(&head.to_le_bytes());
634        }
635        for &next in &self.next {
636            out.extend_from_slice(&next.to_le_bytes());
637        }
638        for &count in &self.counts {
639            out.extend_from_slice(&count.to_le_bytes());
640        }
641    }
642
643    /// Appends the arena's pool section to `out`.
644    ///
645    /// Each page contributes its initialized prefix (`counts[page] ×
646    /// SIZE` bytes) followed by zero padding to [`PAGE_BYTES`]: the
647    /// uninitialized page tails (see [`Arena::insert`] internals) are
648    /// never read, and zero-filling them makes the image canonical and
649    /// leak-free.
650    pub fn dump_pool(&self, out: &mut Vec<u8>) {
651        out.reserve(self.pool.len());
652        for (page, &count) in self.counts.iter().enumerate() {
653            let used = count as usize * T::SIZE;
654            out.extend_from_slice(&self.pool.page(page as u32)[..used]);
655            out.resize(out.len() + (PAGE_BYTES - used), 0);
656        }
657    }
658
659    /// Rebuilds an arena from its two dumped sections.
660    ///
661    /// The input is **untrusted**: every metadata invariant is checked
662    /// before adoption and any inconsistency returns
663    /// [`Error::Corrupt`] — this method never panics on arbitrary bytes.
664    /// Cost is O(pages) on top of the pool copy:
665    ///
666    /// - exact section lengths; `cfg` agreement (shards, mode);
667    /// - per-page slot counts within `PAGE_BYTES / SIZE`;
668    /// - every page reachable exactly once — shard chains and the
669    ///   free-list are walked with a visited bitmap (no cycles, no shared
670    ///   or orphan pages); chain pages are non-empty, free pages empty;
671    /// - chain pages ascend by key range and sit in the shard their first
672    ///   key maps to; the record total matches the page counts.
673    ///
674    /// Slot *content* is not validated beyond the per-page first/last key
675    /// reads — record payloads are semantically validated lazily by the
676    /// owning engine.
677    ///
678    /// # Errors
679    ///
680    /// [`Error::BadSlot`] / [`Error::BadShardCount`] for an invalid `cfg`
681    /// (same gates as [`Arena::new`]); [`Error::Corrupt`] for any image
682    /// inconsistency.
683    pub fn load(cfg: ArenaCfg, meta: &[u8], pool: &[u8]) -> Result<Self, Error> {
684        Self::load_impl(cfg, meta, pool, Paged::owned_from(pool.to_vec()))
685    }
686
687    /// Rebuilds an arena that **borrows** its page pool from a longer-lived
688    /// buffer (a memory-mapped snapshot) instead of copying it.
689    /// Validation is identical to [`Arena::load`] — it reads each page's
690    /// first and last keys, so it touches one OS page per arena page (the
691    /// key-order check needs them), but no pool byte is copied.
692    ///
693    /// # Errors
694    ///
695    /// Same as [`Arena::load`].
696    pub fn load_borrowed(cfg: ArenaCfg, meta: &[u8], pool: &'a [u8]) -> Result<Self, Error> {
697        Self::load_impl(cfg, meta, pool, Paged::borrowed(pool))
698    }
699
700    /// Opens an arena over a borrowed base for the **overlay** write path: the
701    /// base pages are mapped read-only, and the first write to any base page
702    /// copies just that page into owned storage (per-page copy-on-write, see
703    /// [`Paged`](crate::paged)), while pages grown after open live in an owned
704    /// tail. Unlike [`Arena::load_borrowed`] the returned arena is fully
705    /// mutable — inserts, removals and splits work — yet the borrowed base is
706    /// never cloned as a whole and never mutated, so a memory-mapped database
707    /// can be written to while resident only in the pages it actually touches.
708    /// Validation is identical to [`Arena::load`].
709    ///
710    /// # Errors
711    ///
712    /// Same as [`Arena::load`].
713    pub fn load_overlay(cfg: ArenaCfg, meta: &[u8], pool: &'a [u8]) -> Result<Self, Error> {
714        Self::load_impl(cfg, meta, pool, Paged::borrowed(pool))
715    }
716
717    /// Shared body of [`Arena::load`] / [`Arena::load_borrowed`] /
718    /// [`Arena::load_overlay`]: validates the untrusted `meta`/`pool` image and
719    /// adopts `backing` as the pool.
720    fn load_impl(
721        cfg: ArenaCfg,
722        meta: &[u8],
723        pool: &[u8],
724        backing: Paged<'a, PAGE_BYTES>,
725    ) -> Result<Self, Error> {
726        let mut arena = Self::new(cfg)?;
727        if meta.len() < IMAGE_HEADER {
728            return Err(Error::Corrupt("arena meta shorter than its header"));
729        }
730        let shards = u32::from_le_bytes(meta[0..4].try_into().unwrap()) as usize;
731        let pages = u32::from_le_bytes(meta[4..8].try_into().unwrap()) as usize;
732        let free_head = u32::from_le_bytes(meta[8..12].try_into().unwrap());
733        let total = u64::from_le_bytes(meta[12..20].try_into().unwrap());
734        let mode = meta[20];
735        if meta[21..24] != [0u8; 3] {
736            return Err(Error::Corrupt("arena meta reserved bytes must be zero"));
737        }
738        if shards != cfg.shards {
739            return Err(Error::Corrupt("arena meta shard count disagrees with cfg"));
740        }
741        let want_mode = match cfg.mode {
742            ShardMode::Uniform => 0u8,
743            ShardMode::Ordered => 1,
744        };
745        if mode != want_mode {
746            return Err(Error::Corrupt("arena meta shard mode disagrees with cfg"));
747        }
748        if pages as u64 >= u64::from(NONE) {
749            return Err(Error::Corrupt("arena page count overflows the index space"));
750        }
751        let want_meta = IMAGE_HEADER as u64
752            + shards as u64 * PAGE_IDX_BYTES as u64
753            + pages as u64 * PAGE_META_BYTES as u64;
754        if meta.len() as u64 != want_meta {
755            return Err(Error::Corrupt("arena meta length mismatch"));
756        }
757        if pool.len() as u64 != pages as u64 * PAGE_BYTES as u64 {
758            return Err(Error::Corrupt("arena pool length mismatch"));
759        }
760
761        let mut heads = Vec::with_capacity(shards);
762        for i in 0..shards {
763            let at = IMAGE_HEADER + i * PAGE_IDX_BYTES;
764            heads.push(u32::from_le_bytes(
765                meta[at..at + PAGE_IDX_BYTES].try_into().unwrap(),
766            ));
767        }
768        let next_base = IMAGE_HEADER + shards * PAGE_IDX_BYTES;
769        let mut next = Vec::with_capacity(pages);
770        for i in 0..pages {
771            let at = next_base + i * PAGE_IDX_BYTES;
772            next.push(u32::from_le_bytes(
773                meta[at..at + PAGE_IDX_BYTES].try_into().unwrap(),
774            ));
775        }
776        let counts_base = next_base + pages * PAGE_IDX_BYTES;
777        let mut counts = Vec::with_capacity(pages);
778        for i in 0..pages {
779            let at = counts_base + i * COUNT_BYTES;
780            counts.push(u16::from_le_bytes(
781                meta[at..at + COUNT_BYTES].try_into().unwrap(),
782            ));
783        }
784
785        let spp = Self::slots_per_page();
786        if counts.iter().any(|&c| c as usize > spp) {
787            return Err(Error::Corrupt("arena page count exceeds slots per page"));
788        }
789
790        // Reachability walk: every page must appear exactly once across all
791        // shard chains and the free-list; the bitmap catches cycles, shared
792        // pages and (after the walks) orphans.
793        let mut seen = alloc::vec![false; pages];
794        let mut live = 0u64;
795        for (shard, &head) in heads.iter().enumerate() {
796            let mut page = head;
797            let mut prev: Option<u32> = None;
798            while page != NONE {
799                let p = page as usize;
800                if p >= pages {
801                    return Err(Error::Corrupt("arena chain page out of bounds"));
802                }
803                if core::mem::replace(&mut seen[p], true) {
804                    return Err(Error::Corrupt("arena page linked more than once"));
805                }
806                let count = counts[p] as usize;
807                if count == 0 {
808                    return Err(Error::Corrupt("arena chain contains an empty page"));
809                }
810                live += count as u64;
811                let first = &pool[p * PAGE_BYTES..p * PAGE_BYTES + T::KEY_LEN];
812                if arena.shard_of(first) != shard {
813                    return Err(Error::Corrupt("arena page sits in the wrong shard"));
814                }
815                if let Some(pr) = prev {
816                    let last_at =
817                        pr as usize * PAGE_BYTES + (counts[pr as usize] as usize - 1) * T::SIZE;
818                    if pool[last_at..last_at + T::KEY_LEN] >= *first {
819                        return Err(Error::Corrupt("arena chain pages out of key order"));
820                    }
821                }
822                prev = Some(page);
823                page = next[p];
824            }
825        }
826        let mut page = free_head;
827        while page != NONE {
828            let p = page as usize;
829            if p >= pages {
830                return Err(Error::Corrupt("arena free page out of bounds"));
831            }
832            if core::mem::replace(&mut seen[p], true) {
833                return Err(Error::Corrupt("arena page linked more than once"));
834            }
835            if counts[p] != 0 {
836                return Err(Error::Corrupt("arena free page has a nonzero count"));
837            }
838            page = next[p];
839        }
840        if seen.iter().any(|&s| !s) {
841            return Err(Error::Corrupt("arena has an orphan page"));
842        }
843        if live != total {
844            return Err(Error::Corrupt(
845                "arena record total disagrees with page counts",
846            ));
847        }
848
849        arena.pool = backing;
850        arena.heads = heads;
851        arena.next = next;
852        arena.counts = counts;
853        arena.free_head = free_head;
854        arena.total = total as usize;
855        Ok(arena)
856    }
857
858    /// Walks the shard's chain to the page whose key range covers `key`.
859    /// `None` when the shard is empty.
860    fn find_page(&self, shard: usize, key: &[u8]) -> Option<Target> {
861        let head = self.heads[shard];
862        if head == NONE {
863            return None;
864        }
865        let mut prev = NONE;
866        let mut page = head;
867        let mut steps = 0u64;
868        loop {
869            let nxt = self.next[page as usize];
870            // Advance while the next page's range can still contain the key
871            // (its first key is <= key). Peeking a first key touches one
872            // cache line.
873            if nxt == NONE || self.first_key(nxt) > key {
874                break;
875            }
876            steps += COUNT as u64;
877            prev = page;
878            page = nxt;
879        }
880        bump!(self, chain_steps, steps);
881        let _ = steps; // read only by the counters feature
882        Some(Target { prev, page })
883    }
884
885    /// First key of a page. Caller guarantees the page is non-empty (every
886    /// page in a chain holds at least one record — emptied pages are
887    /// unlinked immediately).
888    fn first_key(&self, page: u32) -> &[u8] {
889        &self.pool.page(page)[..T::KEY_LEN]
890    }
891
892    /// Binary search inside a page; wraps the free-function search with the
893    /// page slice resolution.
894    fn search_in(
895        &self,
896        page: u32,
897        count: usize,
898        key: &[u8],
899        cmps: &mut u64,
900    ) -> Result<usize, usize> {
901        search::<T>(self.pool.page(page), count, key, cmps)
902    }
903
904    /// Byte offset of the slot with the given key, if present.
905    fn locate(&self, key: &[u8]) -> Option<usize> {
906        assert_eq!(key.len(), T::KEY_LEN, "key length must equal Slot::KEY_LEN");
907        let shard = self.shard_of(key);
908        let Target { page, .. } = self.find_page(shard, key)?;
909        let count = self.counts[page as usize] as usize;
910        let mut cmps = 0u64;
911        let found = self.search_in(page, count, key, &mut cmps).ok();
912        bump!(self, cmp_ops, cmps);
913        found.map(|pos| page as usize * PAGE_BYTES + pos * T::SIZE)
914    }
915
916    /// Maps a key to its shard. Only the first 8 key bytes participate;
917    /// longer keys sharing an 8-byte prefix land in the same shard (harmless
918    /// for `Ordered`: order across shards is still by prefix).
919    fn shard_of(&self, key: &[u8]) -> usize {
920        let mut pad = [0u8; 8];
921        let n = key.len().min(8);
922        pad[..n].copy_from_slice(&key[..n]);
923        let v = u64::from_be_bytes(pad);
924        let bits = self.cfg.shards.trailing_zeros();
925        if bits == 0 {
926            return 0; // single shard; also avoids the undefined `v >> 64`
927        }
928        let h = match self.cfg.mode {
929            ShardMode::Ordered => v,
930            ShardMode::Uniform => v.wrapping_mul(FIB),
931        };
932        (h >> (64 - bits)) as usize
933    }
934
935    /// Takes a page from the free-list, or grows the pool by one page.
936    fn alloc_page(&mut self) -> Result<u32, Error> {
937        if self.free_head != NONE {
938            let page = self.free_head;
939            self.free_head = self.next[page as usize];
940            self.next[page as usize] = NONE;
941            self.counts[page as usize] = 0;
942            bump!(self, pages_allocated, 1);
943            return Ok(page);
944        }
945
946        let old_len = self.pool.len();
947        let new_len = old_len + PAGE_BYTES;
948        if new_len > self.cfg.max_bytes {
949            return Err(Error::CapacityExceeded {
950                max_bytes: self.cfg.max_bytes,
951            });
952        }
953        let page = (old_len / PAGE_BYTES) as u32;
954        // Grow the owned tail by one page: the whole vector when owned, the
955        // overlay's grown tail when borrowing (the borrowed base is never
956        // resized). `grown_tail_mut` hands back that owned `Vec` directly — no
957        // clone of the base — and the new page's global index is `old_len /
958        // PAGE_BYTES` regardless of which tail holds it.
959        let tail = self.pool.grown_tail_mut();
960        let tail_len = tail.len() + PAGE_BYTES;
961        tail.reserve(PAGE_BYTES);
962        // SAFETY: the new page is left uninitialized on purpose — this is the
963        // one measured unsafe of the crate. Zeroing fresh pages was benched
964        // at 12x slower on the wasm allocation path (wasmtime, 32k pages:
965        // 3889 us zeroed vs 316 us uninit; native: noise) — and wasm is this
966        // structure's primary environment. The invariant making it sound:
967        // *no byte of a page beyond `counts[page] * T::SIZE` is ever read*.
968        // Every read (search, get, iter, range, first_key, shifts) is
969        // bounded by the page count, and a slot's bytes are fully written
970        // before the count is incremented. Consequently `Arena` exposes no
971        // whole-pool reads (no `Clone`/`PartialEq`/`as_bytes`); the snapshot
972        // writer emits only the initialized prefixes of pages.
973        // Routing the same reserve+set_len at the grown tail (not the base)
974        // keeps this the crate's sole `unsafe` — the overlay adds none.
975        #[allow(clippy::uninit_vec)]
976        unsafe {
977            tail.set_len(tail_len);
978        }
979        self.next.push(NONE);
980        self.counts.push(0);
981        bump!(self, pages_allocated, 1);
982        Ok(page)
983    }
984}
985
986/// Binary search over a page's occupied slots, comparing key prefixes.
987///
988/// Free function (not a method) so the borrow of the page slice stays
989/// independent from `&mut self` at call sites. Counting compiles away when
990/// the `counters` feature is off (`COUNT` is a const `false`).
991fn search<T: Slot>(page: &[u8], count: usize, key: &[u8], cmps: &mut u64) -> Result<usize, usize> {
992    let mut lo = 0usize;
993    let mut hi = count;
994    while lo < hi {
995        let mid = lo + (hi - lo) / 2;
996        let off = mid * T::SIZE;
997        // Branchless: adds 0 when the `counters` feature is off, which the
998        // compiler folds away entirely.
999        *cmps += COUNT as u64;
1000        match page[off..off + T::KEY_LEN].cmp(key) {
1001            core::cmp::Ordering::Less => lo = mid + 1,
1002            core::cmp::Ordering::Greater => hi = mid,
1003            core::cmp::Ordering::Equal => return Ok(mid),
1004        }
1005    }
1006    Err(lo)
1007}
1008
1009impl<T: Slot> fmt::Debug for Arena<'_, T> {
1010    /// Summary only — dumping the pool would both flood output and read
1011    /// uninitialized page tails.
1012    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1013        f.debug_struct("Arena")
1014            .field("len", &self.total)
1015            .field("shards", &self.cfg.shards)
1016            .field("mode", &self.cfg.mode)
1017            .field("pages", &(self.pool.len() / PAGE_BYTES))
1018            .field("slot_size", &T::SIZE)
1019            .finish()
1020    }
1021}
1022
1023/// Iterator over all records of an [`Arena`]; see [`Arena::iter`] for
1024/// ordering guarantees.
1025pub struct Iter<'a, T: Slot> {
1026    arena: &'a Arena<'a, T>,
1027    shard: usize,
1028    /// Current chain page; `NONE` means "advance to the next shard's head".
1029    page: u32,
1030    idx: usize,
1031    remaining: usize,
1032}
1033
1034impl<T: Slot> Iterator for Iter<'_, T> {
1035    type Item = T;
1036
1037    fn next(&mut self) -> Option<T> {
1038        loop {
1039            if self.page == NONE {
1040                if self.shard >= self.arena.cfg.shards {
1041                    return None;
1042                }
1043                self.page = self.arena.heads[self.shard];
1044                self.idx = 0;
1045                self.shard += 1;
1046                continue;
1047            }
1048            if self.idx < self.arena.counts[self.page as usize] as usize {
1049                let rel = self.idx * T::SIZE;
1050                self.idx += 1;
1051                self.remaining -= 1;
1052                return Some(T::read(
1053                    &self.arena.pool.page(self.page)[rel..rel + T::SIZE],
1054                ));
1055            }
1056            self.page = self.arena.next[self.page as usize];
1057            self.idx = 0;
1058        }
1059    }
1060
1061    fn size_hint(&self) -> (usize, Option<usize>) {
1062        (self.remaining, Some(self.remaining))
1063    }
1064}
1065
1066impl<T: Slot> ExactSizeIterator for Iter<'_, T> {}
1067
1068impl<'a, T: Slot> IntoIterator for &'a Arena<'a, T> {
1069    type Item = T;
1070    type IntoIter = Iter<'a, T>;
1071
1072    fn into_iter(self) -> Self::IntoIter {
1073        self.iter()
1074    }
1075}
1076
1077/// Iterator over a key range of an [`Arena`]; see [`Arena::range`].
1078pub struct Range<'a, T: Slot> {
1079    arena: &'a Arena<'a, T>,
1080    shard: usize,
1081    /// Current chain page; `NONE` means "advance to the next shard's head".
1082    page: u32,
1083    idx: usize,
1084    /// Exclusive upper bound on key prefixes.
1085    to: &'a [u8],
1086}
1087
1088impl<T: Slot> Iterator for Range<'_, T> {
1089    type Item = T;
1090
1091    fn next(&mut self) -> Option<T> {
1092        loop {
1093            if self.page == NONE {
1094                if self.shard >= self.arena.cfg.shards {
1095                    return None;
1096                }
1097                self.page = self.arena.heads[self.shard];
1098                self.idx = 0;
1099                self.shard += 1;
1100                continue;
1101            }
1102            if self.idx < self.arena.counts[self.page as usize] as usize {
1103                let rel = self.idx * T::SIZE;
1104                let page = self.arena.pool.page(self.page);
1105                if page[rel..rel + T::KEY_LEN] >= *self.to {
1106                    // Keys only grow from here on — the scan is complete.
1107                    return None;
1108                }
1109                self.idx += 1;
1110                return Some(T::read(&page[rel..rel + T::SIZE]));
1111            }
1112            self.page = self.arena.next[self.page as usize];
1113            self.idx = 0;
1114        }
1115    }
1116}