lewis.moe
Nothing to see here, move along
1use crate::block_io::BlockIo;
2use crate::cache::BlockCache;
3use crate::error::FsError;
4use lancer_core::fs::{BTreeNode, BlockRef, FREEMAP_BITS_PER_LEAF, FreemapLeaf, Superblock};
5
6pub struct FreemapAllocator {
7 hint_leaf_key: u64,
8 hint_bit: u32,
9 total_data_blocks: u64,
10 data_region_start: u64,
11 freemap_root: BlockRef,
12}
13
14impl FreemapAllocator {
15 pub fn init(sb: &Superblock) -> Self {
16 let reserved_blocks = compute_reserved_count(sb.total_blocks);
17 let data_region_start = 2 + reserved_blocks;
18 let raw_data_blocks = sb.total_blocks.saturating_sub(data_region_start);
19 let ditto_size = raw_data_blocks / crate::ditto::DITTO_FRACTION;
20 let total_data_blocks = raw_data_blocks.saturating_sub(ditto_size);
21
22 Self {
23 hint_leaf_key: 0,
24 hint_bit: 0,
25 total_data_blocks,
26 data_region_start,
27 freemap_root: sb.freemap_root,
28 }
29 }
30
31 #[cfg(test)]
32 pub fn init_empty() -> Self {
33 Self {
34 hint_leaf_key: 0,
35 hint_bit: 0,
36 total_data_blocks: 0,
37 data_region_start: 2,
38 freemap_root: BlockRef::ZERO,
39 }
40 }
41
42 pub fn set_root(&mut self, root: BlockRef) {
43 self.freemap_root = root;
44 }
45
46 pub fn root(&self) -> BlockRef {
47 self.freemap_root
48 }
49
50 pub fn data_region_start(&self) -> u64 {
51 self.data_region_start
52 }
53
54 pub fn total_data_blocks(&self) -> u64 {
55 self.total_data_blocks
56 }
57
58 pub fn count_allocated_blocks(
59 &self,
60 cache: &mut BlockCache,
61 bio: &mut BlockIo,
62 ) -> Result<u64, FsError> {
63 let total_leaves = self.total_leaves();
64 (0..total_leaves).try_fold(0u64, |allocated, leaf_key| {
65 let leaf_block = find_freemap_leaf(cache, bio, &self.freemap_root, leaf_key)?;
66 let leaf_data = cache.cache_read(bio, leaf_block)?;
67 let leaf = unsafe { &*(leaf_data.as_ptr() as *const FreemapLeaf) };
68 let leaf_base = leaf_key * FREEMAP_BITS_PER_LEAF as u64;
69 let remaining = self.total_data_blocks.saturating_sub(leaf_base);
70 let bits_in_leaf = remaining.min(FREEMAP_BITS_PER_LEAF as u64);
71 let full_bytes = bits_in_leaf / 8;
72 let extra_bits = bits_in_leaf % 8;
73 let byte_count = leaf.bits[..full_bytes as usize]
74 .iter()
75 .fold(0u64, |acc, &b| acc + b.count_ones() as u64);
76 let tail_count = match extra_bits {
77 0 => 0u64,
78 n => {
79 let mask = (1u8 << n) - 1;
80 (leaf.bits[full_bytes as usize] & mask).count_ones() as u64
81 }
82 };
83 Ok(allocated + byte_count + tail_count)
84 })
85 }
86
87 fn total_leaves(&self) -> u64 {
88 self.total_data_blocks
89 .div_ceil(FREEMAP_BITS_PER_LEAF as u64)
90 }
91
92 pub fn alloc_blocks(
93 &mut self,
94 cache: &mut BlockCache,
95 bio: &mut BlockIo,
96 count: u32,
97 ) -> Result<u64, FsError> {
98 let total_leaves = self.total_leaves();
99 if total_leaves == 0 {
100 return Err(FsError::FreemapExhausted);
101 }
102
103 let start_key = self.hint_leaf_key;
104 let hint_bit = self.hint_bit;
105
106 let result = leaf_key_sequence(start_key, total_leaves).try_fold((), |(), leaf_key| {
107 let leaf_block = find_freemap_leaf(cache, bio, &self.freemap_root, leaf_key)
108 .map_err(AllocSignal::Fs)?;
109 let leaf_data = cache.cache_read(bio, leaf_block).map_err(AllocSignal::Fs)?;
110 let leaf = unsafe { &*(leaf_data.as_ptr() as *const FreemapLeaf) };
111
112 let bits_in_leaf = {
113 let leaf_base_block = leaf_key * FREEMAP_BITS_PER_LEAF as u64;
114 let remaining = self.total_data_blocks.saturating_sub(leaf_base_block);
115 remaining.min(FREEMAP_BITS_PER_LEAF as u64) as usize
116 };
117
118 let search_from = match leaf_key == start_key {
119 true => hint_bit as usize,
120 false => 0,
121 };
122
123 match leaf.find_contiguous_free_from(search_from, count as usize) {
124 Some(bit_offset) if bit_offset + count as usize <= bits_in_leaf => {
125 Err(AllocSignal::Found(leaf_key, leaf_block, bit_offset))
126 }
127 _ => Ok(()),
128 }
129 });
130
131 match result {
132 Err(AllocSignal::Found(leaf_key, leaf_block, bit_offset)) => {
133 let leaf_data_mut = cache.cache_get_mut(bio, leaf_block)?;
134 let leaf_mut = unsafe { &mut *(leaf_data_mut.as_mut_ptr() as *mut FreemapLeaf) };
135 leaf_mut.set_range(bit_offset, count as usize);
136
137 self.hint_leaf_key = leaf_key;
138 self.hint_bit = (bit_offset + count as usize) as u32;
139 if self.hint_bit as usize >= FREEMAP_BITS_PER_LEAF {
140 self.hint_leaf_key += 1;
141 self.hint_bit = 0;
142 }
143
144 Ok(self.data_region_start
145 + leaf_key * FREEMAP_BITS_PER_LEAF as u64
146 + bit_offset as u64)
147 }
148 Err(AllocSignal::Fs(e)) => Err(e),
149 Ok(()) => Err(FsError::FreemapExhausted),
150 }
151 }
152
153 pub fn free_blocks(
154 &mut self,
155 cache: &mut BlockCache,
156 bio: &mut BlockIo,
157 start: u64,
158 count: u32,
159 ) -> Result<(), FsError> {
160 if start < self.data_region_start {
161 return Err(FsError::InvalidBlock);
162 }
163 let relative = start - self.data_region_start;
164 let leaf_key = relative / FREEMAP_BITS_PER_LEAF as u64;
165 let bit_offset = (relative % FREEMAP_BITS_PER_LEAF as u64) as usize;
166
167 match bit_offset + count as usize <= FREEMAP_BITS_PER_LEAF {
168 true => {}
169 false => return Err(FsError::InvalidBlock),
170 }
171
172 let leaf_block = find_freemap_leaf(cache, bio, &self.freemap_root, leaf_key)?;
173 let leaf_data = cache.cache_get_mut(bio, leaf_block)?;
174 let leaf = unsafe { &mut *(leaf_data.as_mut_ptr() as *mut FreemapLeaf) };
175
176 match leaf.is_range_allocated(bit_offset, count as usize) {
177 true => {
178 leaf.clear_range(bit_offset, count as usize);
179 Ok(())
180 }
181 false => Err(FsError::DoubleFree),
182 }
183 }
184
185 pub fn alloc_variable(
186 &mut self,
187 cache: &mut BlockCache,
188 bio: &mut BlockIo,
189 size_log2: u8,
190 ) -> Result<u64, FsError> {
191 let block_count = 1u32 << (size_log2.saturating_sub(12));
192 self.alloc_blocks(cache, bio, block_count)
193 }
194}
195
196fn leaf_key_sequence(start: u64, total_leaves: u64) -> impl Iterator<Item = u64> {
197 let mut offset = 0u64;
198 core::iter::from_fn(move || match offset < total_leaves {
199 true => {
200 let key = (start + offset) % total_leaves;
201 offset += 1;
202 Some(key)
203 }
204 false => None,
205 })
206}
207
208enum AllocSignal {
209 Found(u64, u64, usize),
210 Fs(FsError),
211}
212
213const MAX_FREEMAP_DEPTH: usize = 6;
214
215enum FreemapWalk {
216 FoundLeaf(u64),
217 Err(FsError),
218}
219
220impl From<FsError> for FreemapWalk {
221 fn from(e: FsError) -> Self {
222 FreemapWalk::Err(e)
223 }
224}
225
226pub fn find_freemap_leaf_pub(
227 cache: &mut BlockCache,
228 bio: &mut BlockIo,
229 root: &BlockRef,
230 leaf_key: u64,
231) -> Result<u64, FsError> {
232 find_freemap_leaf(cache, bio, root, leaf_key)
233}
234
235fn find_freemap_leaf(
236 cache: &mut BlockCache,
237 bio: &mut BlockIo,
238 root: &BlockRef,
239 leaf_key: u64,
240) -> Result<u64, FsError> {
241 if root.is_null() {
242 return Err(FsError::NotFound);
243 }
244
245 let start = crate::blockref_block_num(root);
246 let result = (0..MAX_FREEMAP_DEPTH).try_fold(start, |current_block, _| {
247 let node_data = cache
248 .cache_read(bio, current_block)
249 .map_err(FreemapWalk::Err)?;
250 let node = unsafe { &*(node_data.as_ptr() as *const BTreeNode) };
251
252 match node.is_valid() {
253 true => {}
254 false => return Err(FreemapWalk::Err(FsError::CorruptNode)),
255 }
256
257 let child_idx = node.find_child_index(leaf_key);
258 match child_idx < node.entry_count() {
259 true => match node.is_leaf() {
260 true => Err(FreemapWalk::FoundLeaf(crate::blockref_block_num(
261 &node.entries[child_idx],
262 ))),
263 false => Ok(crate::blockref_block_num(&node.entries[child_idx])),
264 },
265 false => Err(FreemapWalk::Err(FsError::NotFound)),
266 }
267 });
268
269 match result {
270 Ok(_) => Err(FsError::CorruptNode),
271 Err(FreemapWalk::FoundLeaf(block)) => Ok(block),
272 Err(FreemapWalk::Err(e)) => Err(e),
273 }
274}
275
276pub(crate) fn compute_reserved_count(total_blocks: u64) -> u64 {
277 let data_blocks = total_blocks.saturating_sub(2);
278 let leaves_needed = data_blocks.div_ceil(FREEMAP_BITS_PER_LEAF as u64);
279
280 match leaves_needed {
281 0 => 0,
282 _ => {
283 let level0_nodes = leaves_needed.div_ceil(63);
284
285 let (higher_nodes, higher_depth) =
286 core::iter::successors(Some(level0_nodes), |&level| match level <= 1 {
287 true => None,
288 false => Some(level.div_ceil(63)),
289 })
290 .skip(1)
291 .fold((0u64, 0u64), |(total, d), level_count| {
292 (total + level_count, d + 1)
293 });
294
295 let btree_nodes = level0_nodes + higher_nodes;
296 let tree_depth = higher_depth + 1;
297 let cow_headroom = (tree_depth + 1) * 2;
298 leaves_needed + btree_nodes + cow_headroom
299 }
300 }
301}