disk_layered/
bitmap.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Sector bitmaps for tracking read sectors during layered read operations.
5//!
6//! FUTURE: use real bitmaps instead of bool vectors.
7
8use std::ops::Range;
9
10pub(crate) struct Bitmap {
11    bits: Vec<bool>,
12    sector: u64,
13}
14
15impl Bitmap {
16    pub fn new(sector: u64, len: usize) -> Self {
17        Self {
18            bits: vec![false; len],
19            sector,
20        }
21    }
22
23    pub fn unset_iter(&mut self) -> impl Iterator<Item = SectorBitmapRange<'_>> {
24        let mut n = 0;
25        let sector = self.sector;
26        self.bits
27            .chunk_by_mut(|&a, &b| a == b)
28            .filter_map(move |bits| {
29                let start = n;
30                n += bits.len();
31                if bits.first().is_some_and(|&x| !x) {
32                    Some(SectorBitmapRange {
33                        bits,
34                        start_sector: sector + start as u64,
35                        sector_within_bitmap: start,
36                        set_count: 0,
37                    })
38                } else {
39                    None
40                }
41            })
42    }
43}
44
45pub(crate) struct SectorBitmapRange<'a> {
46    bits: &'a mut [bool],
47    start_sector: u64,
48    sector_within_bitmap: usize,
49    set_count: usize,
50}
51
52impl SectorBitmapRange<'_> {
53    pub fn view(&mut self, len: u64) -> SectorMarker<'_> {
54        SectorMarker {
55            bits: &mut self.bits[..len as usize],
56            set_count: &mut self.set_count,
57            sector_base: self.start_sector,
58        }
59    }
60
61    pub fn start_sector_within_bitmap(&self) -> usize {
62        self.sector_within_bitmap
63    }
64
65    pub fn start_sector(&self) -> u64 {
66        self.start_sector
67    }
68
69    pub fn end_sector(&self) -> u64 {
70        self.start_sector + self.bits.len() as u64
71    }
72
73    pub fn len(&self) -> u64 {
74        self.bits.len() as u64
75    }
76
77    pub fn set_count(&self) -> usize {
78        self.set_count
79    }
80
81    pub(crate) fn unset_iter(&self) -> impl '_ + Iterator<Item = Range<u64>> {
82        let mut n = self.start_sector;
83        self.bits.chunk_by(|&a, &b| a == b).filter_map(move |bits| {
84            let start = n;
85            n += bits.len() as u64;
86            if bits.first().is_some_and(|&x| !x) {
87                Some(start..n)
88            } else {
89                None
90            }
91        })
92    }
93}
94
95/// A type to mark sectors that have been read by a layer as part of a
96/// [`LayerIo::read`](super::LayerIo::read) operation.
97pub struct SectorMarker<'a> {
98    bits: &'a mut [bool],
99    sector_base: u64,
100    set_count: &'a mut usize,
101}
102
103impl SectorMarker<'_> {
104    #[track_caller]
105    fn sector_to_index(&self, sector: u64) -> usize {
106        let i = sector
107            .checked_sub(self.sector_base)
108            .expect("invalid sector");
109        assert!(i < self.bits.len() as u64, "invalid sector");
110        i as usize
111    }
112
113    /// Mark the specified sector number as having been read.
114    #[track_caller]
115    pub fn set(&mut self, sector: u64) {
116        let i = self.sector_to_index(sector);
117        *self.set_count += !self.bits[i] as usize;
118        self.bits[i] = true;
119    }
120
121    /// Mark the range of sectors as having been read.
122    #[track_caller]
123    pub fn set_range(&mut self, range: Range<u64>) {
124        for sector in range {
125            self.set(sector);
126        }
127    }
128
129    /// Mark all the sectors as having been read.
130    pub fn set_all(&mut self) {
131        self.set_range(self.sector_base..self.sector_base + self.bits.len() as u64);
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::Bitmap;
138
139    #[test]
140    fn test_bitmap() {
141        let base = 0x492893;
142        let mut bitmap = Bitmap::new(base, 10);
143        {
144            let mut iter = bitmap.unset_iter();
145            let mut range = iter.next().unwrap();
146            assert!(iter.next().is_none());
147            assert_eq!(range.start_sector(), base);
148            assert_eq!(range.end_sector(), base + 10);
149            assert_eq!(range.len(), 10);
150            assert_eq!(range.set_count(), 0);
151            range.view(6).set_all();
152            assert_eq!(range.set_count(), 6);
153        }
154        {
155            let mut iter = bitmap.unset_iter();
156            let mut range = iter.next().unwrap();
157            assert!(iter.next().is_none());
158            assert_eq!(range.start_sector(), base + 6);
159            assert_eq!(range.end_sector(), base + 10);
160            assert_eq!(range.start_sector_within_bitmap(), 6);
161            assert_eq!(range.len(), 4);
162            range.view(4).set(base + 7);
163            assert_eq!(range.set_count(), 1);
164        }
165        {
166            let mut iter = bitmap.unset_iter();
167            let range = iter.next().unwrap();
168            let range2 = iter.next().unwrap();
169            assert!(iter.next().is_none());
170            assert_eq!(range.start_sector(), base + 6);
171            assert_eq!(range.end_sector(), base + 7);
172            assert_eq!(range2.start_sector(), base + 8);
173            assert_eq!(range2.end_sector(), base + 10);
174        }
175    }
176}