1use crate::BLOCK_SIZE;
6use anyhow::{Error, ensure};
7use std::borrow::Borrow;
8use std::ops::Range;
9
10const TYPE_MASK: u64 = 0xc0000000_00000000;
11const REGULAR: u64 = 0x00000000_00000000;
12const SPARSE: u64 = 0x80000000_00000000;
13
14const MAX_REGULAR_EXTENT_BLOCKS: u64 = 0x3fff_ffff;
21
22const MAX_SPARSE_EXTENT_BLOCKS: u64 = !TYPE_MASK;
24
25#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct Extent {
30 logical_range: Range<u64>,
31 device_offset: Option<u64>,
32}
33
34impl Extent {
35 pub fn logical_range(&self) -> Range<u64> {
37 self.logical_range.clone()
38 }
39
40 pub fn device_offset(&self) -> Option<u64> {
42 self.device_offset
43 }
44
45 pub fn is_sparse(&self) -> bool {
47 self.device_offset.is_none()
48 }
49
50 pub fn len(&self) -> u64 {
52 self.logical_range.end - self.logical_range.start
53 }
54
55 pub fn new(logical_range: Range<u64>, device_offset: Option<u64>) -> Self {
63 Self::try_new(logical_range, device_offset).unwrap()
64 }
65
66 pub fn try_new(logical_range: Range<u64>, device_offset: Option<u64>) -> Result<Self, Error> {
68 ensure!(
69 logical_range.start % BLOCK_SIZE == 0 && logical_range.end % BLOCK_SIZE == 0,
70 "logical_range boundaries must be a multiple of BLOCK_SIZE (4096 bytes), got {:?}",
71 logical_range
72 );
73 ensure!(
74 logical_range.start <= logical_range.end,
75 "logical_range.start must be <= logical_range.end, got {:?}",
76 logical_range
77 );
78
79 let length_blocks = (logical_range.end - logical_range.start) / BLOCK_SIZE;
80 if device_offset.is_some() {
81 ensure!(
82 length_blocks <= MAX_REGULAR_EXTENT_BLOCKS,
83 "Extent length bounds exceed maximum encodeable length"
84 );
85 } else {
86 ensure!(
87 length_blocks <= MAX_SPARSE_EXTENT_BLOCKS,
88 "Extent length bounds exceed maximum encodeable length"
89 );
90 }
91 Ok(Self { logical_range, device_offset })
92 }
93}
94
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
99struct ExtentEntry {
100 end_logical_offset: u64,
101 device_offset: u64, }
103
104impl ExtentEntry {
105 const SPARSE_DEVICE_OFFSET: u64 = u64::MAX;
106
107 fn is_sparse(&self) -> bool {
108 self.device_offset == Self::SPARSE_DEVICE_OFFSET
109 }
110}
111
112fn encode_regular(length_blocks: u32, target_block: u32) -> u64 {
113 REGULAR | ((length_blocks as u64 & MAX_REGULAR_EXTENT_BLOCKS) << 32) | (target_block as u64)
114}
115
116fn encode_sparse(length_blocks: u64) -> u64 {
117 SPARSE | (length_blocks & !TYPE_MASK)
118}
119
120#[derive(Debug, Clone)]
122pub struct ExtentsIterator<'a> {
123 extents: &'a Extents,
124 index: usize,
125}
126
127impl<'a> Iterator for ExtentsIterator<'a> {
128 type Item = Extent;
129
130 fn next(&mut self) -> Option<Self::Item> {
131 if self.index < self.extents.entries.len() {
132 let res = self.extents.entry_to_result(self.index);
133 self.index += 1;
134 Some(res)
135 } else {
136 None
137 }
138 }
139}
140
141#[derive(Debug, Clone, PartialEq, Eq, Default)]
145pub struct Extents {
146 base_device_offset: u64,
147 entries: Box<[ExtentEntry]>,
148}
149
150impl Extents {
151 pub fn base_device_offset(&self) -> u64 {
153 self.base_device_offset
154 }
155
156 pub fn try_new(
159 extents: impl IntoIterator<Item = impl Borrow<Extent>>,
160 base_device_offset: u64,
161 ) -> Result<Self, Error> {
162 let iter = extents.into_iter();
163 let (lower_bound, _) = iter.size_hint();
164 let mut entries = Vec::with_capacity(lower_bound);
165 let mut current_logical_offset = 0u64;
166
167 for extent in iter {
168 let extent = extent.borrow();
169 ensure!(
170 extent.logical_range.start == current_logical_offset,
171 "Extents must be contiguous and start at 0: expected start \
172 {current_logical_offset}, got {}",
173 extent.logical_range.start
174 );
175 ensure!(
176 extent.logical_range.start < extent.logical_range.end,
177 "Extent logical range must be non-empty, got {:?}",
178 extent.logical_range
179 );
180 ensure!(
181 extent.logical_range.start % BLOCK_SIZE == 0
182 && extent.logical_range.end % BLOCK_SIZE == 0,
183 "logical_range boundaries must be a multiple of BLOCK_SIZE ({BLOCK_SIZE} bytes), \
184 got {:?}",
185 extent.logical_range
186 );
187 let length_blocks = extent.len() / BLOCK_SIZE;
188 if let Some(dev_offset) = extent.device_offset {
189 ensure!(
190 dev_offset >= base_device_offset,
191 "device_offset ({dev_offset}) must be >= base_device_offset \
192 ({base_device_offset})"
193 );
194 let relative_offset = dev_offset - base_device_offset;
195 ensure!(
196 relative_offset % BLOCK_SIZE == 0,
197 "Relative device offset ({dev_offset} - {base_device_offset} = \
198 {relative_offset}) must be a multiple of BLOCK_SIZE ({BLOCK_SIZE} bytes)"
199 );
200 let target_block = relative_offset / BLOCK_SIZE;
201 ensure!(
202 target_block <= u32::MAX as u64,
203 "Relative device offset block index exceeds u32::MAX"
204 );
205 ensure!(
206 length_blocks <= MAX_REGULAR_EXTENT_BLOCKS,
207 "Extent length bounds exceed maximum encodeable length"
208 );
209 current_logical_offset = extent.logical_range.end;
210 entries.push(ExtentEntry {
211 end_logical_offset: current_logical_offset,
212 device_offset: dev_offset,
213 });
214 } else {
215 ensure!(
216 length_blocks <= MAX_SPARSE_EXTENT_BLOCKS,
217 "Extent length bounds exceed maximum encodeable length"
218 );
219 current_logical_offset = extent.logical_range.end;
220 entries.push(ExtentEntry {
221 end_logical_offset: current_logical_offset,
222 device_offset: ExtentEntry::SPARSE_DEVICE_OFFSET,
223 });
224 }
225 }
226
227 Ok(Self { base_device_offset, entries: entries.into_boxed_slice() })
228 }
229
230 pub fn encode(&self) -> impl Iterator<Item = u64> + '_ {
233 let mut prev_logical = 0u64;
234 self.entries.iter().map(move |entry| {
235 let length_blocks = (entry.end_logical_offset - prev_logical) / BLOCK_SIZE;
236 prev_logical = entry.end_logical_offset;
237 if entry.is_sparse() {
238 encode_sparse(length_blocks)
239 } else {
240 let relative_offset = entry.device_offset - self.base_device_offset;
241 let target_block = (relative_offset / BLOCK_SIZE) as u32;
242 encode_regular(length_blocks as u32, target_block)
243 }
244 })
245 }
246
247 pub fn encode_extents(extents: &Extents) -> impl Iterator<Item = u64> + '_ {
249 extents.encode()
250 }
251
252 pub fn encode_extents_with_base_offset(extents: &Extents) -> impl Iterator<Item = u64> + '_ {
255 extents.encode()
256 }
257
258 pub fn from_encoded(
263 encoded: impl IntoIterator<Item = u64>,
264 base_device_offset: u64,
265 ) -> Option<Self> {
266 let iter = encoded.into_iter();
267 let (lower_bound, _) = iter.size_hint();
268 let mut entries = Vec::with_capacity(lower_bound);
269 let mut current_logical_offset = 0u64;
270
271 for val in iter {
272 let kind = val & TYPE_MASK;
273 if kind == REGULAR {
274 let length_blocks = ((val & !TYPE_MASK) >> 32) as u64;
275 let target_block = (val & 0xffff_ffff) as u64;
276 let length_bytes = length_blocks.checked_mul(BLOCK_SIZE)?;
277 current_logical_offset = current_logical_offset.checked_add(length_bytes)?;
278 let device_offset =
279 base_device_offset.checked_add(target_block.checked_mul(BLOCK_SIZE)?)?;
280 entries.push(ExtentEntry {
281 end_logical_offset: current_logical_offset,
282 device_offset,
283 });
284 } else if kind == SPARSE {
285 let length_blocks = (val & !TYPE_MASK) as u64;
286 let length_bytes = length_blocks.checked_mul(BLOCK_SIZE)?;
287 current_logical_offset = current_logical_offset.checked_add(length_bytes)?;
288 entries.push(ExtentEntry {
289 end_logical_offset: current_logical_offset,
290 device_offset: ExtentEntry::SPARSE_DEVICE_OFFSET,
291 });
292 } else {
293 return None;
294 }
295 }
296
297 Some(Self { base_device_offset, entries: entries.into_boxed_slice() })
298 }
299
300 pub fn iter_extents(&self, start_offset: u64) -> ExtentsIterator<'_> {
303 let index = self.entries.partition_point(|e| e.end_logical_offset <= start_offset);
304 ExtentsIterator { extents: self, index }
305 }
306
307 pub fn map(&self, offset: u64) -> Option<Extent> {
314 assert!(
315 offset % BLOCK_SIZE == 0,
316 "offset must be a multiple of BLOCK_SIZE (4096 bytes), got {offset}"
317 );
318 let idx = self.entries.partition_point(|e| e.end_logical_offset <= offset);
319 if idx >= self.entries.len() {
320 return None;
321 }
322
323 let entry = &self.entries[idx];
324 let start_logical = self.entry_start_offset(idx);
325 let end_logical = entry.end_logical_offset;
326
327 let offset_within = offset.checked_sub(start_logical)?;
328 let device_offset = if entry.is_sparse() {
329 None
330 } else {
331 Some(entry.device_offset.checked_add(offset_within)?)
332 };
333
334 Some(Extent { logical_range: offset..end_logical, device_offset })
335 }
336
337 pub fn mappings(&self) -> Vec<Extent> {
339 (0..self.entries.len()).map(|i| self.entry_to_result(i)).collect()
340 }
341
342 fn entry_start_offset(&self, idx: usize) -> u64 {
343 if idx == 0 { 0 } else { self.entries[idx - 1].end_logical_offset }
344 }
345
346 fn entry_to_result(&self, idx: usize) -> Extent {
347 let entry = &self.entries[idx];
348 let start_logical = self.entry_start_offset(idx);
349 let end_logical = entry.end_logical_offset;
350 let device_offset = if entry.is_sparse() { None } else { Some(entry.device_offset) };
351 Extent { logical_range: start_logical..end_logical, device_offset }
352 }
353}
354
355#[cfg(test)]
356mod tests {
357 use super::*;
358
359 #[test]
360 fn test_encode_decode_regular() {
361 let extents = Extents::try_new(
362 [
363 Extent::new(0..(4 * BLOCK_SIZE), Some(10 * BLOCK_SIZE)),
364 Extent::new((4 * BLOCK_SIZE)..(6 * BLOCK_SIZE), Some(100 * BLOCK_SIZE)),
365 ],
366 0,
367 )
368 .unwrap();
369 let encoded = Extents::encode_extents(&extents);
370 let extents_container = Extents::from_encoded(encoded, 0).expect("from_encoded failed");
371
372 let decoded = extents_container.mappings();
373 assert_eq!(decoded.len(), 2);
374
375 assert_eq!(decoded[0].logical_range, 0..(4 * BLOCK_SIZE));
376 assert_eq!(decoded[0].device_offset, Some(10 * BLOCK_SIZE));
377
378 assert_eq!(decoded[1].logical_range, (4 * BLOCK_SIZE)..(6 * BLOCK_SIZE));
379 assert_eq!(decoded[1].device_offset, Some(100 * BLOCK_SIZE));
380 }
381
382 #[test]
383 fn test_encode_decode_sparse() {
384 let extents = Extents::try_new(
385 [
386 Extent::new(0..(2 * BLOCK_SIZE), Some(50 * BLOCK_SIZE)),
387 Extent::new((2 * BLOCK_SIZE)..(5 * BLOCK_SIZE), None),
388 Extent::new((5 * BLOCK_SIZE)..(6 * BLOCK_SIZE), Some(200 * BLOCK_SIZE)),
389 ],
390 0,
391 )
392 .unwrap();
393 let encoded = Extents::encode_extents(&extents);
394 let extents_container = Extents::from_encoded(encoded, 0).expect("from_encoded failed");
395
396 let decoded = extents_container.mappings();
397 assert_eq!(decoded.len(), 3);
398
399 assert_eq!(decoded[0].logical_range, 0..(2 * BLOCK_SIZE));
400 assert_eq!(decoded[0].device_offset, Some(50 * BLOCK_SIZE));
401
402 assert_eq!(decoded[1].logical_range, (2 * BLOCK_SIZE)..(5 * BLOCK_SIZE));
403 assert_eq!(decoded[1].device_offset, None);
404
405 assert_eq!(decoded[2].logical_range, (5 * BLOCK_SIZE)..(6 * BLOCK_SIZE));
406 assert_eq!(decoded[2].device_offset, Some(200 * BLOCK_SIZE));
407 }
408
409 #[test]
410 fn test_encode_decode_non_aligned_start() {
411 let base_device_offset = 17408u64; let extents = Extents::try_new(
413 [
414 Extent::new(0..(4 * BLOCK_SIZE), Some(base_device_offset)),
415 Extent::new(
416 (4 * BLOCK_SIZE)..(6 * BLOCK_SIZE),
417 Some(base_device_offset + 100 * BLOCK_SIZE),
418 ),
419 ],
420 base_device_offset,
421 )
422 .unwrap();
423 let encoded = Extents::encode_extents_with_base_offset(&extents);
424 let extents_container =
425 Extents::from_encoded(encoded, base_device_offset).expect("from_encoded failed");
426
427 let decoded = extents_container.mappings();
428 assert_eq!(decoded.len(), 2);
429 assert_eq!(decoded[0].logical_range, 0..(4 * BLOCK_SIZE));
430 assert_eq!(decoded[0].device_offset, Some(base_device_offset));
431 assert_eq!(decoded[1].logical_range, (4 * BLOCK_SIZE)..(6 * BLOCK_SIZE));
432 assert_eq!(decoded[1].device_offset, Some(base_device_offset + 100 * BLOCK_SIZE));
433
434 let mapped = extents_container.map(0).expect("should map at 0");
435 assert_eq!(mapped.device_offset, Some(base_device_offset));
436
437 let mapped_next = extents_container.map(4 * BLOCK_SIZE).expect("should map next");
438 assert_eq!(mapped_next.device_offset, Some(base_device_offset + 100 * BLOCK_SIZE));
439 }
440
441 #[test]
442 fn test_extents_validation_relative_offset_unaligned_fails() {
443 let base_device_offset = 17408u64;
444 let result = Extents::try_new(
446 [Extent::new(0..BLOCK_SIZE, Some(base_device_offset + 500))],
447 base_device_offset,
448 );
449 assert!(result.is_err());
450 assert!(
451 result.unwrap_err().to_string().contains("Relative device offset"),
452 "Error should mention relative device offset"
453 );
454 }
455
456 #[test]
457 fn test_extents_validation_device_offset_smaller_than_base_fails() {
458 let base_device_offset = 17408u64;
459 let result = Extents::try_new([Extent::new(0..BLOCK_SIZE, Some(0))], base_device_offset);
460 assert!(result.is_err());
461 assert!(
462 result.unwrap_err().to_string().contains("must be >= base_device_offset"),
463 "Error should mention dev_offset >= base_device_offset"
464 );
465 }
466
467 #[test]
468 fn test_extents_validation_non_contiguous_logical_fails() {
469 let result = Extents::try_new(
470 [
471 Extent::new(0..BLOCK_SIZE, Some(0)),
472 Extent::new((2 * BLOCK_SIZE)..(3 * BLOCK_SIZE), Some(BLOCK_SIZE)),
473 ],
474 0,
475 );
476 assert!(result.is_err());
477 assert!(
478 result.unwrap_err().to_string().contains("must be contiguous and start at 0"),
479 "Error should mention non-contiguous start"
480 );
481 }
482
483 #[test]
484 fn test_binary_search_map_logical_offset() {
485 let extents = Extents::try_new(
486 [
487 Extent::new(0..(10 * BLOCK_SIZE), Some(100 * BLOCK_SIZE)),
488 Extent::new((10 * BLOCK_SIZE)..(20 * BLOCK_SIZE), Some(200 * BLOCK_SIZE)),
489 Extent::new((20 * BLOCK_SIZE)..(30 * BLOCK_SIZE), Some(300 * BLOCK_SIZE)),
490 ],
491 0,
492 )
493 .unwrap();
494 let encoded = Extents::encode_extents(&extents);
495 let extents_container = Extents::from_encoded(encoded, 0).expect("from_encoded failed");
496
497 let mapped = extents_container.map(0).expect("should map at offset 0");
498 assert_eq!(mapped.logical_range, 0..(10 * BLOCK_SIZE));
499 assert_eq!(mapped.device_offset, Some(100 * BLOCK_SIZE));
500
501 let mapped_mid = extents_container
502 .map(12 * BLOCK_SIZE)
503 .expect("should map inside second extent via binary search");
504 assert_eq!(mapped_mid.logical_range, (12 * BLOCK_SIZE)..(20 * BLOCK_SIZE));
505 assert_eq!(mapped_mid.device_offset, Some(202 * BLOCK_SIZE));
506 }
507
508 #[test]
509 fn test_map_out_of_bounds() {
510 let extents =
511 Extents::try_new([Extent::new(0..(2 * BLOCK_SIZE), Some(10 * BLOCK_SIZE))], 0).unwrap();
512 let encoded = Extents::encode_extents(&extents);
513 let extents_container = Extents::from_encoded(encoded, 0).expect("from_encoded failed");
514
515 assert!(extents_container.map(2 * BLOCK_SIZE).is_none());
516 assert!(extents_container.map(100 * BLOCK_SIZE).is_none());
517 }
518
519 #[test]
520 fn test_binary_search_iter_extents() {
521 let extents = Extents::try_new(
522 [
523 Extent::new(0..(2 * BLOCK_SIZE), Some(10 * BLOCK_SIZE)),
524 Extent::new((2 * BLOCK_SIZE)..(4 * BLOCK_SIZE), Some(20 * BLOCK_SIZE)),
525 Extent::new((4 * BLOCK_SIZE)..(6 * BLOCK_SIZE), Some(30 * BLOCK_SIZE)),
526 ],
527 0,
528 )
529 .unwrap();
530 let encoded = Extents::encode_extents(&extents);
531 let extents_container = Extents::from_encoded(encoded, 0).expect("from_encoded failed");
532
533 let results: Vec<_> = extents_container.iter_extents(3 * BLOCK_SIZE).collect();
534 assert_eq!(results.len(), 2);
535 assert_eq!(results[0].logical_range, (2 * BLOCK_SIZE)..(4 * BLOCK_SIZE));
536 assert_eq!(results[1].logical_range, (4 * BLOCK_SIZE)..(6 * BLOCK_SIZE));
537 }
538
539 #[test]
540 fn test_exact_boundary_queries() {
541 let extents = Extents::try_new(
542 [
543 Extent::new(0..(10 * BLOCK_SIZE), Some(100 * BLOCK_SIZE)),
544 Extent::new((10 * BLOCK_SIZE)..(20 * BLOCK_SIZE), Some(200 * BLOCK_SIZE)),
545 ],
546 0,
547 )
548 .unwrap();
549 let encoded = Extents::encode_extents(&extents);
550 let extents_container = Extents::from_encoded(encoded, 0).expect("from_encoded failed");
551
552 let mapped = extents_container.map(10 * BLOCK_SIZE).expect("should map at exact boundary");
553 assert_eq!(mapped.logical_range, (10 * BLOCK_SIZE)..(20 * BLOCK_SIZE));
554 assert_eq!(mapped.device_offset, Some(200 * BLOCK_SIZE));
555
556 let results: Vec<_> = extents_container.iter_extents(10 * BLOCK_SIZE).collect();
557 assert_eq!(results.len(), 1);
558 assert_eq!(results[0].logical_range, (10 * BLOCK_SIZE)..(20 * BLOCK_SIZE));
559 }
560
561 #[test]
562 #[should_panic(expected = "multiple of BLOCK_SIZE")]
563 fn test_extent_new_unaligned_logical_start_panics() {
564 Extent::new(1..(2 * BLOCK_SIZE), Some(10 * BLOCK_SIZE));
565 }
566
567 #[test]
568 #[should_panic(expected = "multiple of BLOCK_SIZE")]
569 fn test_extent_new_unaligned_logical_end_panics() {
570 Extent::new(0..(2 * BLOCK_SIZE + 1), Some(10 * BLOCK_SIZE));
571 }
572
573 #[test]
574 #[should_panic(expected = "multiple of BLOCK_SIZE")]
575 fn test_map_unaligned_offset_panics() {
576 Extents::default().map(500);
577 }
578
579 #[test]
580 fn test_iter_extents_unaligned_start_offset() {
581 let extents = Extents::try_new(
582 [
583 Extent::new(0..(10 * BLOCK_SIZE), Some(100 * BLOCK_SIZE)),
584 Extent::new((10 * BLOCK_SIZE)..(20 * BLOCK_SIZE), Some(200 * BLOCK_SIZE)),
585 ],
586 0,
587 )
588 .unwrap();
589 let encoded = Extents::encode_extents(&extents);
590 let extents_container = Extents::from_encoded(encoded, 0).expect("from_encoded failed");
591 let results: Vec<_> = extents_container.iter_extents(500).collect();
592 assert_eq!(results.len(), 2);
593 assert_eq!(results[0].logical_range, 0..(10 * BLOCK_SIZE));
594 }
595
596 #[test]
597 fn test_encode_extents_regular_length_overflow_errors() {
598 let result = Extent::try_new(
599 0..((MAX_REGULAR_EXTENT_BLOCKS + 1) * BLOCK_SIZE),
600 Some(10 * BLOCK_SIZE),
601 );
602 assert!(result.is_err());
603 assert_eq!(
604 result.unwrap_err().to_string(),
605 "Extent length bounds exceed maximum encodeable length"
606 );
607 }
608
609 #[test]
610 fn test_from_encoded_unknown_kind_returns_none() {
611 let unknown_descriptor = 0x40000000_00000000;
612 assert!(Extents::from_encoded([unknown_descriptor], 0).is_none());
613 }
614}