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        self.data_mut()[old_len..new_len].fill(0);
186        let old_size = self.size;
187        self.size = size;
188        let clear_limit = core::cmp::min(old_len * BITS_PER_WORD, self.size);
189        if old_size < clear_limit {
190            self.clear(old_size, clear_limit)?;
191        }
192        Ok(())
193    }
194
195    /// Resets the bitmap; clearing and resizing it.
196    ///
197    /// Allocates memory, and can fail.
198    pub fn reset(&mut self, size: usize) -> Result<(), Status> {
199        self.size = size;
200        if size == 0 {
201            return Ok(());
202        }
203        let last = last_idx(size);
204        self.bits.allocate(core::mem::size_of::<usize>() * (last + 1))?;
205        self.clear_all();
206        Ok(())
207    }
208}
209
210impl<S: Storage> Bitmap<usize> for RawBitmapGeneric<S> {
211    fn find(
212        &self,
213        is_set: bool,
214        bitoff: usize,
215        bitmax: usize,
216        run_len: usize,
217    ) -> Result<usize, Status> {
218        if bitmax <= bitoff {
219            return Err(Status::INVALID_ARGS);
220        }
221        let mut start;
222        let mut scan_off = bitoff;
223        loop {
224            let first_mismatch = self.scan(scan_off, bitmax, !is_set);
225            start = match first_mismatch {
226                Some(idx) => idx,
227                None => return Err(Status::NO_RESOURCES),
228            };
229            if bitmax - start < run_len {
230                return Err(Status::NO_RESOURCES);
231            }
232            let first_diff = self.scan(start, start + run_len, is_set);
233            match first_diff {
234                None => return Ok(start),
235                Some(idx) => {
236                    scan_off = idx;
237                }
238            }
239        }
240    }
241
242    fn get(&self, bitoff: usize, bitmax: usize) -> GetResult<usize> {
243        let first_unset = self.scan(bitoff, bitmax, true);
244        GetResult { all_set: first_unset.is_none(), first_unset: first_unset.unwrap_or(bitmax) }
245    }
246
247    fn set(&mut self, bitoff: usize, bitmax: usize) -> Result<(), Status> {
248        if bitoff > bitmax || bitmax > self.size {
249            return Err(Status::INVALID_ARGS);
250        }
251        if bitoff == bitmax {
252            return Ok(());
253        }
254        let first = first_idx(bitoff);
255        let last = last_idx(bitmax);
256        let data = self.data_mut();
257        for (i, word) in data.iter_mut().enumerate().take(last + 1).skip(first) {
258            let mask = get_mask(i == first, i == last, bitoff, bitmax);
259            *word |= mask;
260        }
261        Ok(())
262    }
263
264    fn clear(&mut self, bitoff: usize, bitmax: usize) -> Result<(), Status> {
265        if bitoff > bitmax || bitmax > self.size {
266            return Err(Status::INVALID_ARGS);
267        }
268        if bitoff == bitmax {
269            return Ok(());
270        }
271        let first = first_idx(bitoff);
272        let last = last_idx(bitmax);
273        let data = self.data_mut();
274        for (i, word) in data.iter_mut().enumerate().take(last + 1).skip(first) {
275            let mask = get_mask(i == first, i == last, bitoff, bitmax);
276            *word &= !mask;
277        }
278        Ok(())
279    }
280
281    fn clear_all(&mut self) {
282        if self.size == 0 {
283            return;
284        }
285        let last = last_idx(self.size);
286        self.data_mut()[..=last].fill(0);
287    }
288}