Skip to content

Commit

Permalink
use trailing_zeros for threadset iteration (#3871)
Browse files Browse the repository at this point in the history
(cherry picked from commit ec8ba13)
  • Loading branch information
apfitzge authored and mergify[bot] committed Dec 3, 2024
1 parent d07fc9b commit 7cb0001
Showing 1 changed file with 52 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -421,7 +421,7 @@ impl ThreadSet {

#[inline(always)]
pub(crate) fn contained_threads_iter(self) -> impl Iterator<Item = ThreadId> {
(0..MAX_THREADS).filter(move |thread_id| self.contains(*thread_id))
ThreadSetIterator(self.0)
}

#[inline(always)]
Expand All @@ -430,6 +430,22 @@ impl ThreadSet {
}
}

struct ThreadSetIterator(u64);

impl Iterator for ThreadSetIterator {
type Item = ThreadId;

fn next(&mut self) -> Option<Self::Item> {
if self.0 == 0 {
None
} else {
let thread_id = self.0.trailing_zeros() as ThreadId;
self.0 &= self.0 - 1;
Some(thread_id)
}
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down Expand Up @@ -739,4 +755,39 @@ mod tests {
let any_threads = ThreadSet::any(MAX_THREADS);
assert_eq!(any_threads.num_threads(), MAX_THREADS as u32);
}

#[test]
fn test_thread_set_iter() {
let mut thread_set = ThreadSet::none();
assert!(thread_set.contained_threads_iter().next().is_none());

thread_set.insert(4);
assert_eq!(
thread_set.contained_threads_iter().collect::<Vec<_>>(),
vec![4]
);

thread_set.insert(5);
assert_eq!(
thread_set.contained_threads_iter().collect::<Vec<_>>(),
vec![4, 5]
);
thread_set.insert(63);
assert_eq!(
thread_set.contained_threads_iter().collect::<Vec<_>>(),
vec![4, 5, 63]
);

thread_set.remove(5);
assert_eq!(
thread_set.contained_threads_iter().collect::<Vec<_>>(),
vec![4, 63]
);

let thread_set = ThreadSet::any(64);
assert_eq!(
thread_set.contained_threads_iter().collect::<Vec<_>>(),
(0..64).collect::<Vec<_>>()
);
}
}

0 comments on commit 7cb0001

Please sign in to comment.