diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..531ddd1 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,55 @@ +name: CI + +on: [push, pull_request] + +jobs: + ci: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + rust-toolchain: [nightly] + targets: [x86_64-unknown-linux-gnu, x86_64-unknown-none, riscv64gc-unknown-none-elf, aarch64-unknown-none-softfloat] + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@nightly + with: + toolchain: ${{ matrix.rust-toolchain }} + components: rust-src, clippy, rustfmt + targets: ${{ matrix.targets }} + - name: Check rust version + run: rustc --version --verbose + - name: Check code format + run: cargo fmt --all -- --check + - name: Clippy + run: cargo clippy --target ${{ matrix.targets }} --all-features -- -A clippy::new_without_default + - name: Build + run: cargo build --target ${{ matrix.targets }} --all-features + - name: Unit test + if: ${{ matrix.targets == 'x86_64-unknown-linux-gnu' }} + run: cargo test --target ${{ matrix.targets }} -- --nocapture + + doc: + runs-on: ubuntu-latest + strategy: + fail-fast: false + permissions: + contents: write + env: + default-branch: ${{ format('refs/heads/{0}', github.event.repository.default_branch) }} + RUSTDOCFLAGS: -D rustdoc::broken_intra_doc_links -D missing-docs + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@nightly + - name: Build docs + continue-on-error: ${{ github.ref != env.default-branch && github.event_name != 'pull_request' }} + run: | + cargo doc --no-deps --all-features + printf '' $(cargo tree | head -1 | cut -d' ' -f1) > target/doc/index.html + - name: Deploy to Github Pages + if: ${{ github.ref == env.default-branch }} + uses: JamesIves/github-pages-deploy-action@v4 + with: + single-commit: true + branch: gh-pages + folder: target/doc diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ff78c42 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +/target +/.vscode +.DS_Store +Cargo.lock diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..d938d92 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "cpu_set" +version = "0.1.0" +edition = "2021" + +[dependencies] +bitmaps = { version = "3.2.1", default-features = false } \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..9cd6975 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,489 @@ +//! CPU set library +//! `cpu_set` represents a set of physical CPUs, which is implemented as a bit map, +//! refering to `cpu_set_t` in Linux. +//! Reference: +//! * https://man7.org/linux/man-pages/man2/sched_setaffinity.2.html +//! * https://man7.org/linux/man-pages/man3/CPU_SET.3.html + +use core::hash::{Hash, Hasher}; +use core::ops::*; + +use bitmaps::{BitOps, Bitmap, Bits, BitsImpl}; + +/// A compact array of bits which represents a set of physical CPUs, +/// implemented based on [bitmaps::Bitmap](https://docs.rs/bitmaps/latest/bitmaps/struct.Bitmap.html). +/// +/// The type used to store the cpuset will be the minimum unsigned integer type +/// required to fit the number of bits, from `u8` to `u128`. If the size is 1, +/// `bool` is used. If the size exceeds 128, an array of `u128` will be used, +/// sized as appropriately. The maximum supported size is currently 1024, +/// represented by an array `[u128; 8]`. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct CpuSet +where + BitsImpl<{ SIZE }>: Bits, +{ + value: Bitmap<{ SIZE }>, +} + +impl Hash for CpuSet<{ SIZE }> +where + BitsImpl<{ SIZE }>: Bits, + as Bits>::Store: Hash, +{ + fn hash(&self, state: &mut H) { + self.value.as_value().hash(state) + } +} + +impl PartialOrd for CpuSet<{ SIZE }> +where + BitsImpl<{ SIZE }>: Bits, + as Bits>::Store: PartialOrd, +{ + fn partial_cmp(&self, other: &Self) -> Option { + self.value.as_value().partial_cmp(other.value.as_value()) + } +} + +impl Ord for CpuSet<{ SIZE }> +where + BitsImpl<{ SIZE }>: Bits, + as Bits>::Store: Ord, +{ + fn cmp(&self, other: &Self) -> core::cmp::Ordering { + self.value.as_value().cmp(other.value.as_value()) + } +} + +impl CpuSet<{ SIZE }> +where + BitsImpl: Bits, +{ + /// Construct a cpuset with every bit set to `false`. + #[inline] + pub fn new() -> Self { + Self::default() + } + + /// Construct a cpuset where every bit with index less than `bits` is + /// `true`, and every other bit is `false`. + #[inline] + pub fn mask(bits: usize) -> Self { + debug_assert!(bits <= SIZE); + Self { + value: Bitmap::mask(bits), + } + } + + /// Construct a cpuset from a value of the same type as its backing store. + #[inline] + pub fn from_value(data: as Bits>::Store) -> Self { + Self { + value: Bitmap::from_value(data), + } + } + + pub fn from_usize(value: usize) -> Self { + assert!(value >> SIZE == 0); + + let mut bit_map = Bitmap::new(); + let mut i = 0; + while i < SIZE { + if value & (1 << i) != 0 { + bit_map.set(i, true); + } + i += 1; + } + + Self { value: bit_map } + } + + /// Convert this cpuset into a value of the type of its backing store. + #[inline] + pub fn into_value(self) -> as Bits>::Store { + self.value.into_value() + } + + /// Get a reference to this cpuset's backing store. + #[inline] + pub fn as_value(&self) -> & as Bits>::Store { + self.value.as_value() + } + + #[inline] + pub fn as_bytes(&self) -> &[u8] { + self.value.as_bytes() + } + + /// Count the number of `true` bits in the cpuset. + #[inline] + pub fn len(self) -> usize { + self.value.len() + } + + /// Test if the cpuset contains only `false` bits. + #[inline] + pub fn is_empty(self) -> bool { + self.first_index().is_none() + } + + /// Test if the cpuset contains only `true` bits. + #[inline] + pub fn is_full(self) -> bool { + self.first_false_index().is_none() + } + + /// Get the value of the bit at a given index. + #[inline] + pub fn get(self, index: usize) -> bool { + debug_assert!(index < SIZE); + as Bits>::Store::get(&self.into_value(), index) + } + + /// Set the value of the bit at a given index. + /// + /// Returns the previous value of the bit. + #[inline] + pub fn set(&mut self, index: usize, value: bool) -> bool { + debug_assert!(index < SIZE); + as Bits>::Store::set(&mut self.into_value(), index, value) + } + + /// Find the index of the first `true` bit in the cpuset. + #[inline] + pub fn first_index(self) -> Option { + as Bits>::Store::first_index(&self.into_value()) + } + + /// Find the index of the last `true` bit in the cpuset. + #[inline] + pub fn last_index(self) -> Option { + as Bits>::Store::last_index(&self.into_value()) + } + + /// Find the index of the first `true` bit in the cpuset after `index`. + #[inline] + pub fn next_index(self, index: usize) -> Option { + as Bits>::Store::next_index(&self.into_value(), index) + } + + /// Find the index of the last `true` bit in the cpuset before `index`. + #[inline] + pub fn prev_index(self, index: usize) -> Option { + as Bits>::Store::prev_index(&self.into_value(), index) + } + + /// Find the index of the first `false` bit in the cpuset. + #[inline] + pub fn first_false_index(self) -> Option { + as Bits>::corrected_first_false_index(&self.into_value()) + } + + /// Find the index of the last `false` bit in the cpuset. + #[inline] + pub fn last_false_index(self) -> Option { + as Bits>::corrected_last_false_index(&self.into_value()) + } + + /// Find the index of the first `false` bit in the cpuset after `index`. + #[inline] + pub fn next_false_index(self, index: usize) -> Option { + as Bits>::corrected_next_false_index(&self.into_value(), index) + } + + /// Find the index of the first `false` bit in the cpuset before `index`. + #[inline] + pub fn prev_false_index(self, index: usize) -> Option { + as Bits>::Store::prev_false_index(&self.into_value(), index) + } + + /// Invert all the bits in the cpuset. + #[inline] + pub fn invert(&mut self) { + as Bits>::Store::invert(&mut self.into_value()); + } +} + +impl<'a, const SIZE: usize> IntoIterator for &'a CpuSet<{ SIZE }> +where + BitsImpl<{ SIZE }>: Bits, +{ + type Item = usize; + type IntoIter = Iter<'a, { SIZE }>; + + fn into_iter(self) -> Self::IntoIter { + Iter { + head: None, + tail: Some(SIZE + 1), + data: self, + } + } +} + +impl BitAnd for CpuSet<{ SIZE }> +where + BitsImpl<{ SIZE }>: Bits, +{ + type Output = Self; + fn bitand(self, rhs: Self) -> Self::Output { + Self { + value: self.value.bitand(rhs.value), + } + } +} + +impl BitOr for CpuSet<{ SIZE }> +where + BitsImpl<{ SIZE }>: Bits, +{ + type Output = Self; + fn bitor(self, rhs: Self) -> Self::Output { + Self { + value: self.value.bitor(rhs.value), + } + } +} + +impl BitXor for CpuSet<{ SIZE }> +where + BitsImpl<{ SIZE }>: Bits, +{ + type Output = Self; + fn bitxor(self, rhs: Self) -> Self::Output { + Self { + value: self.value.bitxor(rhs.value), + } + } +} + +impl Not for CpuSet<{ SIZE }> +where + BitsImpl<{ SIZE }>: Bits, +{ + type Output = Self; + fn not(self) -> Self::Output { + Self { + value: self.value.not(), + } + } +} + +impl BitAndAssign for CpuSet<{ SIZE }> +where + BitsImpl<{ SIZE }>: Bits, +{ + fn bitand_assign(&mut self, rhs: Self) { + self.value.bitand_assign(rhs.value) + } +} + +impl BitOrAssign for CpuSet<{ SIZE }> +where + BitsImpl<{ SIZE }>: Bits, +{ + fn bitor_assign(&mut self, rhs: Self) { + self.value.bitor_assign(rhs.value) + } +} + +impl BitXorAssign for CpuSet<{ SIZE }> +where + BitsImpl<{ SIZE }>: Bits, +{ + fn bitxor_assign(&mut self, rhs: Self) { + self.value.bitxor_assign(rhs.value) + } +} + +impl From<[u128; 2]> for CpuSet<256> { + fn from(data: [u128; 2]) -> Self { + CpuSet { value: data.into() } + } +} + +impl From<[u128; 3]> for CpuSet<384> { + fn from(data: [u128; 3]) -> Self { + CpuSet { value: data.into() } + } +} + +impl From<[u128; 4]> for CpuSet<512> { + fn from(data: [u128; 4]) -> Self { + CpuSet { value: data.into() } + } +} + +impl From<[u128; 5]> for CpuSet<640> { + fn from(data: [u128; 5]) -> Self { + CpuSet { value: data.into() } + } +} + +impl From<[u128; 6]> for CpuSet<768> { + fn from(data: [u128; 6]) -> Self { + CpuSet { value: data.into() } + } +} + +impl From<[u128; 7]> for CpuSet<896> { + fn from(data: [u128; 7]) -> Self { + CpuSet { value: data.into() } + } +} + +impl From<[u128; 8]> for CpuSet<1024> { + fn from(data: [u128; 8]) -> Self { + CpuSet { value: data.into() } + } +} + +impl From> for [u128; 2] { + fn from(cpuset: CpuSet<256>) -> Self { + cpuset.into_value() + } +} + +impl From> for [u128; 3] { + fn from(cpuset: CpuSet<384>) -> Self { + cpuset.into_value() + } +} + +impl From> for [u128; 4] { + fn from(cpuset: CpuSet<512>) -> Self { + cpuset.into_value() + } +} + +impl From> for [u128; 5] { + fn from(cpuset: CpuSet<640>) -> Self { + cpuset.into_value() + } +} + +impl From> for [u128; 6] { + fn from(cpuset: CpuSet<768>) -> Self { + cpuset.into_value() + } +} + +impl From> for [u128; 7] { + fn from(cpuset: CpuSet<896>) -> Self { + cpuset.into_value() + } +} + +impl From> for [u128; 8] { + fn from(cpuset: CpuSet<1024>) -> Self { + cpuset.into_value() + } +} + +/// An iterator over the indices in a cpuset which are `true`. +/// +/// This yields a sequence of `usize` indices, not their contents (which are +/// always `true` anyway, by definition). +/// +/// # Examples +/// +/// ```rust +/// # use cpu_set::CpuSet; +/// let mut cpuset: CpuSet<10> = CpuSet::new(); +/// cpuset.set(3, true); +/// cpuset.set(5, true); +/// cpuset.set(8, true); +/// let true_indices: Vec = cpuset.into_iter().collect(); +/// assert_eq!(vec![3, 5, 8], true_indices); +/// ``` +#[derive(Clone, Debug)] +pub struct Iter<'a, const SIZE: usize> +where + BitsImpl: Bits, +{ + head: Option, + tail: Option, + data: &'a CpuSet<{ SIZE }>, +} + +impl<'a, const SIZE: usize> Iterator for Iter<'a, SIZE> +where + BitsImpl<{ SIZE }>: Bits, +{ + type Item = usize; + + fn next(&mut self) -> Option { + let result; + + match self.head { + None => { + result = self.data.first_index(); + } + Some(index) => { + if index >= SIZE { + result = None + } else { + result = self.data.next_index(index); + } + } + } + + if let Some(index) = result { + if let Some(tail) = self.tail { + if tail < index { + self.head = Some(SIZE + 1); + self.tail = None; + return None; + } + } else { + // tail is already done + self.head = Some(SIZE + 1); + return None; + } + + self.head = Some(index); + } else { + self.head = Some(SIZE + 1); + } + + result + } +} + +impl<'a, const SIZE: usize> DoubleEndedIterator for Iter<'a, SIZE> +where + BitsImpl<{ SIZE }>: Bits, +{ + fn next_back(&mut self) -> Option { + let result; + + match self.tail { + None => { + result = None; + } + Some(index) => { + if index >= SIZE { + result = self.data.last_index(); + } else { + result = self.data.prev_index(index); + } + } + } + + if let Some(index) = result { + if let Some(head) = self.head { + if head > index { + self.head = Some(SIZE + 1); + self.tail = None; + return None; + } + } + + self.tail = Some(index); + } else { + self.tail = None; + } + + result + } +}