ic_sqlite_vfs/stable/
memory_layout.rs1use 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}