Skip to main content

ic_sqlite_vfs/sqlite_vfs/
stable_blob.rs

1//! Logical `/main.db` access backed by segmented stable-memory page mapping.
2//!
3//! SQLite sees a contiguous file. Internally, the active superblock points to a
4//! root table. Each root entry points to a 256-page segment table.
5
6use crate::config::{SQLITE_PAGE_SIZE, STABLE_PAGE_SIZE, SUPERBLOCK_SIZE};
7use crate::sqlite_vfs::overlay::{self, Overlay};
8use crate::stable::memory::{self, ContextId, StableMemoryError};
9use crate::stable::meta::{
10    fnv1a64, Superblock, FLAG_CHECKSUM_REFRESHING, FLAG_CHECKSUM_STALE, FLAG_IMPORTING,
11    PAGE_MAP_LAYOUT_VERSION,
12};
13use std::cell::RefCell;
14use std::collections::BTreeMap;
15use std::mem::MaybeUninit;
16
17const CHECKSUM_CHUNK_LEN: u64 = 16 * 1024;
18const PAGE_TABLE_ENTRY_LEN: u64 = 8;
19const SEGMENT_PAGE_COUNT: u64 = 256;
20const SEGMENT_TABLE_BYTES: u64 = SEGMENT_PAGE_COUNT * PAGE_TABLE_ENTRY_LEN;
21const SINGLE_SEGMENT_PAGE_TABLE_BYTES: u64 = SEGMENT_TABLE_BYTES + PAGE_TABLE_ENTRY_LEN;
22const READ_SEGMENT_CACHE_CAPACITY: usize = 8;
23const FILE_PAGE_OFFSET_CACHE_CAPACITY: usize = 64;
24const FILE_PAGE_DATA_CACHE_CAPACITY: usize = 8;
25const COMPACT_MIN_ORPHAN_BYTES: u64 = 16 * 1024 * 1024;
26
27#[derive(Clone, Debug, Eq, PartialEq)]
28pub struct ChecksumRefresh {
29    pub complete: bool,
30    pub checksum: u64,
31    pub scanned_bytes: u64,
32    pub db_size: u64,
33}
34
35#[derive(Clone, Debug, Eq, PartialEq)]
36pub struct StorageStats {
37    pub layout_version: u64,
38    pub page_count: u64,
39    pub page_table_bytes: u64,
40    pub active_bytes: u64,
41    pub allocated_bytes: u64,
42    pub orphan_bytes_estimate: u64,
43    pub orphan_ratio_basis_points: u64,
44    pub compact_recommended: bool,
45}
46
47#[derive(Clone, Copy, Debug, Eq, PartialEq)]
48pub(crate) enum StableBlobFailpoint {
49    OverlayWrite,
50    OverlayTruncate,
51    CommitCapacity,
52    CommitChunkWrite,
53    CommitPageTableWrite,
54    CommitSuperblockStore,
55}
56
57thread_local! {
58    #[cfg(test)]
59    static FAILPOINTS: RefCell<BTreeMap<ContextId, StableBlobFailpoint>> = const { RefCell::new(BTreeMap::new()) };
60    static READ_TABLE_CACHE: RefCell<Vec<(ContextId, ReadTableCache)>> = const { RefCell::new(Vec::new()) };
61    static COMMIT_SEGMENT_CACHE: RefCell<Vec<(ContextId, CommitSegmentCache)>> = const { RefCell::new(Vec::new()) };
62}
63
64#[derive(Clone, Copy, Debug, Eq, PartialEq)]
65struct ReadCacheKey {
66    page_table_offset: u64,
67    page_count: u64,
68    db_size: u64,
69    last_tx_id: u64,
70}
71
72#[derive(Debug)]
73struct ReadTableCache {
74    key: Option<ReadCacheKey>,
75    root: Vec<u64>,
76    segments: Vec<CachedSegment>,
77}
78
79#[derive(Debug)]
80struct CachedSegment {
81    segment_no: u64,
82    table: Vec<u64>,
83}
84
85#[derive(Debug)]
86struct CommitSegmentCache {
87    segment_no: u64,
88    segment_offset: u64,
89    table: Vec<u64>,
90}
91
92impl ReadTableCache {
93    fn new() -> Self {
94        Self {
95            key: None,
96            root: Vec::new(),
97            segments: Vec::new(),
98        }
99    }
100
101    fn clear(&mut self) {
102        self.key = None;
103        self.root.clear();
104        self.segments.clear();
105    }
106
107    fn ensure_key(&mut self, key: ReadCacheKey) {
108        if self.key == Some(key) {
109            return;
110        }
111        self.clear();
112        self.key = Some(key);
113    }
114
115    #[inline(always)]
116    fn segment_page_offset(&mut self, segment_no: u64, index: usize) -> Option<u64> {
117        if self.segments.is_empty() {
118            return None;
119        }
120        if self.segments.len() == 1 {
121            let segment = &self.segments[0];
122            if segment.segment_no == segment_no {
123                return Some(segment.table[index]);
124            }
125            return None;
126        }
127        let position = self
128            .segments
129            .iter()
130            .position(|segment| segment.segment_no == segment_no)?;
131        let offset = Some(self.segments[position].table[index]);
132        if position + 1 != self.segments.len() {
133            let segment = self.segments.remove(position);
134            self.segments.push(segment);
135        }
136        offset
137    }
138
139    fn insert_segment(&mut self, segment_no: u64, table: Vec<u64>) {
140        if let Some(position) = self
141            .segments
142            .iter()
143            .position(|segment| segment.segment_no == segment_no)
144        {
145            self.segments.remove(position);
146        }
147        self.segments.push(CachedSegment { segment_no, table });
148        while self.segments.len() > READ_SEGMENT_CACHE_CAPACITY {
149            self.segments.remove(0);
150        }
151    }
152}
153
154#[derive(Debug)]
155pub(crate) struct PageOffsetCache {
156    entries: Vec<(u64, u64)>,
157    pages: Vec<(u64, Vec<u8>)>,
158}
159
160impl PageOffsetCache {
161    pub(crate) fn new() -> Self {
162        Self {
163            entries: Vec::with_capacity(FILE_PAGE_OFFSET_CACHE_CAPACITY),
164            pages: Vec::new(),
165        }
166    }
167
168    fn get(&self, page_no: u64) -> Option<u64> {
169        match self.entries.as_slice() {
170            [] => None,
171            [(cached_page, physical)] => (*cached_page == page_no).then_some(*physical),
172            entries => {
173                for (cached_page, physical) in entries {
174                    if *cached_page == page_no {
175                        return Some(*physical);
176                    }
177                }
178                None
179            }
180        }
181    }
182
183    fn insert(&mut self, page_no: u64, physical: u64) {
184        if self.entries.len() == FILE_PAGE_OFFSET_CACHE_CAPACITY {
185            self.entries.remove(0);
186        }
187        self.entries.push((page_no, physical));
188    }
189
190    #[inline(always)]
191    fn copy_page_slice(&self, page_no: u64, in_page: usize, dst: &mut [u8]) -> bool {
192        if self.pages.is_empty() {
193            return false;
194        }
195        if self.pages.len() == 1 {
196            let (cached_page, page) = &self.pages[0];
197            if *cached_page == page_no {
198                let end = in_page + dst.len();
199                dst.copy_from_slice(&page[in_page..end]);
200                return true;
201            }
202            return false;
203        }
204        for (cached_page, page) in &self.pages {
205            if *cached_page == page_no {
206                let end = in_page + dst.len();
207                dst.copy_from_slice(&page[in_page..end]);
208                return true;
209            }
210        }
211        false
212    }
213
214    fn insert_page(&mut self, page_no: u64, page: Vec<u8>) {
215        if self.pages.len() == FILE_PAGE_DATA_CACHE_CAPACITY {
216            self.pages.remove(0);
217        }
218        self.pages.push((page_no, page));
219    }
220}
221
222#[cfg(test)]
223pub(crate) fn set_failpoint(failpoint: StableBlobFailpoint) {
224    if let Ok(context) = memory::active_context_id() {
225        FAILPOINTS.with(|slot| {
226            slot.borrow_mut().insert(context, failpoint);
227        });
228    }
229}
230
231#[cfg(test)]
232pub(crate) fn clear_failpoint() {
233    FAILPOINTS.with(|slot| slot.borrow_mut().clear());
234}
235
236pub(crate) fn ensure_page_map_layout() -> Result<(), StableMemoryError> {
237    let block = Superblock::load()?;
238    if block.layout_version >= PAGE_MAP_LAYOUT_VERSION {
239        return Ok(());
240    }
241    Err(StableMemoryError::UnsupportedLayoutVersion(
242        block.layout_version,
243    ))
244}
245
246pub(crate) fn begin_update() -> Result<u64, StableMemoryError> {
247    let block = Superblock::load()?;
248    if block.layout_version < PAGE_MAP_LAYOUT_VERSION {
249        return Err(StableMemoryError::UnsupportedLayoutVersion(
250            block.layout_version,
251        ));
252    }
253    if block.is_importing() {
254        return Err(StableMemoryError::ImportAlreadyStarted);
255    }
256    overlay::begin(block.db_size)?;
257    Ok(block.db_size)
258}
259
260pub(crate) fn rollback_update() {
261    overlay::rollback();
262}
263
264#[doc(hidden)]
265pub fn invalidate_read_cache() {
266    READ_TABLE_CACHE.with(|cache| cache.borrow_mut().clear());
267    COMMIT_SEGMENT_CACHE.with(|cache| cache.borrow_mut().clear());
268}
269
270pub(crate) fn commit_update() -> Result<(), StableMemoryError> {
271    let Some(overlay) = overlay::take() else {
272        return Ok(());
273    };
274    if overlay.is_empty() {
275        return Ok(());
276    }
277    commit_overlay(overlay, true)
278}
279
280pub(crate) fn read_at(offset: u64, dst: &mut [u8]) -> Result<bool, StableMemoryError> {
281    if let Some(result) = overlay::read_at(offset, dst) {
282        return result;
283    }
284    read_base_at(offset, dst)
285}
286
287pub(crate) fn read_base_at(offset: u64, dst: &mut [u8]) -> Result<bool, StableMemoryError> {
288    if dst.is_empty() {
289        return Ok(true);
290    }
291    let block = Superblock::load()?;
292    read_base_at_with_block(&block, offset, dst)
293}
294
295pub(crate) fn read_base_at_with_block(
296    block: &Superblock,
297    offset: u64,
298    dst: &mut [u8],
299) -> Result<bool, StableMemoryError> {
300    if dst.is_empty() {
301        return Ok(true);
302    }
303    if offset >= block.db_size {
304        dst.fill(0);
305        return Ok(false);
306    }
307    let requested = u64::try_from(dst.len()).map_err(|_| StableMemoryError::OffsetOverflow)?;
308    if requested <= block.db_size - offset {
309        read_logical_range(block, offset, dst)?;
310        return Ok(true);
311    }
312    let copied = requested.min(block.db_size - offset);
313    let copied_len = usize::try_from(copied).map_err(|_| StableMemoryError::OffsetOverflow)?;
314    read_logical_range(block, offset, &mut dst[..copied_len])?;
315    dst[copied_len..].fill(0);
316    Ok(copied == requested)
317}
318
319#[inline(always)]
320pub(crate) fn read_base_at_with_page_cache(
321    block: &Superblock,
322    offset: u64,
323    dst: &mut [u8],
324    page_offsets: &mut PageOffsetCache,
325) -> Result<bool, StableMemoryError> {
326    if dst.is_empty() {
327        return Ok(true);
328    }
329    if offset >= block.db_size {
330        dst.fill(0);
331        return Ok(false);
332    }
333    let requested = u64::try_from(dst.len()).map_err(|_| StableMemoryError::OffsetOverflow)?;
334    if requested <= block.db_size - offset {
335        read_logical_range_with_page_cache(block, offset, dst, page_offsets)?;
336        return Ok(true);
337    }
338    let copied = requested.min(block.db_size - offset);
339    let copied_len = usize::try_from(copied).map_err(|_| StableMemoryError::OffsetOverflow)?;
340    read_logical_range_with_page_cache(block, offset, &mut dst[..copied_len], page_offsets)?;
341    dst[copied_len..].fill(0);
342    Ok(copied == requested)
343}
344
345pub(crate) fn read_base_page(page_no: u64) -> Result<Vec<u8>, StableMemoryError> {
346    let block = Superblock::load()?;
347    let mut page = zero_page();
348    if page_no >= active_page_count(&block)? {
349        return Ok(page);
350    }
351    let physical = page_offset_for(&block, page_no)?;
352    if physical != 0 {
353        #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
354        crate::read_metrics::record_stable_data_read(page.len());
355        memory::read_preallocated(physical, &mut page)?;
356    }
357    Ok(page)
358}
359
360pub(crate) fn write_at(offset: u64, bytes: &[u8]) -> Result<(), StableMemoryError> {
361    if let Some(result) = overlay::write_at(offset, bytes) {
362        hit_failpoint(StableBlobFailpoint::OverlayWrite)?;
363        return result;
364    }
365    if bytes.is_empty() {
366        return Ok(());
367    }
368    ensure_page_map_layout()?;
369    let mut direct = Overlay::new(Superblock::load()?.db_size);
370    direct.write_at(offset, bytes)?;
371    commit_overlay(direct, false)
372}
373
374pub(crate) fn truncate(size: u64) -> Result<(), StableMemoryError> {
375    if let Some(result) = overlay::truncate(size) {
376        hit_failpoint(StableBlobFailpoint::OverlayTruncate)?;
377        return result;
378    }
379    ensure_page_map_layout()?;
380    let mut direct = Overlay::new(Superblock::load()?.db_size);
381    direct.truncate(size)?;
382    if direct.is_empty() {
383        return Ok(());
384    }
385    commit_overlay(direct, false)
386}
387
388pub(crate) fn file_size() -> Result<u64, StableMemoryError> {
389    if let Some(size) = overlay::file_size() {
390        return Ok(size);
391    }
392    Ok(Superblock::load()?.db_size)
393}
394
395pub fn export_chunk(offset: u64, len: u64) -> Result<Vec<u8>, StableMemoryError> {
396    reject_during_update()?;
397    let block = Superblock::load()?;
398    if offset >= block.db_size {
399        return Ok(Vec::new());
400    }
401    let copied = len.min(block.db_size - offset);
402    let copied_len = usize::try_from(copied).map_err(|_| StableMemoryError::OffsetOverflow)?;
403    let mut out = vec![0_u8; copied_len];
404    read_logical_range(&block, offset, &mut out)?;
405    Ok(out)
406}
407
408pub fn import_chunk(offset: u64, bytes: &[u8]) -> Result<(), StableMemoryError> {
409    reject_during_update()?;
410    let mut block = Superblock::load()?;
411    if !block.is_importing() {
412        return Err(StableMemoryError::ImportNotStarted);
413    }
414    let len = u64::try_from(bytes.len()).map_err(|_| StableMemoryError::OffsetOverflow)?;
415    if offset != block.import_written_until {
416        return Err(StableMemoryError::ImportOutOfOrder {
417            offset,
418            expected: block.import_written_until,
419        });
420    }
421    let end = checked_add(offset, len)?;
422    if end > block.import_total_size {
423        return Err(StableMemoryError::ImportOutOfBounds {
424            offset,
425            len,
426            db_size: block.import_total_size,
427        });
428    }
429    memory::write(import_offset(&block, offset)?, bytes)?;
430    block.import_written_until = end;
431    block.store()?;
432    invalidate_read_cache();
433    Ok(())
434}
435
436pub fn begin_import(total_size: u64, expected_checksum: u64) -> Result<(), StableMemoryError> {
437    reject_during_update()?;
438    let mut block = Superblock::load()?;
439    if block.is_importing() {
440        return Err(StableMemoryError::ImportAlreadyStarted);
441    }
442    let import_base_offset = append_base()?;
443    checked_add(import_base_offset, total_size)?;
444    block.flags |= FLAG_IMPORTING;
445    block.clear_checksum_refresh();
446    block.import_expected_checksum = expected_checksum;
447    block.import_written_until = 0;
448    block.import_total_size = total_size;
449    block.import_base_offset = import_base_offset;
450    block.store()?;
451    invalidate_read_cache();
452    Ok(())
453}
454
455pub fn finish_import() -> Result<(), StableMemoryError> {
456    reject_during_update()?;
457    let mut block = Superblock::load()?;
458    if !block.is_importing() {
459        return Err(StableMemoryError::ImportNotStarted);
460    }
461    if block.import_written_until != block.import_total_size {
462        return Err(StableMemoryError::ImportIncomplete {
463            written_until: block.import_written_until,
464            db_size: block.import_total_size,
465        });
466    }
467    let checksum = checksum_physical_range(block.import_base_offset, block.import_total_size)?;
468    if checksum != block.import_expected_checksum {
469        let expected = block.import_expected_checksum;
470        clear_import(&mut block)?;
471        return Err(StableMemoryError::ChecksumMismatch {
472            expected,
473            actual: checksum,
474        });
475    }
476    let entries = imported_page_table(&block)?;
477    let (root_offset, root_len) = write_segmented_tables(&entries)?;
478    block.db_size = block.import_total_size;
479    block.db_base_offset = block.import_base_offset;
480    block.page_table_offset = root_offset;
481    block.page_count = root_len;
482    block.layout_version = PAGE_MAP_LAYOUT_VERSION;
483    block.flags &= !FLAG_IMPORTING;
484    block.flags &= !FLAG_CHECKSUM_STALE;
485    block.clear_checksum_refresh();
486    block.checksum = checksum;
487    block.import_expected_checksum = 0;
488    block.import_written_until = 0;
489    block.import_total_size = 0;
490    block.import_base_offset = 0;
491    block.store()?;
492    invalidate_read_cache();
493    Ok(())
494}
495
496pub fn cancel_import() -> Result<(), StableMemoryError> {
497    reject_during_update()?;
498    let mut block = Superblock::load()?;
499    if !block.is_importing() {
500        return Err(StableMemoryError::ImportNotStarted);
501    }
502    clear_import(&mut block)
503}
504
505pub fn refresh_checksum() -> Result<u64, StableMemoryError> {
506    reject_during_update()?;
507    let checksum = checksum()?;
508    let mut block = Superblock::load()?;
509    block.checksum = checksum;
510    block.flags &= !FLAG_CHECKSUM_STALE;
511    block.clear_checksum_refresh();
512    block.store()?;
513    invalidate_read_cache();
514    Ok(checksum)
515}
516
517pub fn refresh_checksum_chunk(max_bytes: u64) -> Result<ChecksumRefresh, StableMemoryError> {
518    reject_during_update()?;
519    if max_bytes == 0 {
520        return Err(StableMemoryError::ChecksumRefreshChunkEmpty);
521    }
522
523    let mut block = Superblock::load()?;
524    if block.is_importing() {
525        return Err(StableMemoryError::ImportAlreadyStarted);
526    }
527    if !block.is_checksum_refreshing() {
528        block.flags |= FLAG_CHECKSUM_REFRESHING;
529        block.checksum_refresh_offset = 0;
530        block.checksum_refresh_hash = fnv1a64(&[]);
531        block.checksum_refresh_tx_id = block.last_tx_id;
532    }
533    if block.checksum_refresh_tx_id != block.last_tx_id {
534        block.clear_checksum_refresh();
535        block.store()?;
536        invalidate_read_cache();
537        return refresh_checksum_chunk(max_bytes);
538    }
539
540    let start = block.checksum_refresh_offset;
541    let end = block.db_size.min(start.saturating_add(max_bytes));
542    let mut offset = start;
543    let mut hash = block.checksum_refresh_hash;
544    while offset < end {
545        let len = (end - offset).min(CHECKSUM_CHUNK_LEN);
546        let copied_len = usize::try_from(len).map_err(|_| StableMemoryError::OffsetOverflow)?;
547        let mut bytes = vec![0_u8; copied_len];
548        read_logical_range(&block, offset, &mut bytes)?;
549        hash = fold_fnv1a64(hash, &bytes);
550        offset += len;
551    }
552
553    block.checksum_refresh_offset = offset;
554    block.checksum_refresh_hash = hash;
555    if offset == block.db_size {
556        block.checksum = hash;
557        block.flags &= !FLAG_CHECKSUM_STALE;
558        block.clear_checksum_refresh();
559    }
560    let out = ChecksumRefresh {
561        complete: offset == block.db_size,
562        checksum: hash,
563        scanned_bytes: offset,
564        db_size: block.db_size,
565    };
566    block.store()?;
567    invalidate_read_cache();
568    Ok(out)
569}
570
571pub fn checksum() -> Result<u64, StableMemoryError> {
572    reject_during_update()?;
573    let block = Superblock::load()?;
574    checksum_logical_range(&block, block.db_size)
575}
576
577pub fn compact() -> Result<(), StableMemoryError> {
578    reject_during_update()?;
579    ensure_page_map_layout()?;
580    let block = Superblock::load()?;
581    let table = read_page_table(&block)?;
582    let mut compacted = Vec::with_capacity(table.len());
583    let mut cursor = append_base()?;
584    let non_zero_pages = table.iter().filter(|offset| **offset != 0).count();
585    let data_bytes = u64::try_from(non_zero_pages)
586        .map_err(|_| StableMemoryError::OffsetOverflow)?
587        .checked_mul(page_size())
588        .ok_or(StableMemoryError::OffsetOverflow)?;
589    memory::ensure_capacity(checked_add(cursor, data_bytes)?)?;
590
591    for offset in table {
592        if offset == 0 {
593            compacted.push(0);
594            continue;
595        }
596        let mut page = zero_page();
597        memory::read_preallocated(offset, &mut page)?;
598        memory::write_preallocated(cursor, &page)?;
599        compacted.push(cursor);
600        cursor = checked_add(cursor, page_size())?;
601    }
602
603    let (root_offset, root_len) = write_segmented_tables(&compacted)?;
604    Superblock::store_page_map_without_tx(root_offset, root_len, block.db_size)?;
605    invalidate_read_cache();
606    Ok(())
607}
608
609pub fn storage_stats() -> Result<StorageStats, StableMemoryError> {
610    let block = Superblock::load()?;
611    let table = read_page_table(&block)?;
612    let non_zero_pages = u64::try_from(table.iter().filter(|offset| **offset != 0).count())
613        .map_err(|_| StableMemoryError::OffsetOverflow)?;
614    let segment_count = active_segment_count(&block)?;
615    let root_bytes = root_table_bytes(segment_count)?;
616    let segment_bytes = segment_count
617        .checked_mul(segment_table_bytes()?)
618        .ok_or(StableMemoryError::OffsetOverflow)?;
619    let page_table_bytes = checked_add(root_bytes, segment_bytes)?;
620    let active_bytes = SUPERBLOCK_SIZE
621        .checked_add(non_zero_pages.saturating_mul(page_size()))
622        .and_then(|value| value.checked_add(page_table_bytes))
623        .ok_or(StableMemoryError::OffsetOverflow)?;
624    let allocated_bytes = memory::size_pages()
625        .checked_mul(STABLE_PAGE_SIZE)
626        .ok_or(StableMemoryError::OffsetOverflow)?;
627    let orphan_bytes_estimate = allocated_bytes.saturating_sub(active_bytes);
628    let orphan_ratio_basis_points = orphan_bytes_estimate
629        .saturating_mul(10_000)
630        .checked_div(active_bytes)
631        .unwrap_or(0);
632    Ok(StorageStats {
633        layout_version: block.layout_version,
634        page_count: active_page_count(&block)?,
635        page_table_bytes,
636        active_bytes,
637        allocated_bytes,
638        orphan_bytes_estimate,
639        orphan_ratio_basis_points,
640        compact_recommended: orphan_bytes_estimate >= active_bytes
641            && orphan_bytes_estimate >= COMPACT_MIN_ORPHAN_BYTES,
642    })
643}
644
645pub(crate) fn page_count_for_size(size: u64) -> Result<u64, StableMemoryError> {
646    Ok(size.div_ceil(page_size()))
647}
648
649#[cfg(test)]
650pub(crate) fn debug_root_table_for_tests() -> Result<Vec<u64>, StableMemoryError> {
651    let block = Superblock::load()?;
652    read_root_table(&block)
653}
654
655fn commit_overlay(overlay: Overlay, advance_tx: bool) -> Result<(), StableMemoryError> {
656    hit_failpoint(StableBlobFailpoint::CommitCapacity)?;
657    let profile_enabled = commit_profile_enabled();
658    let block = Superblock::load()?;
659    let overlay_size = overlay.size();
660    let final_page_count = page_count_for_size(overlay_size)?;
661    let data_cursor = append_base()?;
662    debug_assert!(overlay
663        .dirty_pages()
664        .iter()
665        .all(|(page_no, _)| *page_no < final_page_count));
666    let dirty_pages = overlay.dirty_pages();
667    if let [(page_no, page)] = dirty_pages {
668        if overlay_size >= block.db_size
669            && *page_no < final_page_count
670            && final_page_count <= SEGMENT_PAGE_COUNT
671        {
672            let build_profile_start = commit_profile_start(profile_enabled);
673            let options = SinglePageCommitOptions {
674                advance_tx,
675                overlay_size,
676                data_cursor,
677                profile_enabled,
678                build_profile_start,
679            };
680            return commit_single_segment_page_overlay(&block, *page_no, page, options);
681        }
682    }
683
684    let final_segment_count = segment_count_for_pages(final_page_count)?;
685    let profile_start = commit_profile_start(profile_enabled);
686    let mut root = read_commit_root_table(&block)?;
687    commit_profile_record_load(profile_start);
688
689    let build_profile_start = commit_profile_start(profile_enabled);
690    let root_len =
691        usize::try_from(final_segment_count).map_err(|_| StableMemoryError::OffsetOverflow)?;
692    if root.len() != root_len {
693        root.resize(root_len, 0);
694    }
695
696    if let [(page_no, page)] = dirty_pages {
697        if overlay_size >= block.db_size && *page_no < final_page_count {
698            let options = SinglePageCommitOptions {
699                advance_tx,
700                overlay_size,
701                data_cursor,
702                profile_enabled,
703                build_profile_start,
704            };
705            return commit_single_page_overlay(
706                &block,
707                final_segment_count,
708                root,
709                *page_no,
710                page,
711                options,
712            );
713        }
714    }
715
716    let mut segment_updates = BTreeMap::<u64, Vec<u64>>::new();
717    let mut page_cursor = data_cursor;
718
719    for (page_no, _) in dirty_pages {
720        if *page_no >= final_page_count {
721            continue;
722        }
723        let segment_no = segment_no(*page_no);
724        let index = segment_index(*page_no)?;
725        let table = load_segment_for_update(&block, &root, &mut segment_updates, segment_no)?;
726        table[index] = page_cursor;
727        page_cursor = checked_add(page_cursor, page_size())?;
728    }
729
730    if overlay_size < block.db_size {
731        clear_truncated_tail(&block, &root, &mut segment_updates, final_page_count)?;
732    }
733    commit_profile_record_build_segments(build_profile_start);
734
735    let mut table_cursor = page_cursor;
736    let root_entries_len = final_segment_count;
737    let segment_table_writes = segment_updates.len();
738    let segment_table_bytes = u64::try_from(segment_table_writes)
739        .map_err(|_| StableMemoryError::OffsetOverflow)?
740        .checked_mul(segment_table_bytes()?)
741        .ok_or(StableMemoryError::OffsetOverflow)?;
742    let page_table_bytes = checked_add(segment_table_bytes, root_table_bytes(root_entries_len)?)?;
743    let profile_start = commit_profile_start(profile_enabled);
744    memory::ensure_capacity(checked_add(table_cursor, page_table_bytes)?)?;
745    commit_profile_record_capacity(profile_start);
746
747    let profile_start = commit_profile_start(profile_enabled);
748    let mut cursor = data_cursor;
749    for (_, page) in dirty_pages {
750        hit_failpoint(StableBlobFailpoint::CommitChunkWrite)?;
751        write_commit_page(cursor, page, profile_enabled)?;
752        cursor = checked_add(cursor, page_size())?;
753    }
754    commit_profile_record_page_write(profile_start);
755
756    hit_failpoint(StableBlobFailpoint::CommitPageTableWrite)?;
757    let profile_start = commit_profile_start(profile_enabled);
758    for (segment_no, table) in segment_updates {
759        let offset = write_commit_segment_table_at(&table, &mut table_cursor, profile_enabled)?;
760        let index = usize::try_from(segment_no).map_err(|_| StableMemoryError::OffsetOverflow)?;
761        root[index] = offset;
762    }
763    let root_offset = write_commit_root_table_at(&root, &mut table_cursor, profile_enabled)?;
764    commit_profile_record_table_write(profile_start);
765
766    hit_failpoint(StableBlobFailpoint::CommitSuperblockStore)?;
767    let profile_start = commit_profile_start(profile_enabled);
768    let result = store_commit_page_map(
769        advance_tx,
770        root_offset,
771        root_entries_len,
772        overlay_size,
773        profile_enabled,
774    );
775    commit_profile_record_superblock_store(profile_start);
776    result
777}
778
779#[derive(Clone, Copy)]
780struct SinglePageCommitOptions {
781    advance_tx: bool,
782    overlay_size: u64,
783    data_cursor: u64,
784    profile_enabled: bool,
785    build_profile_start: Option<u64>,
786}
787
788fn commit_single_page_overlay(
789    block: &Superblock,
790    final_segment_count: u64,
791    mut root: Vec<u64>,
792    page_no: u64,
793    page: &[u8],
794    options: SinglePageCommitOptions,
795) -> Result<(), StableMemoryError> {
796    let segment_no = segment_no(page_no);
797    let index = segment_index(page_no)?;
798    let mut table = read_commit_segment_table(block, &root, segment_no)?;
799    table[index] = options.data_cursor;
800    let page_cursor = checked_add(options.data_cursor, page_size())?;
801    commit_profile_record_build_segments(options.build_profile_start);
802
803    let root_entries_len = final_segment_count;
804    let page_table_bytes =
805        checked_add(segment_table_bytes()?, root_table_bytes(root_entries_len)?)?;
806    let profile_start = commit_profile_start(options.profile_enabled);
807    memory::ensure_capacity(checked_add(page_cursor, page_table_bytes)?)?;
808    commit_profile_record_capacity(profile_start);
809
810    hit_failpoint(StableBlobFailpoint::CommitChunkWrite)?;
811    let profile_start = commit_profile_start(options.profile_enabled);
812    write_commit_page(options.data_cursor, page, options.profile_enabled)?;
813    commit_profile_record_page_write(profile_start);
814
815    hit_failpoint(StableBlobFailpoint::CommitPageTableWrite)?;
816    let profile_start = commit_profile_start(options.profile_enabled);
817    let mut table_cursor = page_cursor;
818    let offset = write_commit_segment_table_at(&table, &mut table_cursor, options.profile_enabled)?;
819    let root_offset = if final_segment_count == 1 {
820        write_commit_root_table_at(&[offset], &mut table_cursor, options.profile_enabled)?
821    } else {
822        let root_index =
823            usize::try_from(segment_no).map_err(|_| StableMemoryError::OffsetOverflow)?;
824        root[root_index] = offset;
825        write_commit_root_table_at(&root, &mut table_cursor, options.profile_enabled)?
826    };
827    commit_profile_record_table_write(profile_start);
828
829    hit_failpoint(StableBlobFailpoint::CommitSuperblockStore)?;
830    let profile_start = commit_profile_start(options.profile_enabled);
831    let result = store_commit_page_map(
832        options.advance_tx,
833        root_offset,
834        root_entries_len,
835        options.overlay_size,
836        options.profile_enabled,
837    );
838    commit_profile_record_superblock_store(profile_start);
839    if result.is_ok() {
840        cache_commit_segment_table(segment_no, offset, table);
841    }
842    result
843}
844
845fn commit_single_segment_page_overlay(
846    block: &Superblock,
847    page_no: u64,
848    page: &[u8],
849    options: SinglePageCommitOptions,
850) -> Result<(), StableMemoryError> {
851    let index = segment_index(page_no)?;
852    let root = read_commit_root_table(block)?;
853    let mut table = read_commit_segment_table(block, &root, 0)?;
854    table[index] = options.data_cursor;
855    let page_cursor = checked_add(options.data_cursor, page_size())?;
856    commit_profile_record_build_segments(options.build_profile_start);
857
858    let profile_start = commit_profile_start(options.profile_enabled);
859    memory::ensure_capacity(checked_add(page_cursor, SINGLE_SEGMENT_PAGE_TABLE_BYTES)?)?;
860    commit_profile_record_capacity(profile_start);
861
862    hit_failpoint(StableBlobFailpoint::CommitChunkWrite)?;
863    let profile_start = commit_profile_start(options.profile_enabled);
864    memory::write_prechecked(options.data_cursor, page)?;
865    commit_profile_record_page_write(profile_start);
866
867    hit_failpoint(StableBlobFailpoint::CommitPageTableWrite)?;
868    let profile_start = commit_profile_start(options.profile_enabled);
869    let mut table_cursor = page_cursor;
870    let offset = write_commit_segment_table_at(&table, &mut table_cursor, options.profile_enabled)?;
871    let root_offset =
872        write_commit_root_table_at(&[offset], &mut table_cursor, options.profile_enabled)?;
873    commit_profile_record_table_write(profile_start);
874
875    hit_failpoint(StableBlobFailpoint::CommitSuperblockStore)?;
876    let profile_start = commit_profile_start(options.profile_enabled);
877    let result = store_commit_page_map(
878        options.advance_tx,
879        root_offset,
880        1,
881        options.overlay_size,
882        options.profile_enabled,
883    );
884    commit_profile_record_superblock_store(profile_start);
885    if result.is_ok() {
886        cache_commit_segment_table(0, offset, table);
887    }
888    result
889}
890
891#[cfg(any(test, debug_assertions, feature = "bench-profile"))]
892#[inline(always)]
893fn commit_profile_enabled() -> bool {
894    crate::read_metrics::metrics_enabled()
895}
896
897#[cfg(not(any(test, debug_assertions, feature = "bench-profile")))]
898#[inline(always)]
899fn commit_profile_enabled() -> bool {
900    false
901}
902
903#[cfg(any(test, debug_assertions, feature = "bench-profile"))]
904#[inline(always)]
905fn commit_profile_start(enabled: bool) -> Option<u64> {
906    if enabled {
907        Some(crate::read_metrics::instruction_counter())
908    } else {
909        None
910    }
911}
912
913#[cfg(not(any(test, debug_assertions, feature = "bench-profile")))]
914#[inline(always)]
915fn commit_profile_start(_enabled: bool) -> Option<u64> {
916    None
917}
918
919macro_rules! commit_profile_recorder {
920    ($name:ident, $record:ident) => {
921        #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
922        #[inline(always)]
923        fn $name(start: Option<u64>) {
924            if let Some(start) = start {
925                crate::read_metrics::$record(
926                    crate::read_metrics::instruction_counter().saturating_sub(start),
927                );
928            }
929        }
930
931        #[cfg(not(any(test, debug_assertions, feature = "bench-profile")))]
932        #[inline(always)]
933        fn $name(_start: Option<u64>) {}
934    };
935}
936
937commit_profile_recorder!(commit_profile_record_load, record_commit_load);
938commit_profile_recorder!(
939    commit_profile_record_build_segments,
940    record_commit_build_segments
941);
942commit_profile_recorder!(commit_profile_record_capacity, record_commit_capacity);
943commit_profile_recorder!(commit_profile_record_page_write, record_commit_page_write);
944commit_profile_recorder!(commit_profile_record_table_write, record_commit_table_write);
945commit_profile_recorder!(
946    commit_profile_record_superblock_store,
947    record_commit_superblock_store
948);
949
950#[inline(always)]
951fn write_commit_page(
952    offset: u64,
953    page: &[u8],
954    profile_enabled: bool,
955) -> Result<(), StableMemoryError> {
956    if profile_enabled {
957        memory::write_prechecked(offset, page)
958    } else {
959        memory::write_prechecked_unmetered(offset, page)
960    }
961}
962
963fn store_commit_page_map(
964    advance_tx: bool,
965    root_offset: u64,
966    root_entries_len: u64,
967    overlay_size: u64,
968    profile_enabled: bool,
969) -> Result<(), StableMemoryError> {
970    match (advance_tx, profile_enabled) {
971        (true, true) => Superblock::commit_page_map(root_offset, root_entries_len, overlay_size),
972        (true, false) => {
973            Superblock::commit_page_map_unmetered(root_offset, root_entries_len, overlay_size)
974        }
975        (false, true) => {
976            Superblock::store_page_map_without_tx(root_offset, root_entries_len, overlay_size)
977        }
978        (false, false) => Superblock::store_page_map_without_tx_unmetered(
979            root_offset,
980            root_entries_len,
981            overlay_size,
982        ),
983    }
984}
985
986fn load_segment_for_update<'a>(
987    block: &Superblock,
988    root: &[u64],
989    updates: &'a mut BTreeMap<u64, Vec<u64>>,
990    segment_no: u64,
991) -> Result<&'a mut Vec<u64>, StableMemoryError> {
992    match updates.entry(segment_no) {
993        std::collections::btree_map::Entry::Occupied(entry) => Ok(entry.into_mut()),
994        std::collections::btree_map::Entry::Vacant(entry) => {
995            let table = read_segment_table(block, root, segment_no)?;
996            Ok(entry.insert(table))
997        }
998    }
999}
1000
1001fn clear_truncated_tail(
1002    block: &Superblock,
1003    root: &[u64],
1004    updates: &mut BTreeMap<u64, Vec<u64>>,
1005    final_page_count: u64,
1006) -> Result<(), StableMemoryError> {
1007    let old_page_count = active_page_count(block)?;
1008    if final_page_count >= old_page_count || final_page_count == 0 {
1009        return Ok(());
1010    }
1011    let boundary_segment = segment_no(final_page_count);
1012    if boundary_segment >= segment_count_for_pages(final_page_count)? {
1013        return Ok(());
1014    }
1015    let start = segment_index(final_page_count)?;
1016    if start == 0 {
1017        return Ok(());
1018    }
1019    let table = load_segment_for_update(block, root, updates, boundary_segment)?;
1020    table[start..].fill(0);
1021    Ok(())
1022}
1023
1024fn reject_during_update() -> Result<(), StableMemoryError> {
1025    if overlay::is_active() {
1026        Err(StableMemoryError::UpdateInProgress)
1027    } else {
1028        Ok(())
1029    }
1030}
1031
1032fn read_logical_range(
1033    block: &Superblock,
1034    offset: u64,
1035    dst: &mut [u8],
1036) -> Result<(), StableMemoryError> {
1037    if dst.is_empty() {
1038        return Ok(());
1039    }
1040    let in_page =
1041        usize::try_from(offset % page_size()).map_err(|_| StableMemoryError::OffsetOverflow)?;
1042    if dst.len() <= page_len() - in_page {
1043        return read_logical_page_slice(block, offset / page_size(), in_page, dst);
1044    }
1045
1046    let mut copied_total = 0_usize;
1047    while copied_total < dst.len() {
1048        let absolute = checked_add(
1049            offset,
1050            u64::try_from(copied_total).map_err(|_| StableMemoryError::OffsetOverflow)?,
1051        )?;
1052        let page_no = absolute / page_size();
1053        let in_page = usize::try_from(absolute % page_size())
1054            .map_err(|_| StableMemoryError::OffsetOverflow)?;
1055        let copied = (page_len() - in_page).min(dst.len() - copied_total);
1056        read_logical_page_slice(
1057            block,
1058            page_no,
1059            in_page,
1060            &mut dst[copied_total..copied_total + copied],
1061        )?;
1062        copied_total += copied;
1063    }
1064    Ok(())
1065}
1066
1067fn read_logical_range_with_page_cache(
1068    block: &Superblock,
1069    offset: u64,
1070    dst: &mut [u8],
1071    page_offsets: &mut PageOffsetCache,
1072) -> Result<(), StableMemoryError> {
1073    let in_page =
1074        usize::try_from(offset % page_size()).map_err(|_| StableMemoryError::OffsetOverflow)?;
1075    if dst.len() <= page_len() - in_page {
1076        return read_logical_page_slice_with_page_cache(
1077            block,
1078            offset / page_size(),
1079            in_page,
1080            dst,
1081            page_offsets,
1082        );
1083    }
1084
1085    let mut copied_total = 0_usize;
1086    while copied_total < dst.len() {
1087        let absolute = checked_add(
1088            offset,
1089            u64::try_from(copied_total).map_err(|_| StableMemoryError::OffsetOverflow)?,
1090        )?;
1091        let page_no = absolute / page_size();
1092        let in_page = usize::try_from(absolute % page_size())
1093            .map_err(|_| StableMemoryError::OffsetOverflow)?;
1094        let copied = (page_len() - in_page).min(dst.len() - copied_total);
1095        read_logical_page_slice_with_page_cache(
1096            block,
1097            page_no,
1098            in_page,
1099            &mut dst[copied_total..copied_total + copied],
1100            page_offsets,
1101        )?;
1102        copied_total += copied;
1103    }
1104    Ok(())
1105}
1106
1107fn read_logical_page_slice(
1108    block: &Superblock,
1109    page_no: u64,
1110    in_page: usize,
1111    dst: &mut [u8],
1112) -> Result<(), StableMemoryError> {
1113    let physical = page_offset_for(block, page_no)?;
1114    if physical == 0 {
1115        dst.fill(0);
1116        return Ok(());
1117    }
1118    let stable_offset = checked_add(
1119        physical,
1120        u64::try_from(in_page).map_err(|_| StableMemoryError::OffsetOverflow)?,
1121    )?;
1122    #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1123    crate::read_metrics::record_stable_data_read(dst.len());
1124    memory::read_preallocated(stable_offset, dst)
1125}
1126
1127#[inline(always)]
1128fn read_logical_page_slice_with_page_cache(
1129    block: &Superblock,
1130    page_no: u64,
1131    in_page: usize,
1132    dst: &mut [u8],
1133    page_offsets: &mut PageOffsetCache,
1134) -> Result<(), StableMemoryError> {
1135    if dst.len() < page_len() && page_offsets.copy_page_slice(page_no, in_page, dst) {
1136        return Ok(());
1137    }
1138    let physical = match page_offsets.get(page_no) {
1139        Some(physical) => physical,
1140        None => {
1141            let physical = if block.page_table_offset == 0 {
1142                0
1143            } else {
1144                cached_page_offset_for(block, page_no)?
1145            };
1146            page_offsets.insert(page_no, physical);
1147            physical
1148        }
1149    };
1150    if physical == 0 {
1151        dst.fill(0);
1152        return Ok(());
1153    }
1154    if in_page == 0 && dst.len() == page_len() {
1155        #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1156        crate::read_metrics::record_stable_data_read(dst.len());
1157        return memory::read_preallocated(physical, dst);
1158    }
1159    if dst.len() < page_len() {
1160        let mut page = zero_page();
1161        #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1162        crate::read_metrics::record_stable_data_read(page.len());
1163        memory::read_preallocated(physical, &mut page)?;
1164        let end = in_page + dst.len();
1165        dst.copy_from_slice(&page[in_page..end]);
1166        page_offsets.insert_page(page_no, page);
1167        return Ok(());
1168    }
1169    let stable_offset = checked_add(
1170        physical,
1171        u64::try_from(in_page).map_err(|_| StableMemoryError::OffsetOverflow)?,
1172    )?;
1173    #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1174    crate::read_metrics::record_stable_data_read(dst.len());
1175    memory::read_preallocated(stable_offset, dst)
1176}
1177
1178fn page_offset_for(block: &Superblock, page_no: u64) -> Result<u64, StableMemoryError> {
1179    if page_no >= active_page_count(block)? || block.page_table_offset == 0 {
1180        return Ok(0);
1181    }
1182    cached_page_offset_for(block, page_no)
1183}
1184
1185fn read_page_table(block: &Superblock) -> Result<Vec<u64>, StableMemoryError> {
1186    let root = read_root_table(block)?;
1187    let count = active_page_count(block)?;
1188    let capacity = usize::try_from(count).map_err(|_| StableMemoryError::OffsetOverflow)?;
1189    let mut entries = Vec::with_capacity(capacity);
1190    for segment_no in 0..segment_count_for_pages(count)? {
1191        let table = read_segment_table(block, &root, segment_no)?;
1192        for entry in table {
1193            if entries.len() == capacity {
1194                break;
1195            }
1196            entries.push(entry);
1197        }
1198    }
1199    Ok(entries)
1200}
1201
1202fn cached_page_offset_for(block: &Superblock, page_no: u64) -> Result<u64, StableMemoryError> {
1203    let context = memory::active_context_id()?;
1204    let key = read_cache_key(block);
1205    let segment_no = segment_no(page_no);
1206    let index = segment_index(page_no)?;
1207    READ_TABLE_CACHE.with(|cache| {
1208        let mut caches = cache.borrow_mut();
1209        let cache = match read_table_cache_index(&caches, context) {
1210            Some(index) => &mut caches[index].1,
1211            None => {
1212                caches.push((context, ReadTableCache::new()));
1213                &mut caches
1214                    .last_mut()
1215                    .ok_or(StableMemoryError::OffsetOverflow)?
1216                    .1
1217            }
1218        };
1219        cache.ensure_key(key);
1220        if cache.root.is_empty() {
1221            #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1222            crate::read_metrics::record_page_table_root_miss();
1223            cache.root = read_root_table(block)?;
1224        } else {
1225            #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1226            crate::read_metrics::record_page_table_root_hit();
1227        }
1228        let root_index =
1229            usize::try_from(segment_no).map_err(|_| StableMemoryError::OffsetOverflow)?;
1230        let segment_offset = cache.root[root_index];
1231        if segment_offset == 0 {
1232            return Ok(0);
1233        }
1234        if let Some(offset) = cache.segment_page_offset(segment_no, index) {
1235            #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1236            crate::read_metrics::record_page_table_segment_hit();
1237            return Ok(offset);
1238        }
1239        #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1240        crate::read_metrics::record_page_table_segment_miss();
1241        let table = read_segment_table_at(segment_offset)?;
1242        let offset = table[index];
1243        cache.insert_segment(segment_no, table);
1244        Ok(offset)
1245    })
1246}
1247
1248fn read_table_cache_index(
1249    caches: &[(ContextId, ReadTableCache)],
1250    context: ContextId,
1251) -> Option<usize> {
1252    caches
1253        .iter()
1254        .position(|(stored_context, _)| *stored_context == context)
1255}
1256
1257fn read_root_table(block: &Superblock) -> Result<Vec<u64>, StableMemoryError> {
1258    if block.page_count == 0 {
1259        return Ok(Vec::new());
1260    }
1261    let entries_len =
1262        usize::try_from(block.page_count).map_err(|_| StableMemoryError::OffsetOverflow)?;
1263    read_u64_table_at(block.page_table_offset, entries_len)
1264}
1265
1266fn read_commit_root_table(block: &Superblock) -> Result<Vec<u64>, StableMemoryError> {
1267    read_root_table(block)
1268}
1269
1270fn read_segment_table(
1271    _block: &Superblock,
1272    root: &[u64],
1273    segment_no: u64,
1274) -> Result<Vec<u64>, StableMemoryError> {
1275    let index = usize::try_from(segment_no).map_err(|_| StableMemoryError::OffsetOverflow)?;
1276    let Some(offset) = root.get(index).copied() else {
1277        return Ok(vec![0_u64; segment_page_count_usize()]);
1278    };
1279    if offset == 0 {
1280        return Ok(vec![0_u64; segment_page_count_usize()]);
1281    }
1282    read_segment_table_at(offset)
1283}
1284
1285fn read_commit_segment_table(
1286    _block: &Superblock,
1287    root: &[u64],
1288    segment_no: u64,
1289) -> Result<Vec<u64>, StableMemoryError> {
1290    let index = usize::try_from(segment_no).map_err(|_| StableMemoryError::OffsetOverflow)?;
1291    let Some(offset) = root.get(index).copied() else {
1292        return Ok(vec![0_u64; segment_page_count_usize()]);
1293    };
1294    if offset == 0 {
1295        return Ok(vec![0_u64; segment_page_count_usize()]);
1296    }
1297    read_commit_segment_table_at(segment_no, offset)
1298}
1299
1300fn read_commit_segment_table_at(
1301    segment_no: u64,
1302    offset: u64,
1303) -> Result<Vec<u64>, StableMemoryError> {
1304    if offset == 0 {
1305        return Ok(vec![0_u64; segment_page_count_usize()]);
1306    }
1307    if let Some(table) = take_commit_segment_table(segment_no, offset) {
1308        return Ok(table);
1309    }
1310    read_segment_table_at(offset)
1311}
1312
1313fn take_commit_segment_table(segment_no: u64, segment_offset: u64) -> Option<Vec<u64>> {
1314    let Ok(context) = memory::active_context_id() else {
1315        return None;
1316    };
1317    COMMIT_SEGMENT_CACHE.with(|cache| {
1318        let mut cache = cache.borrow_mut();
1319        if cache.len() == 1 {
1320            let (stored_context, cached) = &cache[0];
1321            if *stored_context == context
1322                && cached.segment_no == segment_no
1323                && cached.segment_offset == segment_offset
1324            {
1325                return cache.pop().map(|(_, cached)| cached.table);
1326            }
1327            return None;
1328        }
1329        cache
1330            .iter()
1331            .position(|(stored_context, cached)| {
1332                *stored_context == context
1333                    && cached.segment_no == segment_no
1334                    && cached.segment_offset == segment_offset
1335            })
1336            .map(|position| cache.remove(position).1.table)
1337    })
1338}
1339
1340fn cache_commit_segment_table(segment_no: u64, segment_offset: u64, table: Vec<u64>) {
1341    let Ok(context) = memory::active_context_id() else {
1342        return;
1343    };
1344    COMMIT_SEGMENT_CACHE.with(|cache| {
1345        let mut cache = cache.borrow_mut();
1346        if cache.is_empty() {
1347            cache.push((
1348                context,
1349                CommitSegmentCache {
1350                    segment_no,
1351                    segment_offset,
1352                    table,
1353                },
1354            ));
1355            return;
1356        }
1357        if cache.len() == 1 {
1358            let (stored_context, cached) = &mut cache[0];
1359            if *stored_context == context {
1360                cached.segment_no = segment_no;
1361                cached.segment_offset = segment_offset;
1362                cached.table = table;
1363                return;
1364            }
1365        } else if let Some((_, cached)) = cache
1366            .iter_mut()
1367            .find(|(stored_context, _)| *stored_context == context)
1368        {
1369            cached.segment_no = segment_no;
1370            cached.segment_offset = segment_offset;
1371            cached.table = table;
1372            return;
1373        }
1374        cache.push((
1375            context,
1376            CommitSegmentCache {
1377                segment_no,
1378                segment_offset,
1379                table,
1380            },
1381        ));
1382    });
1383}
1384
1385fn read_segment_table_at(offset: u64) -> Result<Vec<u64>, StableMemoryError> {
1386    read_u64_table_at(offset, segment_page_count_usize())
1387}
1388
1389fn write_segmented_tables(entries: &[u64]) -> Result<(u64, u64), StableMemoryError> {
1390    if entries.is_empty() {
1391        return Ok((0, 0));
1392    }
1393    let root_len = segment_count_for_pages(entries_len_u64(entries)?)?;
1394    let mut cursor = append_base()?;
1395    let segment_bytes = root_len
1396        .checked_mul(segment_table_bytes()?)
1397        .ok_or(StableMemoryError::OffsetOverflow)?;
1398    let page_table_bytes = checked_add(segment_bytes, root_table_bytes(root_len)?)?;
1399    memory::ensure_capacity(checked_add(cursor, page_table_bytes)?)?;
1400    let mut root = Vec::with_capacity(
1401        usize::try_from(root_len).map_err(|_| StableMemoryError::OffsetOverflow)?,
1402    );
1403    for segment_no in 0..root_len {
1404        let start = usize::try_from(
1405            segment_no
1406                .checked_mul(SEGMENT_PAGE_COUNT)
1407                .ok_or(StableMemoryError::OffsetOverflow)?,
1408        )
1409        .map_err(|_| StableMemoryError::OffsetOverflow)?;
1410        let mut table = vec![0_u64; segment_page_count_usize()];
1411        for (offset, entry) in entries[start..]
1412            .iter()
1413            .take(segment_page_count_usize())
1414            .enumerate()
1415        {
1416            table[offset] = *entry;
1417        }
1418        root.push(write_segment_table_at(&table, &mut cursor)?);
1419    }
1420    let root_offset = write_root_table_at(&root, &mut cursor)?;
1421    Ok((root_offset, entries_len_u64(&root)?))
1422}
1423
1424#[inline(always)]
1425fn write_segment_table_at(entries: &[u64], cursor: &mut u64) -> Result<u64, StableMemoryError> {
1426    if entries.len() == segment_page_count_usize() {
1427        return write_u64_table_at(entries, cursor);
1428    }
1429
1430    let mut table = vec![0_u64; segment_page_count_usize()];
1431    for (index, entry) in entries.iter().take(segment_page_count_usize()).enumerate() {
1432        table[index] = *entry;
1433    }
1434    write_u64_table_at(&table, cursor)
1435}
1436
1437fn write_root_table_at(entries: &[u64], cursor: &mut u64) -> Result<u64, StableMemoryError> {
1438    write_u64_table_at(entries, cursor)
1439}
1440
1441#[inline(always)]
1442fn write_commit_segment_table_at(
1443    entries: &[u64],
1444    cursor: &mut u64,
1445    profile_enabled: bool,
1446) -> Result<u64, StableMemoryError> {
1447    if profile_enabled {
1448        write_segment_table_at(entries, cursor)
1449    } else {
1450        write_segment_table_at_unmetered(entries, cursor)
1451    }
1452}
1453
1454#[inline(always)]
1455fn write_commit_root_table_at(
1456    entries: &[u64],
1457    cursor: &mut u64,
1458    profile_enabled: bool,
1459) -> Result<u64, StableMemoryError> {
1460    if profile_enabled {
1461        write_root_table_at(entries, cursor)
1462    } else {
1463        write_u64_table_at_unmetered(entries, cursor)
1464    }
1465}
1466
1467fn write_segment_table_at_unmetered(
1468    entries: &[u64],
1469    cursor: &mut u64,
1470) -> Result<u64, StableMemoryError> {
1471    if entries.len() == segment_page_count_usize() {
1472        return write_u64_table_at_unmetered(entries, cursor);
1473    }
1474
1475    let mut table = vec![0_u64; segment_page_count_usize()];
1476    for (index, entry) in entries.iter().take(segment_page_count_usize()).enumerate() {
1477        table[index] = *entry;
1478    }
1479    write_u64_table_at_unmetered(&table, cursor)
1480}
1481
1482fn write_u64_table_at(entries: &[u64], cursor: &mut u64) -> Result<u64, StableMemoryError> {
1483    if entries.is_empty() {
1484        return Ok(0);
1485    }
1486    let offset = *cursor;
1487    let byte_len = entries
1488        .len()
1489        .checked_mul(8)
1490        .ok_or(StableMemoryError::OffsetOverflow)?;
1491    #[cfg(target_endian = "little")]
1492    {
1493        // SAFETY: page-table encoding is little-endian u64 and the target is little-endian.
1494        let bytes = unsafe { std::slice::from_raw_parts(entries.as_ptr().cast::<u8>(), byte_len) };
1495        memory::write_prechecked(offset, bytes)?;
1496        *cursor = checked_add(
1497            offset,
1498            u64::try_from(byte_len).map_err(|_| StableMemoryError::OffsetOverflow)?,
1499        )?;
1500        Ok(offset)
1501    }
1502
1503    #[cfg(not(target_endian = "little"))]
1504    {
1505        let mut bytes = vec![0_u8; byte_len];
1506        for (chunk, entry) in bytes.chunks_exact_mut(8).zip(entries) {
1507            chunk.copy_from_slice(&entry.to_le_bytes());
1508        }
1509        memory::write_prechecked(offset, &bytes)?;
1510        *cursor = checked_add(
1511            offset,
1512            u64::try_from(byte_len).map_err(|_| StableMemoryError::OffsetOverflow)?,
1513        )?;
1514        Ok(offset)
1515    }
1516}
1517
1518fn read_u64_table_at(offset: u64, entries_len: usize) -> Result<Vec<u64>, StableMemoryError> {
1519    if entries_len == 0 {
1520        return Ok(Vec::new());
1521    }
1522    let byte_len = entries_len
1523        .checked_mul(8)
1524        .ok_or(StableMemoryError::OffsetOverflow)?;
1525    #[cfg(target_endian = "little")]
1526    {
1527        let mut entries = Vec::<MaybeUninit<u64>>::with_capacity(entries_len);
1528        unsafe {
1529            entries.set_len(entries_len);
1530        }
1531        // SAFETY: the buffer has `entries_len` u64 slots. The stable-memory read
1532        // fills every byte before conversion to initialized `u64` values.
1533        let bytes =
1534            unsafe { std::slice::from_raw_parts_mut(entries.as_mut_ptr().cast::<u8>(), byte_len) };
1535        memory::read_preallocated(offset, bytes)?;
1536        let ptr = entries.as_mut_ptr().cast::<u64>();
1537        let len = entries.len();
1538        let capacity = entries.capacity();
1539        std::mem::forget(entries);
1540        // SAFETY: all bytes were just initialized by `read_preallocated`, and
1541        // every bit pattern is valid for `u64`.
1542        unsafe { Ok(Vec::from_raw_parts(ptr, len, capacity)) }
1543    }
1544
1545    #[cfg(not(target_endian = "little"))]
1546    {
1547        let mut bytes = vec![0_u8; byte_len];
1548        memory::read_preallocated(offset, &mut bytes)?;
1549        decode_u64_table(&bytes)
1550    }
1551}
1552
1553fn write_u64_table_at_unmetered(
1554    entries: &[u64],
1555    cursor: &mut u64,
1556) -> Result<u64, StableMemoryError> {
1557    if entries.is_empty() {
1558        return Ok(0);
1559    }
1560    let offset = *cursor;
1561    let byte_len = entries
1562        .len()
1563        .checked_mul(8)
1564        .ok_or(StableMemoryError::OffsetOverflow)?;
1565    #[cfg(target_endian = "little")]
1566    {
1567        // SAFETY: page-table encoding is little-endian u64 and the target is little-endian.
1568        let bytes = unsafe { std::slice::from_raw_parts(entries.as_ptr().cast::<u8>(), byte_len) };
1569        memory::write_prechecked_unmetered(offset, bytes)?;
1570        *cursor = checked_add(
1571            offset,
1572            u64::try_from(byte_len).map_err(|_| StableMemoryError::OffsetOverflow)?,
1573        )?;
1574        Ok(offset)
1575    }
1576
1577    #[cfg(not(target_endian = "little"))]
1578    {
1579        let mut bytes = vec![0_u8; byte_len];
1580        for (chunk, entry) in bytes.chunks_exact_mut(8).zip(entries) {
1581            chunk.copy_from_slice(&entry.to_le_bytes());
1582        }
1583        memory::write_prechecked_unmetered(offset, &bytes)?;
1584        *cursor = checked_add(
1585            offset,
1586            u64::try_from(byte_len).map_err(|_| StableMemoryError::OffsetOverflow)?,
1587        )?;
1588        Ok(offset)
1589    }
1590}
1591
1592#[cfg(not(target_endian = "little"))]
1593fn decode_u64_table(bytes: &[u8]) -> Result<Vec<u64>, StableMemoryError> {
1594    if !bytes.len().is_multiple_of(8) {
1595        return Err(StableMemoryError::OffsetOverflow);
1596    }
1597    let mut entries = Vec::with_capacity(bytes.len() / 8);
1598    for chunk in bytes.chunks_exact(8) {
1599        entries.push(u64::from_le_bytes([
1600            chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7],
1601        ]));
1602    }
1603    Ok(entries)
1604}
1605
1606fn imported_page_table(block: &Superblock) -> Result<Vec<u64>, StableMemoryError> {
1607    let count = page_count_for_size(block.import_total_size)?;
1608    let capacity = usize::try_from(count).map_err(|_| StableMemoryError::OffsetOverflow)?;
1609    let mut entries = Vec::with_capacity(capacity);
1610    for page_no in 0..count {
1611        entries.push(checked_add(
1612            block.import_base_offset,
1613            page_no
1614                .checked_mul(page_size())
1615                .ok_or(StableMemoryError::OffsetOverflow)?,
1616        )?);
1617    }
1618    Ok(entries)
1619}
1620
1621fn checksum_logical_range(block: &Superblock, len: u64) -> Result<u64, StableMemoryError> {
1622    let mut offset = 0_u64;
1623    let mut hash = fnv1a64(&[]);
1624    while offset < len {
1625        let chunk_len = (len - offset).min(CHECKSUM_CHUNK_LEN);
1626        let copied_len =
1627            usize::try_from(chunk_len).map_err(|_| StableMemoryError::OffsetOverflow)?;
1628        let mut bytes = vec![0_u8; copied_len];
1629        read_logical_range(block, offset, &mut bytes)?;
1630        hash = fold_fnv1a64(hash, &bytes);
1631        offset += chunk_len;
1632    }
1633    Ok(hash)
1634}
1635
1636fn checksum_physical_range(base_offset: u64, len: u64) -> Result<u64, StableMemoryError> {
1637    let mut offset = 0_u64;
1638    let mut hash = fnv1a64(&[]);
1639    while offset < len {
1640        let chunk_len = (len - offset).min(CHECKSUM_CHUNK_LEN);
1641        let copied_len =
1642            usize::try_from(chunk_len).map_err(|_| StableMemoryError::OffsetOverflow)?;
1643        let mut bytes = vec![0_u8; copied_len];
1644        memory::read_preallocated(checked_add(base_offset, offset)?, &mut bytes)?;
1645        hash = fold_fnv1a64(hash, &bytes);
1646        offset += chunk_len;
1647    }
1648    Ok(hash)
1649}
1650
1651fn clear_import(block: &mut Superblock) -> Result<(), StableMemoryError> {
1652    block.flags &= !FLAG_IMPORTING;
1653    block.import_expected_checksum = 0;
1654    block.import_written_until = 0;
1655    block.import_total_size = 0;
1656    block.import_base_offset = 0;
1657    block.store()?;
1658    invalidate_read_cache();
1659    Ok(())
1660}
1661
1662fn import_offset(block: &Superblock, offset: u64) -> Result<u64, StableMemoryError> {
1663    checked_add(block.import_base_offset, offset)
1664}
1665
1666fn active_page_count(block: &Superblock) -> Result<u64, StableMemoryError> {
1667    page_count_for_size(block.db_size)
1668}
1669
1670fn active_segment_count(block: &Superblock) -> Result<u64, StableMemoryError> {
1671    Ok(block.page_count)
1672}
1673
1674fn read_cache_key(block: &Superblock) -> ReadCacheKey {
1675    ReadCacheKey {
1676        page_table_offset: block.page_table_offset,
1677        page_count: block.page_count,
1678        db_size: block.db_size,
1679        last_tx_id: block.last_tx_id,
1680    }
1681}
1682
1683fn segment_count_for_pages(page_count: u64) -> Result<u64, StableMemoryError> {
1684    Ok(page_count.div_ceil(SEGMENT_PAGE_COUNT))
1685}
1686
1687fn segment_no(page_no: u64) -> u64 {
1688    page_no / SEGMENT_PAGE_COUNT
1689}
1690
1691fn segment_index(page_no: u64) -> Result<usize, StableMemoryError> {
1692    usize::try_from(page_no % SEGMENT_PAGE_COUNT).map_err(|_| StableMemoryError::OffsetOverflow)
1693}
1694
1695fn segment_page_count_usize() -> usize {
1696    usize::try_from(SEGMENT_PAGE_COUNT).expect("segment page count fits usize")
1697}
1698
1699fn segment_table_len() -> usize {
1700    segment_page_count_usize() * 8
1701}
1702
1703fn segment_table_bytes() -> Result<u64, StableMemoryError> {
1704    u64::try_from(segment_table_len()).map_err(|_| StableMemoryError::OffsetOverflow)
1705}
1706
1707fn root_table_bytes(entry_count: u64) -> Result<u64, StableMemoryError> {
1708    entry_count
1709        .checked_mul(PAGE_TABLE_ENTRY_LEN)
1710        .ok_or(StableMemoryError::OffsetOverflow)
1711}
1712
1713fn entries_len_u64<T>(entries: &[T]) -> Result<u64, StableMemoryError> {
1714    u64::try_from(entries.len()).map_err(|_| StableMemoryError::OffsetOverflow)
1715}
1716
1717fn append_base() -> Result<u64, StableMemoryError> {
1718    memory::size_pages()
1719        .checked_mul(STABLE_PAGE_SIZE)
1720        .ok_or(StableMemoryError::OffsetOverflow)
1721}
1722
1723fn page_size() -> u64 {
1724    u64::from(SQLITE_PAGE_SIZE)
1725}
1726
1727fn page_len() -> usize {
1728    usize::try_from(SQLITE_PAGE_SIZE).expect("SQLite page size fits usize")
1729}
1730
1731fn zero_page() -> Vec<u8> {
1732    vec![0_u8; page_len()]
1733}
1734
1735fn checked_add(left: u64, right: u64) -> Result<u64, StableMemoryError> {
1736    left.checked_add(right)
1737        .ok_or(StableMemoryError::OffsetOverflow)
1738}
1739
1740fn fold_fnv1a64(mut hash: u64, bytes: &[u8]) -> u64 {
1741    for byte in bytes {
1742        hash ^= u64::from(*byte);
1743        hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
1744    }
1745    hash
1746}
1747
1748#[cfg(test)]
1749fn hit_failpoint(failpoint: StableBlobFailpoint) -> Result<(), StableMemoryError> {
1750    let Ok(context) = memory::active_context_id() else {
1751        return Ok(());
1752    };
1753    FAILPOINTS.with(|slot| {
1754        let mut slot = slot.borrow_mut();
1755        if slot.get(&context).copied() == Some(failpoint) {
1756            slot.remove(&context);
1757            Err(StableMemoryError::Failpoint(failpoint.name()))
1758        } else {
1759            Ok(())
1760        }
1761    })
1762}
1763
1764#[cfg(not(test))]
1765fn hit_failpoint(_failpoint: StableBlobFailpoint) -> Result<(), StableMemoryError> {
1766    Ok(())
1767}
1768
1769#[cfg(test)]
1770impl StableBlobFailpoint {
1771    fn name(self) -> &'static str {
1772        match self {
1773            Self::OverlayWrite => "before overlay write",
1774            Self::OverlayTruncate => "before overlay truncate",
1775            Self::CommitCapacity => "before commit capacity",
1776            Self::CommitChunkWrite => "before commit page write",
1777            Self::CommitPageTableWrite => "before commit page table write",
1778            Self::CommitSuperblockStore => "before commit superblock store",
1779        }
1780    }
1781}
1782
1783#[cfg(test)]
1784mod tests {
1785    use super::*;
1786    use proptest::prelude::*;
1787    use proptest::test_runner::{Config, TestRunner};
1788    use std::collections::BTreeSet;
1789
1790    #[test]
1791    fn layout_math_matches_expected_boundaries() {
1792        assert_eq!(page_count_for_size(0).unwrap(), 0);
1793        assert_eq!(page_count_for_size(1).unwrap(), 1);
1794        assert_eq!(page_count_for_size(page_size()).unwrap(), 1);
1795        assert_eq!(page_count_for_size(page_size() + 1).unwrap(), 2);
1796
1797        assert_eq!(segment_count_for_pages(0).unwrap(), 0);
1798        assert_eq!(segment_count_for_pages(1).unwrap(), 1);
1799        assert_eq!(segment_count_for_pages(SEGMENT_PAGE_COUNT).unwrap(), 1);
1800        assert_eq!(segment_count_for_pages(SEGMENT_PAGE_COUNT + 1).unwrap(), 2);
1801
1802        assert_eq!(segment_no(SEGMENT_PAGE_COUNT), 1);
1803        assert_eq!(segment_index(SEGMENT_PAGE_COUNT - 1).unwrap(), 255);
1804        assert_eq!(segment_index(SEGMENT_PAGE_COUNT).unwrap(), 0);
1805        assert_eq!(root_table_bytes(2).unwrap(), 16);
1806    }
1807
1808    #[test]
1809    fn layout_math_rejects_u64_max_overflow_boundaries() {
1810        assert!(matches!(
1811            root_table_bytes(u64::MAX),
1812            Err(StableMemoryError::OffsetOverflow)
1813        ));
1814        assert!(matches!(
1815            checked_add(u64::MAX, 1),
1816            Err(StableMemoryError::OffsetOverflow)
1817        ));
1818
1819        let mut block = Superblock::fresh();
1820        block.import_base_offset = u64::MAX;
1821        assert!(matches!(
1822            import_offset(&block, 1),
1823            Err(StableMemoryError::OffsetOverflow)
1824        ));
1825
1826        block.import_base_offset = u64::MAX - page_size() + 1;
1827        block.import_total_size = page_size() + 1;
1828        assert!(matches!(
1829            imported_page_table(&block),
1830            Err(StableMemoryError::OffsetOverflow)
1831        ));
1832    }
1833
1834    #[test]
1835    fn pbt_layout_math_matches_verus_model() {
1836        let mut runner = TestRunner::new(Config {
1837            cases: 512,
1838            ..Config::default()
1839        });
1840
1841        runner
1842            .run(
1843                &(
1844                    boundary_size_strategy(),
1845                    boundary_page_strategy(),
1846                    boundary_entry_strategy(),
1847                ),
1848                |(size, page_no, entries)| {
1849                    let page_count = page_count_for_size(size).unwrap();
1850                    let page_size = u128::from(page_size());
1851                    if size == 0 {
1852                        prop_assert_eq!(page_count, 0);
1853                    } else {
1854                        prop_assert!(u128::from(page_count - 1) * page_size < u128::from(size));
1855                        prop_assert!(u128::from(size) <= u128::from(page_count) * page_size);
1856                    }
1857
1858                    let segment_count = segment_count_for_pages(page_count).unwrap();
1859                    if page_count == 0 {
1860                        prop_assert_eq!(segment_count, 0);
1861                    } else {
1862                        prop_assert!(
1863                            u128::from(segment_count - 1) * u128::from(SEGMENT_PAGE_COUNT)
1864                                < u128::from(page_count)
1865                        );
1866                        prop_assert!(
1867                            u128::from(page_count)
1868                                <= u128::from(segment_count) * u128::from(SEGMENT_PAGE_COUNT)
1869                        );
1870                    }
1871
1872                    let index = segment_index(page_no).unwrap();
1873                    prop_assert!(index < segment_page_count_usize());
1874                    prop_assert_eq!(
1875                        u128::from(segment_no(page_no)) * u128::from(SEGMENT_PAGE_COUNT)
1876                            + index as u128,
1877                        u128::from(page_no)
1878                    );
1879
1880                    match root_table_bytes(entries) {
1881                        Ok(bytes) => prop_assert_eq!(bytes, entries * PAGE_TABLE_ENTRY_LEN),
1882                        Err(StableMemoryError::OffsetOverflow) => {
1883                            prop_assert!(entries.checked_mul(PAGE_TABLE_ENTRY_LEN).is_none());
1884                        }
1885                        Err(error) => return Err(TestCaseError::fail(error.to_string())),
1886                    }
1887                    Ok(())
1888                },
1889            )
1890            .unwrap();
1891    }
1892
1893    fn boundary_size_strategy() -> impl Strategy<Value = u64> {
1894        let page = page_size();
1895        let segment_bytes = SEGMENT_PAGE_COUNT * page;
1896        prop_oneof![
1897            any::<u64>(),
1898            prop::sample::select(boundary_values(&[
1899                0,
1900                1,
1901                page - 1,
1902                page,
1903                page + 1,
1904                segment_bytes - 1,
1905                segment_bytes,
1906                segment_bytes + 1,
1907                u64::MAX,
1908            ])),
1909        ]
1910    }
1911
1912    fn boundary_page_strategy() -> impl Strategy<Value = u64> {
1913        prop_oneof![
1914            any::<u64>(),
1915            prop::sample::select(boundary_values(&[
1916                0,
1917                1,
1918                SEGMENT_PAGE_COUNT - 1,
1919                SEGMENT_PAGE_COUNT,
1920                SEGMENT_PAGE_COUNT + 1,
1921                u64::MAX,
1922            ])),
1923        ]
1924    }
1925
1926    fn boundary_entry_strategy() -> impl Strategy<Value = u64> {
1927        let max_without_overflow = u64::MAX / PAGE_TABLE_ENTRY_LEN;
1928        prop_oneof![
1929            any::<u64>(),
1930            prop::sample::select(boundary_values(&[
1931                0,
1932                1,
1933                SEGMENT_PAGE_COUNT - 1,
1934                SEGMENT_PAGE_COUNT,
1935                SEGMENT_PAGE_COUNT + 1,
1936                max_without_overflow - 1,
1937                max_without_overflow,
1938                max_without_overflow + 1,
1939                u64::MAX - 1,
1940                u64::MAX,
1941            ])),
1942        ]
1943    }
1944
1945    fn boundary_values(values: &[u64]) -> Vec<u64> {
1946        values
1947            .iter()
1948            .flat_map(|value| [value.saturating_sub(1), *value, value.saturating_add(1)])
1949            .collect()
1950    }
1951
1952    #[test]
1953    fn fnv_fold_matches_one_pass_for_multiple_partitions() {
1954        let bytes: Vec<u8> = (0..97)
1955            .map(|index| (index as u8).wrapping_mul(37).wrapping_add(11))
1956            .collect();
1957        let expected = fnv1a64(&bytes);
1958
1959        for split in [0_usize, 1, 2, 7, 31, 64, bytes.len()] {
1960            let split = split.min(bytes.len());
1961            let mut hash = fnv1a64(&[]);
1962            hash = fold_fnv1a64(hash, &bytes[..split]);
1963            hash = fold_fnv1a64(hash, &bytes[split..]);
1964            assert_eq!(hash, expected);
1965        }
1966
1967        let mut hash = fnv1a64(&[]);
1968        for chunk in bytes.chunks(13) {
1969            hash = fold_fnv1a64(hash, chunk);
1970        }
1971        assert_eq!(hash, expected);
1972    }
1973
1974    #[test]
1975    #[serial_test::serial]
1976    fn page_map_commit_tracks_dirty_page_offsets() {
1977        crate::stable::memory::reset_for_tests();
1978        crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
1979        invalidate_read_cache();
1980
1981        let page_zero = vec![1_u8; page_len()];
1982        let page_later = vec![2_u8; page_len()];
1983        let later_page_no = SEGMENT_PAGE_COUNT + 1;
1984        write_at(0, &page_zero).unwrap();
1985        write_at(later_page_no * page_size(), &page_later).unwrap();
1986
1987        let block = Superblock::load().unwrap();
1988        let root = read_root_table(&block).unwrap();
1989        let table = read_page_table(&block).unwrap();
1990        let expected_pages = active_page_count(&block).unwrap();
1991        let expected_segments = segment_count_for_pages(expected_pages).unwrap();
1992
1993        assert_eq!(root.len() as u64, expected_segments);
1994        assert_eq!(table.len() as u64, expected_pages);
1995        assert_ne!(table[0], 0);
1996        assert_ne!(table[later_page_no as usize], 0);
1997
1998        let old_page_zero_offset = table[0];
1999        let updated_page_zero = vec![3_u8; page_len()];
2000        write_at(0, &updated_page_zero).unwrap();
2001        let updated_table = read_page_table(&Superblock::load().unwrap()).unwrap();
2002        let mut out = vec![0_u8; page_len()];
2003        read_base_at(0, &mut out).unwrap();
2004
2005        assert_ne!(updated_table[0], old_page_zero_offset);
2006        assert_eq!(out, updated_page_zero);
2007    }
2008
2009    #[test]
2010    #[serial_test::serial]
2011    fn page_map_commit_tracks_multi_segment_dirty_and_clean_pages() {
2012        crate::stable::memory::reset_for_tests();
2013        crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2014        invalidate_read_cache();
2015
2016        let clean_page_no = 1;
2017        let later_page_no = SEGMENT_PAGE_COUNT + 1;
2018        write_at(0, &vec![1_u8; page_len()]).unwrap();
2019        write_at(clean_page_no * page_size(), &vec![2_u8; page_len()]).unwrap();
2020        write_at(later_page_no * page_size(), &vec![3_u8; page_len()]).unwrap();
2021
2022        let before = Superblock::load().unwrap();
2023        let before_root = read_root_table(&before).unwrap();
2024        let before_table = read_page_table(&before).unwrap();
2025
2026        begin_update().unwrap();
2027        write_at(0, &vec![4_u8; page_len()]).unwrap();
2028        write_at(later_page_no * page_size(), &vec![5_u8; page_len()]).unwrap();
2029        commit_update().unwrap();
2030
2031        let after = Superblock::load().unwrap();
2032        let after_root = read_root_table(&after).unwrap();
2033        let after_table = read_page_table(&after).unwrap();
2034
2035        assert_eq!(after_root.len(), after.page_count as usize);
2036        assert_eq!(after_root.len(), before_root.len());
2037        assert_ne!(after_table[0], before_table[0]);
2038        assert_eq!(
2039            after_table[clean_page_no as usize],
2040            before_table[clean_page_no as usize]
2041        );
2042        assert_ne!(
2043            after_table[later_page_no as usize],
2044            before_table[later_page_no as usize]
2045        );
2046    }
2047
2048    #[test]
2049    #[serial_test::serial]
2050    fn page_map_commit_zeroes_truncated_tail_slots() {
2051        crate::stable::memory::reset_for_tests();
2052        crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2053        invalidate_read_cache();
2054
2055        write_at(0, &vec![1_u8; page_len()]).unwrap();
2056        write_at(page_size(), &vec![2_u8; page_len()]).unwrap();
2057        write_at(2 * page_size(), &vec![3_u8; page_len()]).unwrap();
2058        truncate(page_size()).unwrap();
2059
2060        let block = Superblock::load().unwrap();
2061        let root = read_root_table(&block).unwrap();
2062        let segment = read_segment_table(&block, &root, 0).unwrap();
2063
2064        assert_eq!(block.db_size, page_size());
2065        assert!(segment[0] != 0);
2066        assert_eq!(segment[1], 0);
2067        assert_eq!(segment[2], 0);
2068    }
2069
2070    #[test]
2071    #[serial_test::serial]
2072    fn compact_keeps_zero_pages_and_densifies_offsets_across_segments() {
2073        crate::stable::memory::reset_for_tests();
2074        crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2075        invalidate_read_cache();
2076
2077        let later_page_no = SEGMENT_PAGE_COUNT + 2;
2078        let first_page = vec![7_u8; page_len()];
2079        let later_page = vec![9_u8; page_len()];
2080        write_at(0, &first_page).unwrap();
2081        write_at(later_page_no * page_size(), &later_page).unwrap();
2082
2083        compact().unwrap();
2084
2085        let block = Superblock::load().unwrap();
2086        let root = read_root_table(&block).unwrap();
2087        let table = read_page_table(&block).unwrap();
2088        let mut first_out = vec![0_u8; page_len()];
2089        let mut later_out = vec![0_u8; page_len()];
2090
2091        read_base_at(0, &mut first_out).unwrap();
2092        read_base_at(later_page_no * page_size(), &mut later_out).unwrap();
2093
2094        assert_eq!(root.len() as u64, block.page_count);
2095        assert_eq!(table.len() as u64, active_page_count(&block).unwrap());
2096        assert_ne!(table[0], 0);
2097        assert_eq!(table[1], 0);
2098        assert_eq!(table[later_page_no as usize], table[0] + page_size());
2099        assert_eq!(first_out, first_page);
2100        assert_eq!(later_out, later_page);
2101    }
2102
2103    #[test]
2104    #[serial_test::serial]
2105    fn single_segment_fast_path_preserves_table_after_expand_only_commit() {
2106        crate::stable::memory::reset_for_tests();
2107        crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2108        invalidate_read_cache();
2109
2110        write_at(0, &[0]).unwrap();
2111        truncate(page_size() * 4).unwrap();
2112        truncate(page_size() * 4 + 1).unwrap();
2113
2114        let block = Superblock::load().unwrap();
2115        let table = read_page_table(&block).unwrap();
2116        let mut first = [1_u8; 1];
2117        let mut expanded_tail = [1_u8; 1];
2118
2119        read_base_at(0, &mut first).unwrap();
2120        read_base_at(page_size() * 4, &mut expanded_tail).unwrap();
2121
2122        assert_eq!(block.db_size, page_size() * 4 + 1);
2123        assert_ne!(table[0], 0);
2124        assert_eq!(table[1], 0);
2125        assert_ne!(table[4], 0);
2126        assert_eq!(first, [0]);
2127        assert_eq!(expanded_tail, [0]);
2128    }
2129
2130    #[test]
2131    #[serial_test::serial]
2132    fn page_table_u64_encoding_is_little_endian_and_round_trips() {
2133        crate::stable::memory::reset_for_tests();
2134        crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2135        invalidate_read_cache();
2136
2137        let entries = [
2138            0_u64,
2139            1,
2140            0x0102_0304_0506_0708,
2141            0xf1f2_f3f4_f5f6_f7f8,
2142            u64::MAX,
2143        ];
2144        let mut cursor = 128_u64;
2145        let expected_len = u64::try_from(entries.len() * 8).unwrap();
2146        crate::stable::memory::ensure_capacity(cursor + expected_len).unwrap();
2147
2148        let offset = write_u64_table_at(&entries, &mut cursor).unwrap();
2149        let decoded = read_u64_table_at(offset, entries.len()).unwrap();
2150        let mut encoded = vec![0_u8; entries.len() * 8];
2151        crate::stable::memory::read_preallocated(offset, &mut encoded).unwrap();
2152        let expected = entries
2153            .iter()
2154            .flat_map(|entry| entry.to_le_bytes())
2155            .collect::<Vec<_>>();
2156
2157        assert_eq!(offset, 128);
2158        assert_eq!(cursor, 128 + expected_len);
2159        assert_eq!(decoded, entries);
2160        assert_eq!(encoded, expected);
2161
2162        let mut empty_cursor = cursor;
2163        assert_eq!(write_u64_table_at(&[], &mut empty_cursor).unwrap(), 0);
2164        assert_eq!(empty_cursor, cursor);
2165        assert!(read_u64_table_at(cursor, 0).unwrap().is_empty());
2166    }
2167
2168    #[test]
2169    #[serial_test::serial]
2170    fn pbt_page_table_u64_encoding_round_trips() {
2171        let mut runner = TestRunner::new(Config {
2172            cases: 128,
2173            ..Config::default()
2174        });
2175
2176        runner
2177            .run(
2178                &proptest::collection::vec(any::<u64>(), 0..=512),
2179                |entries| {
2180                    crate::stable::memory::reset_for_tests();
2181                    crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2182                    invalidate_read_cache();
2183
2184                    let mut cursor = 128_u64;
2185                    let byte_len = entries.len().checked_mul(8).unwrap();
2186                    let end = cursor + u64::try_from(byte_len).unwrap();
2187                    crate::stable::memory::ensure_capacity(end).unwrap();
2188
2189                    let offset = write_u64_table_at(&entries, &mut cursor).unwrap();
2190                    let decoded = read_u64_table_at(offset, entries.len()).unwrap();
2191                    prop_assert_eq!(decoded, entries.clone());
2192                    prop_assert_eq!(cursor, end);
2193
2194                    let mut encoded = vec![0_u8; byte_len];
2195                    crate::stable::memory::read_preallocated(offset, &mut encoded).unwrap();
2196                    let expected = entries
2197                        .iter()
2198                        .flat_map(|entry| entry.to_le_bytes())
2199                        .collect::<Vec<_>>();
2200                    prop_assert_eq!(encoded, expected);
2201                    Ok(())
2202                },
2203            )
2204            .unwrap();
2205    }
2206
2207    #[test]
2208    #[serial_test::serial]
2209    fn pbt_compact_preserves_sparse_page_model() {
2210        let mut runner = TestRunner::new(Config {
2211            cases: 32,
2212            ..Config::default()
2213        });
2214
2215        runner
2216            .run(
2217                &proptest::collection::vec(prop::option::of(any::<u8>()), 0..=300),
2218                |pages| {
2219                    crate::stable::memory::reset_for_tests();
2220                    crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2221                    invalidate_read_cache();
2222
2223                    let active_len = pages
2224                        .iter()
2225                        .rposition(Option::is_some)
2226                        .map(|index| index + 1)
2227                        .unwrap_or(0);
2228                    for (page_no, byte) in pages.iter().take(active_len).enumerate() {
2229                        if let Some(byte) = byte {
2230                            write_at(
2231                                u64::try_from(page_no).unwrap() * page_size(),
2232                                &vec![*byte; page_len()],
2233                            )
2234                            .unwrap();
2235                        }
2236                    }
2237
2238                    compact().unwrap();
2239                    let block = Superblock::load().unwrap();
2240                    prop_assert_eq!(
2241                        block.db_size,
2242                        u64::try_from(active_len).unwrap() * page_size()
2243                    );
2244                    let table = read_page_table(&block).unwrap();
2245                    prop_assert_eq!(table.len(), active_len);
2246
2247                    let mut first_compacted_offset = None;
2248                    let mut non_zero_seen = 0_u64;
2249                    for (page_no, byte) in pages.iter().take(active_len).enumerate() {
2250                        let entry = table[page_no];
2251                        let mut page = vec![0_u8; page_len()];
2252                        read_base_at(u64::try_from(page_no).unwrap() * page_size(), &mut page)
2253                            .unwrap();
2254
2255                        if let Some(byte) = byte {
2256                            let base = *first_compacted_offset.get_or_insert(entry);
2257                            prop_assert_ne!(entry, 0);
2258                            prop_assert_eq!(entry, base + non_zero_seen * page_size());
2259                            prop_assert_eq!(page, vec![*byte; page_len()]);
2260                            non_zero_seen += 1;
2261                        } else {
2262                            prop_assert_eq!(entry, 0);
2263                            prop_assert_eq!(page, vec![0_u8; page_len()]);
2264                        }
2265                    }
2266                    Ok(())
2267                },
2268            )
2269            .unwrap();
2270    }
2271
2272    #[derive(Clone, Debug)]
2273    enum BlobOp {
2274        Write { offset: u64, len: usize, byte: u8 },
2275        Truncate { size: u64 },
2276        Compact,
2277    }
2278
2279    #[test]
2280    #[serial_test::serial]
2281    fn pbt_blob_operations_match_logical_model_across_compact() {
2282        let mut runner = TestRunner::new(Config {
2283            cases: 48,
2284            ..Config::default()
2285        });
2286
2287        runner
2288            .run(&blob_operation_sequence(), |operations| {
2289                crate::stable::memory::reset_for_tests();
2290                crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2291                invalidate_read_cache();
2292
2293                let mut model = Vec::new();
2294                let mut materialized = BTreeSet::new();
2295                assert_blob_model(&model, &materialized, false)?;
2296
2297                for operation in operations {
2298                    let compacted = apply_blob_op(operation, &mut model, &mut materialized)?;
2299                    assert_blob_model(&model, &materialized, compacted)?;
2300                }
2301                Ok(())
2302            })
2303            .unwrap();
2304    }
2305
2306    fn blob_operation_sequence() -> impl Strategy<Value = Vec<BlobOp>> {
2307        let write = (blob_offset_strategy(), blob_len_strategy(), any::<u8>())
2308            .prop_map(|(offset, len, byte)| BlobOp::Write { offset, len, byte });
2309        let truncate = blob_offset_strategy().prop_map(|size| BlobOp::Truncate { size });
2310        proptest::collection::vec(prop_oneof![write, truncate, Just(BlobOp::Compact)], 0..=48)
2311    }
2312
2313    fn blob_offset_strategy() -> impl Strategy<Value = u64> {
2314        let limit = blob_model_limit();
2315        let page = page_size();
2316        let segment = SEGMENT_PAGE_COUNT * page;
2317        prop_oneof![
2318            0_u64..=limit,
2319            prop::sample::select(boundary_values(&[
2320                0,
2321                1,
2322                page - 1,
2323                page,
2324                page + 1,
2325                segment - 1,
2326                segment,
2327                segment + 1,
2328                limit - 1,
2329                limit,
2330            ]))
2331            .prop_map(move |value| value.min(limit)),
2332        ]
2333    }
2334
2335    fn blob_len_strategy() -> impl Strategy<Value = usize> {
2336        prop_oneof![
2337            0_usize..=(page_len() * 2 + 17),
2338            prop::sample::select(vec![
2339                0,
2340                1,
2341                page_len() - 1,
2342                page_len(),
2343                page_len() + 1,
2344                page_len() * 2 + 1,
2345            ]),
2346        ]
2347    }
2348
2349    fn blob_model_limit() -> u64 {
2350        (SEGMENT_PAGE_COUNT + 3) * page_size()
2351    }
2352
2353    fn apply_blob_op(
2354        operation: BlobOp,
2355        model: &mut Vec<u8>,
2356        materialized: &mut BTreeSet<u64>,
2357    ) -> Result<bool, TestCaseError> {
2358        match operation {
2359            BlobOp::Write { offset, len, byte } => {
2360                let len = len.min(usize::try_from(blob_model_limit() - offset).unwrap());
2361                let bytes = vec![byte; len];
2362                write_at(offset, &bytes).map_err(|error| TestCaseError::fail(error.to_string()))?;
2363                if len == 0 {
2364                    return Ok(false);
2365                }
2366
2367                let start = usize::try_from(offset).unwrap();
2368                let end = start + len;
2369                if model.len() < start {
2370                    model.resize(start, 0);
2371                }
2372                if model.len() < end {
2373                    model.resize(end, 0);
2374                }
2375                model[start..end].copy_from_slice(&bytes);
2376                mark_materialized_range(offset, len, materialized);
2377                Ok(false)
2378            }
2379            BlobOp::Truncate { size } => {
2380                truncate(size).map_err(|error| TestCaseError::fail(error.to_string()))?;
2381                let new_len = usize::try_from(size).unwrap();
2382                model.resize(new_len, 0);
2383                let active_pages = page_count_for_size(size)
2384                    .map_err(|error| TestCaseError::fail(error.to_string()))?;
2385                materialized.retain(|page_no| *page_no < active_pages);
2386                if size > 0 && !size.is_multiple_of(page_size()) {
2387                    materialized.insert(size / page_size());
2388                }
2389                Ok(false)
2390            }
2391            BlobOp::Compact => {
2392                compact().map_err(|error| TestCaseError::fail(error.to_string()))?;
2393                Ok(true)
2394            }
2395        }
2396    }
2397
2398    fn mark_materialized_range(offset: u64, len: usize, materialized: &mut BTreeSet<u64>) {
2399        let end = offset + u64::try_from(len).unwrap();
2400        let first_page = offset / page_size();
2401        let last_page = (end - 1) / page_size();
2402        for page_no in first_page..=last_page {
2403            materialized.insert(page_no);
2404        }
2405    }
2406
2407    fn assert_blob_model(
2408        model: &[u8],
2409        materialized: &BTreeSet<u64>,
2410        expect_compacted: bool,
2411    ) -> Result<(), TestCaseError> {
2412        let block = Superblock::load().map_err(|error| TestCaseError::fail(error.to_string()))?;
2413        prop_assert_eq!(block.db_size, u64::try_from(model.len()).unwrap());
2414
2415        if !model.is_empty() {
2416            let mut out = vec![0_u8; model.len()];
2417            read_base_at(0, &mut out).map_err(|error| TestCaseError::fail(error.to_string()))?;
2418            prop_assert_eq!(out, model);
2419        }
2420
2421        let mut tail = vec![1_u8; 32];
2422        read_base_at(u64::try_from(model.len()).unwrap(), &mut tail)
2423            .map_err(|error| TestCaseError::fail(error.to_string()))?;
2424        prop_assert_eq!(tail, vec![0_u8; 32]);
2425
2426        let table =
2427            read_page_table(&block).map_err(|error| TestCaseError::fail(error.to_string()))?;
2428        let active_pages = page_count_for_size(u64::try_from(model.len()).unwrap())
2429            .map_err(|error| TestCaseError::fail(error.to_string()))?;
2430        prop_assert_eq!(table.len(), usize::try_from(active_pages).unwrap());
2431
2432        let mut first_compacted_offset = None;
2433        let mut non_zero_seen = 0_u64;
2434        for (index, entry) in table.iter().enumerate() {
2435            let page_no = u64::try_from(index).unwrap();
2436            if materialized.contains(&page_no) {
2437                prop_assert_ne!(*entry, 0);
2438                if expect_compacted {
2439                    let base = *first_compacted_offset.get_or_insert(*entry);
2440                    prop_assert_eq!(*entry, base + non_zero_seen * page_size());
2441                }
2442                non_zero_seen += 1;
2443            } else {
2444                prop_assert_eq!(*entry, 0);
2445            }
2446        }
2447        Ok(())
2448    }
2449
2450    #[test]
2451    #[serial_test::serial]
2452    fn read_metrics_separate_table_cache_from_data_reads() {
2453        crate::stable::memory::reset_for_tests();
2454        crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2455        invalidate_read_cache();
2456
2457        let page = vec![7_u8; page_len()];
2458        write_at(0, &page).unwrap();
2459        invalidate_read_cache();
2460        crate::read_metrics::reset_read_metrics();
2461
2462        let first = read_base_page(0).unwrap();
2463        let second = read_base_page(0).unwrap();
2464        let metrics = crate::read_metrics::read_metrics_snapshot();
2465
2466        assert_eq!(first, page);
2467        assert_eq!(second, page);
2468        assert!(metrics.stable_data_read_calls >= 2);
2469        assert!(metrics.stable_data_read_bytes >= page_size() * 2);
2470        assert!(metrics.page_table_root_misses >= 1);
2471        assert!(metrics.page_table_root_hits >= 1);
2472        assert!(metrics.page_table_segment_misses >= 1);
2473        assert!(metrics.page_table_segment_hits >= 1);
2474        #[cfg(feature = "bench-profile")]
2475        assert!(metrics.superblock_loads <= 1);
2476        #[cfg(not(feature = "bench-profile"))]
2477        assert_eq!(metrics.superblock_loads, 0);
2478    }
2479
2480    #[test]
2481    #[serial_test::serial]
2482    fn page_offset_cache_reuses_page_data_for_small_reads() {
2483        crate::stable::memory::reset_for_tests();
2484        crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2485        invalidate_read_cache();
2486
2487        let page = vec![9_u8; page_len()];
2488        write_at(0, &page).unwrap();
2489        let block = Superblock::load().unwrap();
2490        let mut cache = PageOffsetCache::new();
2491        let mut first = [0_u8; 16];
2492        let mut second = [0_u8; 16];
2493
2494        crate::read_metrics::reset_read_metrics();
2495        read_base_at_with_page_cache(&block, 0, &mut first, &mut cache).unwrap();
2496        read_base_at_with_page_cache(&block, 8, &mut second, &mut cache).unwrap();
2497        let metrics = crate::read_metrics::read_metrics_snapshot();
2498
2499        assert_eq!(first, [9_u8; 16]);
2500        assert_eq!(second, [9_u8; 16]);
2501        assert_eq!(metrics.stable_data_read_calls, 1);
2502        assert_eq!(metrics.stable_data_read_bytes, page_size());
2503    }
2504}