1use 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
95pub 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 #[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 #[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 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}