use trailing_zeros for threadset iteration (#3871)

core/src/banking_stage/transaction_scheduler/thread_aware_account_locks.rs

@@ -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)]
@@ -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::*;
@@ -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<_>>()
+        );
+    }
 }