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}