Skip to main content

bitmap/
raw_bitmap.rs

1// Copyright 2026 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use crate::bitmap::{Bitmap, GetResult};
6use crate::storage::Storage;
7use zx_status::Status;
8
9const BITS_PER_WORD: usize = usize::BITS as usize;
10
11const fn last_idx(bitmax: usize) -> usize {
12    (bitmax - 1) / BITS_PER_WORD
13}
14
15const fn first_idx(bitoff: usize) -> usize {
16    bitoff / BITS_PER_WORD
17}
18
19fn get_mask(first: bool, last: bool, off: usize, max: usize) -> usize {
20    let ones = !0usize;
21    let mut mask = ones;
22    if first {
23        mask &= ones << (off % BITS_PER_WORD);
24    }
25    if last {
26        mask &= ones >> ((BITS_PER_WORD - (max % BITS_PER_WORD)) % BITS_PER_WORD);
27    }
28    mask
29}
30
31fn mask_bits(data: usize, idx: usize, bitoff: usize, bitmax: usize, is_set: bool) -> usize {
32    let mask = get_mask(idx == first_idx(bitoff), idx == last_idx(bitmax), bitoff, bitmax);
33    if is_set { !(!mask | data) } else { mask & data }
34}
35
36/// A simple bitmap backed by generic storage.
37#[derive(Debug, Default)]
38pub struct RawBitmapGeneric<S: Storage> {
39    bits: S,
40    size: usize,
41}
42
43impl<S: Storage> RawBitmapGeneric<S> {
44    /// Create a new raw bitmap with the given storage.
45    pub const fn new(bits: S) -> Self {
46        Self { bits, size: 0 }
47    }
48
49    /// Returns the size of this bitmap.
50    pub fn size(&self) -> usize {
51        self.size
52    }
53
54    fn data(&self) -> &[usize] {
55        self.bits.get_data()
56    }
57
58    fn data_mut(&mut self) -> &mut [usize] {
59        self.bits.get_data_mut()
60    }
61
62    /// Access the underlying storage read-only.
63    pub fn storage(&self) -> &S {
64        &self.bits
65    }
66
67    /// Access the underlying storage mutably.
68    pub fn storage_mut(&mut self) -> &mut S {
69        &mut self.bits
70    }
71
72    /// Shrinks the accessible portion of the bitmap, without re-allocating
73    /// the underlying storage.
74    ///
75    /// This is useful for programs which require underlying bitmap storage
76    /// to be aligned to a certain size (initialized via Reset), but want to
77    /// restrict access to a smaller portion of the bitmap (via Shrink).
78    pub fn shrink(&mut self, size: usize) -> Result<(), Status> {
79        if size > self.size {
80            return Err(Status::NO_MEMORY);
81        }
82        self.size = size;
83        Ok(())
84    }
85
86    /// Returns true if all bits in the range `[bitoff, bitmax)` match `is_set`.
87    ///
88    /// If they do not match, returns false and the index of the first bit that doesn't match.
89    /// An empty region (i.e. `bitoff >= bitmax`) will return true.
90    pub fn scan(&self, bitoff: usize, mut bitmax: usize, is_set: bool) -> Option<usize> {
91        bitmax = core::cmp::min(bitmax, self.size);
92        if bitoff >= bitmax {
93            return None;
94        }
95        let data = self.data();
96        let mut i = first_idx(bitoff);
97        let last = last_idx(bitmax);
98        loop {
99            let word = data[i];
100            let masked = mask_bits(word, i, bitoff, bitmax, is_set);
101            if masked != 0 {
102                let first_match = i * BITS_PER_WORD + (masked.trailing_zeros() as usize);
103                return Some(first_match);
104            }
105            if i == last {
106                return None;
107            }
108            i += 1;
109        }
110    }
111
112    /// Returns true if all bits in the range `[bitoff, bitmax)` match `is_set`, scanning in reverse.
113    ///
114    /// If they do not match, returns false and the index of the last bit that doesn't match.
115    /// An empty region (i.e. `bitoff >= bitmax`) will return true.
116    pub fn reverse_scan(&self, bitoff: usize, mut bitmax: usize, is_set: bool) -> Option<usize> {
117        bitmax = core::cmp::min(bitmax, self.size);
118        if bitoff >= bitmax {
119            return None;
120        }
121        let data = self.data();
122        let mut i = last_idx(bitmax);
123        let first = first_idx(bitoff);
124        loop {
125            let word = data[i];
126            let masked = mask_bits(word, i, bitoff, bitmax, is_set);
127            if masked != 0 {
128                let last_match = (i + 1) * BITS_PER_WORD - ((masked.leading_zeros() as usize) + 1);
129                return Some(last_match);
130            }
131            if i == first {
132                return None;
133            }
134            i -= 1;
135        }
136    }
137
138    /// Finds the last run of `run_len` `is_set` bits, in `[bitoff, bitmax)`.
139    ///
140    /// Returns the start of the run, or `Status::NO_RESOURCES` if not found.
141    pub fn reverse_find(
142        &self,
143        is_set: bool,
144        bitoff: usize,
145        bitmax: usize,
146        run_len: usize,
147    ) -> Result<usize, Status> {
148        if bitmax <= bitoff {
149            return Err(Status::INVALID_ARGS);
150        }
151        let mut scan_max = bitmax;
152        loop {
153            let first_mismatch = self.reverse_scan(bitoff, scan_max, !is_set);
154            let start = match first_mismatch {
155                Some(idx) => idx + 1,
156                None => return Err(Status::NO_RESOURCES),
157            };
158            if start - bitoff < run_len {
159                return Err(Status::NO_RESOURCES);
160            }
161            let first_diff = self.reverse_scan(start - run_len, start, is_set);
162            match first_diff {
163                None => return Ok(start - run_len),
164                Some(idx) => {
165                    scan_max = idx;
166                }
167            }
168        }
169    }
170
171    /// Increases the bitmap size.
172    pub fn grow(&mut self, size: usize) -> Result<(), Status> {
173        if !S::SUPPORTS_GROW {
174            return Err(Status::NO_RESOURCES);
175        }
176        if size < self.size {
177            return Err(Status::INVALID_ARGS);
178        } else if size == self.size {
179            return Ok(());
180        }
181        let old_len = if self.size == 0 { 0 } else { last_idx(self.size) + 1 };
182        let new_len = last_idx(size) + 1;
183        let new_bitsize = core::mem::size_of::<usize>() * new_len;
184        self.bits.grow(new_bitsize)?;
185        let data = self.data_mut();
186        for i in old_len..new_len {
187            data[i] = 0;
188        }
189        let old_size = self.size;
190        self.size = size;
191        let clear_limit = core::cmp::min(old_len * BITS_PER_WORD, self.size);
192        if old_size < clear_limit {
193            self.clear(old_size, clear_limit)?;
194        }
195        Ok(())
196    }
197
198    /// Resets the bitmap; clearing and resizing it.
199    ///
200    /// Allocates memory, and can fail.
201    pub fn reset(&mut self, size: usize) -> Result<(), Status> {
202        self.size = size;
203        if size == 0 {
204            return Ok(());
205        }
206        let last = last_idx(size);
207        self.bits.allocate(core::mem::size_of::<usize>() * (last + 1))?;
208        self.clear_all();
209        Ok(())
210    }
211}
212
213impl<S: Storage> Bitmap<usize> for RawBitmapGeneric<S> {
214    fn find(
215        &self,
216        is_set: bool,
217        bitoff: usize,
218        bitmax: usize,
219        run_len: usize,
220    ) -> Result<usize, Status> {
221        if bitmax <= bitoff {
222            return Err(Status::INVALID_ARGS);
223        }
224        let mut start;
225        let mut scan_off = bitoff;
226        loop {
227            let first_mismatch = self.scan(scan_off, bitmax, !is_set);
228            start = match first_mismatch {
229                Some(idx) => idx,
230                None => return Err(Status::NO_RESOURCES),
231            };
232            if bitmax - start < run_len {
233                return Err(Status::NO_RESOURCES);
234            }
235            let first_diff = self.scan(start, start + run_len, is_set);
236            match first_diff {
237                None => return Ok(start),
238                Some(idx) => {
239                    scan_off = idx;
240                }
241            }
242        }
243    }
244
245    fn get(&self, bitoff: usize, bitmax: usize) -> GetResult<usize> {
246        let first_unset = self.scan(bitoff, bitmax, true);
247        GetResult { all_set: first_unset.is_none(), first_unset: first_unset.unwrap_or(bitmax) }
248    }
249
250    fn set(&mut self, bitoff: usize, bitmax: usize) -> Result<(), Status> {
251        if bitoff > bitmax || bitmax > self.size {
252            return Err(Status::INVALID_ARGS);
253        }
254        if bitoff == bitmax {
255            return Ok(());
256        }
257        let first = first_idx(bitoff);
258        let last = last_idx(bitmax);
259        let data = self.data_mut();
260        for i in first..=last {
261            let mask = get_mask(i == first, i == last, bitoff, bitmax);
262            data[i] |= mask;
263        }
264        Ok(())
265    }
266
267    fn clear(&mut self, bitoff: usize, bitmax: usize) -> Result<(), Status> {
268        if bitoff > bitmax || bitmax > self.size {
269            return Err(Status::INVALID_ARGS);
270        }
271        if bitoff == bitmax {
272            return Ok(());
273        }
274        let first = first_idx(bitoff);
275        let last = last_idx(bitmax);
276        let data = self.data_mut();
277        for i in first..=last {
278            let mask = get_mask(i == first, i == last, bitoff, bitmax);
279            data[i] &= !mask;
280        }
281        Ok(())
282    }
283
284    fn clear_all(&mut self) {
285        if self.size == 0 {
286            return;
287        }
288        let last = last_idx(self.size);
289        let data = self.data_mut();
290        for i in 0..=last {
291            data[i] = 0;
292        }
293    }
294}