Merge pull request #82 from woshiluo/woshiluo/hartmask-iter

spec: add some methods for `HartMask`, add `HartIds` structure as iteration of `HartMask`

sbi-spec/src/binary.rs

@@ -761,6 +761,20 @@ impl SbiRet {
     }
 }
 
+/// Check if the implementation can contains the provided `bit`.
+#[inline]
+pub(crate) const fn valid_bit(base: usize, bit: usize) -> bool {
+    if bit < base {
+        // invalid index, under minimum range.
+        false
+    } else if (bit - base) >= usize::BITS as usize {
+        // invalid index, over max range.
+        false
+    } else {
+        true
+    }
+}
+
 /// Check if the implementation contains the provided `bit`.
 ///
 /// ## Parameters
@@ -774,11 +788,7 @@ pub(crate) const fn has_bit(mask: usize, base: usize, ignore: usize, bit: usize)
     if base == ignore {
         // ignore the `mask`, consider all `bit`s as set.
         true
-    } else if bit < base {
-        // invalid index, under minimum range.
-        false
-    } else if (bit - base) >= usize::BITS as usize {
-        // invalid index, over max range.
+    } else if !valid_bit(base, bit) {
         false
     } else {
         // index is in range, check if it is set in the mask.
@@ -788,12 +798,28 @@ pub(crate) const fn has_bit(mask: usize, base: usize, ignore: usize, bit: usize)
 
 /// Hart mask structure in SBI function calls.
 #[repr(C)]
-#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
 pub struct HartMask {
     hart_mask: usize,
     hart_mask_base: usize,
 }
 
+/// Iteration for HartMask, from low to high.
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
+pub struct HartIds {
+    inner: HartMask,
+    visited_mask: usize,
+}
+
+/// Error of mask modification.
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
+pub enum MaskError {
+    /// This mask has been ignored.
+    Ignored,
+    /// Request bit is invalid.
+    InvalidBit,
+}
+
 impl HartMask {
     /// Special value to ignore the `mask`, and consider all `bit`s as set.
     pub const IGNORE_MASK: usize = usize::MAX;
@@ -829,6 +855,64 @@ impl HartMask {
             hart_id,
         )
     }
+
+    /// Add a hart id to this [HartMark].
+    /// Returns error when hart_id is invalid.
+    #[inline]
+    pub fn insert(&mut self, hart_id: usize) -> Result<(), MaskError> {
+        if self.hart_mask_base == Self::IGNORE_MASK {
+            Ok(())
+        } else if valid_bit(self.hart_mask_base, hart_id) {
+            self.hart_mask |= 1usize << (hart_id - self.hart_mask_base);
+            Ok(())
+        } else {
+            Err(MaskError::InvalidBit)
+        }
+    }
+
+    /// Remove a hart id from this [HartMark].
+    /// Returns error when hart_id is invalid, or it has been ignored.
+    #[inline]
+    pub fn remove(&mut self, hart_id: usize) -> Result<(), MaskError> {
+        if self.hart_mask_base == Self::IGNORE_MASK {
+            Err(MaskError::Ignored)
+        } else if valid_bit(self.hart_mask_base, hart_id) {
+            self.hart_mask &= !(1usize << (hart_id - self.hart_mask_base));
+            Ok(())
+        } else {
+            Err(MaskError::InvalidBit)
+        }
+    }
+
+    /// Returns [HartIds] of self.
+    #[inline]
+    pub const fn iter(&self) -> HartIds {
+        HartIds {
+            inner: HartMask {
+                hart_mask: self.hart_mask,
+                hart_mask_base: self.hart_mask_base,
+            },
+            visited_mask: 0,
+        }
+    }
+}
+
+impl Iterator for HartIds {
+    type Item = usize;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        let non_visited_mask = (!self.visited_mask) & (self.inner.hart_mask);
+        if non_visited_mask == 0 {
+            None
+        } else {
+            let low_bit = non_visited_mask.trailing_zeros();
+            let hart_id = usize::try_from(low_bit).unwrap() + self.inner.hart_mask_base;
+            self.visited_mask |= 1usize << low_bit;
+
+            Some(hart_id)
+        }
+    }
 }
 
 /// Counter index mask structure in SBI function calls for the `PMU` extension §11.
@@ -1052,6 +1136,19 @@ mod tests {
             assert!(mask.has_bit(i));
         }
         assert!(mask.has_bit(usize::MAX));
+        // Test HartIds
+        let mut mask_iter = HartMask::from_mask_base(0b101011, 1).iter();
+        assert_eq!(mask_iter.next(), Some(1usize));
+        assert_eq!(mask_iter.next(), Some(2usize));
+        assert_eq!(mask_iter.next(), Some(4usize));
+        assert_eq!(mask_iter.next(), Some(6usize));
+        assert_eq!(mask_iter.next(), None);
+        let mut mask = HartMask::from_mask_base(0, 1);
+        assert!(!mask.has_bit(1));
+        assert!(mask.insert(1).is_ok());
+        assert!(mask.has_bit(1));
+        assert!(mask.remove(1).is_ok());
+        assert!(!mask.has_bit(1));
     }
 
     #[test]