hv1_structs/
proc_mask.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Structures for working with processor masks.
5
6/// A set of processor IDs, stored as a sparse array of 64-bit masks.
7#[derive(Copy, Clone)]
8pub struct ProcessorSet<'a> {
9    valid_masks: u64,
10    masks: &'a [u64],
11}
12
13impl<'a> ProcessorSet<'a> {
14    /// Attempts to create a ProcessorSet from a HV_GENERIC_SET_SPARSE_4K format HV_GENERIC_SET.
15    pub fn from_generic_set(format: u64, rest: &'a [u64]) -> Option<Self> {
16        if format != hvdef::hypercall::HV_GENERIC_SET_SPARSE_4K {
17            return None;
18        }
19        let &[valid_masks, ref masks @ ..] = rest else {
20            return None;
21        };
22        Self::from_processor_masks(valid_masks, masks)
23    }
24
25    /// Attempts to create a ProcessorSet from a set of processor masks.
26    pub fn from_processor_masks(valid_masks: u64, masks: &'a [u64]) -> Option<Self> {
27        let mask_count = valid_masks.count_ones();
28        if masks.len() != mask_count as usize {
29            return None;
30        }
31        Some(Self { valid_masks, masks })
32    }
33
34    /// Returns the set as an iterator of u64s, suitable for collecting and
35    /// using as raw HV_GENERIC_SET_SPARSE_4K in a hypercall.
36    pub fn as_generic_set(&self) -> impl Iterator<Item = u64> + use<'_> {
37        std::iter::once(self.valid_masks).chain(self.masks.iter().copied())
38    }
39
40    /// Returns true if the set is empty.
41    pub fn is_empty(&self) -> bool {
42        self.valid_masks == 0 || self.count() == 0
43    }
44
45    /// Returns the number of processors in the set.
46    pub fn count(&self) -> usize {
47        self.masks.iter().map(|x| x.count_ones() as usize).sum()
48    }
49
50    /// Returns an iterator over the processor IDs in the set.
51    pub fn iter(&self) -> ProcessorSetIter<'a> {
52        self.into_iter()
53    }
54}
55
56impl<'a> IntoIterator for ProcessorSet<'a> {
57    type Item = u32;
58    type IntoIter = ProcessorSetIter<'a>;
59
60    fn into_iter(self) -> Self::IntoIter {
61        ProcessorSetIter {
62            bit: 0,
63            mask: 0,
64            remaining_valid: self.valid_masks,
65            masks: self.masks,
66        }
67    }
68}
69
70/// An iterator over the processor IDs in a ProcessorSet.
71pub struct ProcessorSetIter<'a> {
72    bit: u32,
73    mask: u64,
74    remaining_valid: u64,
75    masks: &'a [u64],
76}
77
78impl Iterator for ProcessorSetIter<'_> {
79    type Item = u32;
80
81    fn next(&mut self) -> Option<u32> {
82        while self.mask == 0 {
83            self.mask = *self.masks.first()?;
84            self.masks = &self.masks[1..];
85            self.bit = self.remaining_valid.trailing_zeros();
86            self.remaining_valid &= !(1 << self.bit);
87        }
88        let proc = self.mask.trailing_zeros();
89        self.mask &= !(1 << proc);
90        Some(self.bit * 64 + proc)
91    }
92}
93
94impl std::iter::FusedIterator for ProcessorSetIter<'_> {}
95
96#[cfg(test)]
97mod test {
98    use super::*;
99
100    #[test]
101    // Values taken from the Hypervisor Functional Specification
102    fn test_processor_set() {
103        let set = ProcessorSet::from_processor_masks(0x5, &[0x21, 0x4]).unwrap();
104        assert_eq!(set.count(), 3);
105
106        let mut iter = set.into_iter();
107        assert_eq!(iter.next(), Some(0));
108        assert_eq!(iter.next(), Some(5));
109        assert_eq!(iter.next(), Some(130));
110        assert_eq!(iter.next(), None);
111    }
112}