refaults_vmo/
atomic_vec.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 std::cmp::min;
6use std::sync::OnceLock;
7use std::sync::atomic::{AtomicU64, Ordering};
8
9const BITS: u64 = u64::BITS as u64;
10
11fn new_words(nwords: u64) -> Box<[AtomicU64]> {
12    (0..nwords).map(|_| AtomicU64::new(0)).collect()
13}
14
15mod private {
16    /// Prevents other crates from implementing Overflow, as it is tightly coupled with
17    /// [`AtomicBitVec`] internals.
18    pub trait Sealed {}
19}
20
21/// Holds the words of an [`AtomicBitVec`] that don't fit in its initial allocation.
22///
23/// This is what makes an `AtomicBitVec` growable or not: [`NoOverflow`] is a zero-sized type, so a
24/// vector that is never grown doesn't pay for the ability to grow.
25pub trait Overflow: private::Sealed + Default {
26    /// Returns the word at `index`, counted from the first word past the initial allocation,
27    /// allocating it if needed, or `None` if the vector cannot grow.
28    fn word(&self, index: u64) -> Option<&AtomicU64>;
29}
30
31/// Overflow for a bit vector that cannot grow. See [`Overflow`].
32#[derive(Default)]
33pub struct NoOverflow;
34
35impl private::Sealed for NoOverflow {}
36
37impl Overflow for NoOverflow {
38    fn word(&self, _index: u64) -> Option<&AtomicU64> {
39        None
40    }
41}
42
43/// Overflow for a bit vector that can grow. See [`Overflow`].
44///
45/// Words are held in a chain of segments, appended when a word past the end of the chain is
46/// requested. Segments are neither moved nor freed once published, so references to the words they
47/// hold stay valid for as long as the vector. `OnceLock` publishes a segment to the other threads
48/// walking the chain: one that doesn't observe it enters `get_or_init`, which waits for the
49/// initialization it missed.
50#[derive(Default)]
51pub struct Growable {
52    head: OnceLock<Box<Segment>>,
53}
54
55impl private::Sealed for Growable {}
56
57impl Overflow for Growable {
58    fn word(&self, mut index: u64) -> Option<&AtomicU64> {
59        // A segment that holds `index` holds at least as many words as the whole chain, since
60        // `index` counts from the end of the initial allocation. Growing a vector repeatedly
61        // therefore allocates a logarithmic number of segments.
62        let nwords = index + 1;
63        let mut next = &self.head;
64        loop {
65            // Append a segment if we reached the end of the chain. If another thread wins the
66            // race, we adopt its segment, which may be too short, and keep walking.
67            let segment = next.get_or_init(|| Box::new(Segment::new(nwords)));
68            match segment.words.get(index as usize) {
69                Some(word) => return Some(word),
70                None => index -= segment.words.len() as u64,
71            }
72            next = &segment.next;
73        }
74    }
75}
76
77/// A chunk of storage appended to a [`Growable`] bit vector.
78struct Segment {
79    words: Box<[AtomicU64]>,
80    next: OnceLock<Box<Segment>>,
81}
82
83impl Segment {
84    fn new(nwords: u64) -> Self {
85        Self { words: new_words(nwords), next: OnceLock::new() }
86    }
87}
88
89/// An atomic bit-vector.
90///
91/// By default the vector has a fixed size, and setting a bit past its end panics.
92/// `AtomicBitVec<Growable>` instead grows to hold that bit, which appends storage rather than
93/// reallocating it, so that a word handed out to a thread stays valid while another thread grows
94/// the vector.
95pub struct AtomicBitVec<O: Overflow = NoOverflow> {
96    /// Words `0..storage.len()`. Allocated once, and never reallocated.
97    storage: Box<[AtomicU64]>,
98    /// The words past the end of `storage`. Zero-sized unless the vector is growable.
99    overflow: O,
100}
101
102impl<O: Overflow> AtomicBitVec<O> {
103    /// Creates a new `AtomicBitVec` holding at least `nbits` bits, all set to false.
104    pub fn new(nbits: u64) -> Self {
105        Self { storage: new_words(nbits.div_ceil(BITS)), overflow: O::default() }
106    }
107
108    /// Sets the bits between `start_bit` (included) and `end_bit` (excluded), and returns the
109    /// number of bits that were already set. A growable vector grows to hold `end_bit`.
110    pub fn test_and_set_range(&self, start_bit: u64, end_bit: u64) -> u64 {
111        if start_bit >= end_bit {
112            return 0;
113        }
114
115        let mut counter = 0;
116        let mut current_bit = start_bit;
117
118        while current_bit < end_bit {
119            let current_word_index = current_bit / BITS;
120            let current_word_start_bit = current_word_index * BITS;
121            let mask =
122                Self::get_mask(current_bit % BITS, min(end_bit - current_word_start_bit, BITS));
123            let old_word = self.word(current_word_index).fetch_or(mask, Ordering::Relaxed);
124            counter += (old_word & mask).count_ones();
125            current_bit = current_word_start_bit + BITS;
126        }
127
128        counter.into()
129    }
130
131    /// Copy the first `nbits` bits. Test only.
132    pub fn to_vec(&self, nbits: u64) -> Vec<bool> {
133        let fetch = |bit: u64| {
134            let bit_mask = 1 << (bit % BITS);
135            self.word(bit / BITS).load(Ordering::Relaxed) & bit_mask != 0
136        };
137        (0..nbits).map(fetch).collect()
138    }
139
140    /// Returns the word at `index`, which a fixed-size vector must hold, and which a growable one
141    /// grows to hold.
142    fn word(&self, index: u64) -> &AtomicU64 {
143        match self.storage.get(index as usize) {
144            Some(word) => word,
145            None => self.overflow_word(index - self.storage.len() as u64),
146        }
147    }
148
149    /// Returns the overflow word at `index`, counted from the end of `storage`. Cold because most
150    /// files will not be resized.
151    #[cold]
152    fn overflow_word(&self, index: u64) -> &AtomicU64 {
153        self.overflow.word(index).expect("word index outside of the vector")
154    }
155
156    fn get_mask(start_bit: u64, end_bit: u64) -> u64 {
157        let left_mask = u64::MAX << start_bit;
158        let right_mask = u64::MAX >> (BITS - end_bit);
159        left_mask & right_mask
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    #[test]
168    fn test_new() {
169        let vec = AtomicBitVec::<NoOverflow>::new(0);
170        assert_eq!(vec.storage.len(), 0);
171
172        let vec = AtomicBitVec::<NoOverflow>::new(1);
173        assert_eq!(vec.storage.len(), 1);
174        assert_eq!(vec.storage[0].load(Ordering::Relaxed), 0);
175
176        let vec = AtomicBitVec::<NoOverflow>::new(64);
177        assert_eq!(vec.storage.len(), 1);
178        assert_eq!(vec.storage[0].load(Ordering::Relaxed), 0);
179
180        let vec = AtomicBitVec::<NoOverflow>::new(65);
181        assert_eq!(vec.storage.len(), 2);
182        assert_eq!(vec.storage[0].load(Ordering::Relaxed), 0);
183        assert_eq!(vec.storage[1].load(Ordering::Relaxed), 0);
184    }
185
186    #[test]
187    fn test_size() {
188        // A vector that cannot grow pays nothing for the ability to grow: `FxBlob` holds one, and
189        // sits exactly on its own size budget.
190        assert!(size_of::<AtomicBitVec>() <= 16);
191        assert!(size_of::<AtomicBitVec<Growable>>() <= 32);
192    }
193
194    #[test]
195    fn test_test_and_set() {
196        let vec = AtomicBitVec::<NoOverflow>::new(320);
197
198        // Set bits and check they were not set before.
199        assert_eq!(vec.test_and_set_range(10, 20), 0);
200        // Check they are set now.
201        assert_eq!(vec.test_and_set_range(10, 20), 10);
202
203        // Check another range, partially overlapping.
204        assert_eq!(vec.test_and_set_range(15, 25), 5);
205
206        // A range across two words.
207        assert_eq!(vec.test_and_set_range(60, 70), 0);
208        // Only 10 more.
209        assert_eq!(vec.test_and_set_range(55, 75), 10);
210
211        // Large range.
212        assert_eq!(vec.test_and_set_range(50, 300), 20);
213        assert_eq!(vec.test_and_set_range(50, 300), 250);
214    }
215
216    #[test]
217    fn test_test_and_set2() {
218        let vec = AtomicBitVec::<NoOverflow>::new(320);
219        assert_eq!(vec.test_and_set_range(64, 128), 0);
220        assert_eq!(vec.test_and_set_range(64, 128), 64);
221    }
222
223    #[test]
224    fn test_test_and_set3() {
225        let vec = AtomicBitVec::<NoOverflow>::new(150);
226
227        assert_eq!(vec.test_and_set_range(0, 1), 0);
228        assert_eq!(vec.test_and_set_range(63, 64), 0);
229        // 0 and 63 are already set.
230        assert_eq!(vec.test_and_set_range(0, 64), 2);
231        assert_eq!(vec.test_and_set_range(64, 65), 0);
232        // 63 and 64 are already set.
233        assert_eq!(vec.test_and_set_range(63, 65), 2);
234
235        // 63 and 64 are already set.
236        assert_eq!(vec.test_and_set_range(63, 128), 2);
237        // 63 to 127 (included) are already set.
238        assert_eq!(vec.test_and_set_range(63, 129), 65);
239    }
240
241    // A fixed-size vector panics once a write reaches a word it does not hold.
242    #[test]
243    #[should_panic]
244    fn test_and_set_out_of_bounds() {
245        let vec = AtomicBitVec::<NoOverflow>::new(10);
246        vec.test_and_set_range(64, 65);
247    }
248
249    #[test]
250    fn test_grow() {
251        // A growable vector does not panic when a write reaches the end, but it grows.
252        let vec = AtomicBitVec::<Growable>::new(10);
253        vec.test_and_set_range(0, 10);
254        assert_eq!(vec.test_and_set_range(10, 64), 0);
255        assert_eq!(vec.test_and_set_range(0, 64), 64);
256
257        // Growing past it appends a segment.
258        assert_eq!(vec.test_and_set_range(64, 65), 0);
259        assert_eq!(vec.test_and_set_range(0, 65), 65);
260    }
261
262    #[test]
263    fn test_grow_across_segments() {
264        let vec = AtomicBitVec::<Growable>::new(64);
265
266        // Set a range spanning the original storage and two new segments.
267        assert_eq!(vec.test_and_set_range(60, 200), 0);
268
269        // The count works, even across segments.
270        assert_eq!(vec.test_and_set_range(60, 200), 140);
271
272        // Bits outside of the range were not touched.
273        assert_eq!(vec.test_and_set_range(0, 60), 0);
274    }
275
276    #[test]
277    fn test_concurrent_grow() {
278        const THREADS: u64 = 4;
279        const ROUNDS: u64 = 1000;
280
281        let vec = std::sync::Arc::new(AtomicBitVec::<Growable>::new(1));
282        let threads: Vec<_> = (0..THREADS)
283            .map(|_| {
284                let vec = vec.clone();
285                std::thread::spawn(move || {
286                    // Every thread races to grow the vector and to set the bits it exposes.
287                    // Count the bits that this thread set, i.e. that it found unset.
288                    let mut bits_set = 0;
289                    for round in 1..=ROUNDS {
290                        bits_set += round - vec.test_and_set_range(0, round);
291                    }
292                    bits_set
293                })
294            })
295            .collect();
296
297        let bits_set: u64 = threads.into_iter().map(|thread| thread.join().unwrap()).sum();
298
299        // Every bit goes from unset to set exactly once, and exactly one thread sees it happen.
300        assert_eq!(bits_set, ROUNDS);
301        assert!(vec.to_vec(ROUNDS).into_iter().all(|bit| bit));
302    }
303}
refaults_vmo/atomic_vec.rs

refaults_vmo/
atomic_vec.rs
