Skip to main content

ic_sqlite_vfs/stable/
memory_layout.rs

1//! Shared constants and helpers for the MemoryManager-compatible layout.
2//!
3//! Kept separate so the forked manager stays small and the byte format remains
4//! visible in one place.
5
6use crate::config::STABLE_PAGE_SIZE;
7use crate::stable::raw_memory::Memory;
8use std::cell::Cell;
9
10pub(super) const MAGIC: &[u8; 3] = b"MGR";
11pub(super) const LAYOUT_VERSION: u8 = 1;
12pub(super) const MAX_NUM_MEMORIES: u8 = 255;
13pub(super) const MAX_NUM_BUCKETS: u64 = 32_768;
14pub(super) const BUCKET_SIZE_IN_PAGES: u64 = 128;
15pub(super) const UNALLOCATED_BUCKET_MARKER: u8 = MAX_NUM_MEMORIES;
16pub(super) const BUCKETS_OFFSET_IN_PAGES: u64 = 1;
17pub(super) const BUCKETS_OFFSET_IN_BYTES: u64 = BUCKETS_OFFSET_IN_PAGES * STABLE_PAGE_SIZE;
18pub(super) const HEADER_RESERVED_BYTES: usize = 32;
19pub(super) const HEADER_SIZE: u64 = 3 + 1 + 2 + 2 + HEADER_RESERVED_BYTES as u64 + 255 * 8;
20
21#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
22pub struct MemoryId(pub(super) u8);
23
24impl MemoryId {
25    pub const fn new(id: u8) -> Self {
26        assert!(id != UNALLOCATED_BUCKET_MARKER);
27        Self(id)
28    }
29}
30
31#[derive(Clone, Copy)]
32pub(super) struct BucketId(pub(super) u16);
33
34#[derive(Clone, Copy)]
35pub(super) struct VirtualSegment {
36    address: u64,
37    length: u64,
38}
39
40impl VirtualSegment {
41    pub(super) fn new(address: u64, length: u64) -> Self {
42        Self { address, length }
43    }
44
45    fn contains(self, other: Self) -> bool {
46        virtual_segment_contains(self.address, self.length, other.address, other.length)
47    }
48
49    fn address(self) -> u64 {
50        self.address
51    }
52}
53
54#[derive(Clone)]
55pub(super) struct BucketCache {
56    bucket: Cell<VirtualSegment>,
57    real_address: Cell<u64>,
58}
59
60impl BucketCache {
61    pub(super) fn new() -> Self {
62        Self {
63            bucket: Cell::new(VirtualSegment::new(0, 0)),
64            real_address: Cell::new(0),
65        }
66    }
67
68    pub(super) fn get(&self, segment: VirtualSegment) -> Option<u64> {
69        let bucket = self.bucket.get();
70        bucket
71            .contains(segment)
72            .then(|| self.real_address.get() + (segment.address() - bucket.address()))
73    }
74
75    pub(super) fn store(&self, bucket: VirtualSegment, real_address: u64) {
76        self.bucket.set(bucket);
77        self.real_address.set(real_address);
78    }
79}
80
81pub(super) fn bucket_allocations_address(id: BucketId) -> u64 {
82    HEADER_SIZE + u64::from(id.0)
83}
84
85pub(super) fn virtual_segment_contains(
86    address: u64,
87    length: u64,
88    other_address: u64,
89    other_length: u64,
90) -> bool {
91    if length > u64::MAX - address {
92        return false;
93    }
94    if other_length > u64::MAX - other_address {
95        return false;
96    }
97    let end = address + length;
98    let other_end = other_address + other_length;
99    address <= other_address && other_end <= end
100}
101
102pub(super) fn write_growing<M: Memory>(memory: &M, offset: u64, bytes: &[u8]) {
103    let end = offset
104        .checked_add(bytes.len() as u64)
105        .expect("offset overflow");
106    let size_bytes = memory
107        .size()
108        .checked_mul(STABLE_PAGE_SIZE)
109        .expect("memory size overflow");
110    if end > size_bytes {
111        let pages = (end - size_bytes).div_ceil(STABLE_PAGE_SIZE);
112        assert!(memory.grow(pages) >= 0, "stable memory grow failed");
113    }
114    memory.write(offset, bytes);
115}
116
117pub(super) fn read_u64(bytes: &[u8]) -> u64 {
118    u64::from_le_bytes(bytes.try_into().expect("u64 field has 8 bytes"))
119}
120
121#[cfg(test)]
122mod tests {
123    use super::{virtual_segment_contains, BucketCache, VirtualSegment};
124
125    #[test]
126    fn virtual_segment_contains_rejects_overflowing_ranges() {
127        let segment = VirtualSegment::new(100, 50);
128
129        assert!(segment.contains(VirtualSegment::new(120, 10)));
130        assert!(!segment.contains(VirtualSegment::new(u64::MAX - 4, 8)));
131        assert!(
132            !VirtualSegment::new(u64::MAX - 4, 8).contains(VirtualSegment::new(u64::MAX - 3, 1))
133        );
134        assert_eq!(
135            segment.contains(VirtualSegment::new(120, 10)),
136            virtual_segment_contains(100, 50, 120, 10)
137        );
138    }
139
140    #[test]
141    fn bucket_cache_rejects_overflowing_segment_hit() {
142        let cache = BucketCache::new();
143
144        cache.store(VirtualSegment::new(0, 128), 1_000);
145
146        assert_eq!(cache.get(VirtualSegment::new(8, 8)), Some(1_008));
147        assert_eq!(cache.get(VirtualSegment::new(u64::MAX - 4, 8)), None);
148    }
149}