diff --git a/.cargo/config.toml b/.cargo/config.toml index 20f8233cf..2eea954df 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -9,7 +9,6 @@ test_awkernel_lib = "test --package awkernel_lib --no-default-features --feature test_awkernel_async_lib = "test --package awkernel_async_lib --no-default-features --features std" test_awkernel_drivers = "test --package awkernel_drivers --no-default-features --features std" test_smoltcp = "test --package smoltcp --no-default-features --features std --features awkernel" -test_affinity_btree_queue = "test --package affinity_btree_queue" test_rd_gen_to_dags = "test --package rd_gen_to_dags --no-default-features --features std --features milliseconds" clippy_x86 = "clippy --package awkernel --no-default-features --features x86 --target targets/x86_64-kernel.json -Zbuild-std=core,alloc,compiler_builtins -Zbuild-std-features=compiler-builtins-mem -Zjson-target-spec" diff --git a/Cargo.toml b/Cargo.toml index 65a89b195..b752d950b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,6 @@ [workspace] members = [ - "affinity_btree_queue", "kernel", "userland", "x86bootdisk", @@ -22,3 +21,6 @@ members = [ ] resolver = "2" + +[workspace.dependencies] +affinity_btree_queue = { version = "0.1.1" } diff --git a/Makefile b/Makefile index 28a00f6b5..b051e1de5 100644 --- a/Makefile +++ b/Makefile @@ -260,7 +260,6 @@ test: FORCE cargo test_awkernel_async_lib -- --nocapture cargo test_awkernel_drivers cargo test_smoltcp - cargo test_affinity_btree_queue cargo test_rd_gen_to_dags # Format diff --git a/affinity_btree_queue/Cargo.toml b/affinity_btree_queue/Cargo.toml deleted file mode 100644 index ed0444f1a..000000000 --- a/affinity_btree_queue/Cargo.toml +++ /dev/null @@ -1,12 +0,0 @@ -[package] -name = "affinity_btree_queue" -version = "0.1.1" -edition = "2024" -license = "MIT OR Apache-2.0" -description = "An affinity-aware priority queue backed by an augmented classical (CLRS-style) B-tree." -readme = "README.md" -keywords = ["no-std", "priority-queue"] -categories = ["data-structures"] -repository = "https://github.com/ytakano/affinity_btree_queue" - -[dependencies] diff --git a/affinity_btree_queue/README.md b/affinity_btree_queue/README.md deleted file mode 100644 index 3176296c2..000000000 --- a/affinity_btree_queue/README.md +++ /dev/null @@ -1,159 +0,0 @@ -# Affinity BTree Queue - -An affinity-aware priority queue backed by an augmented classical (CLRS-style) B-tree, written in Rust. - -Each entry carries `(priority, affinity: CpuMask, value)`. -`pop_for_cpu(cpu)` returns the highest-priority entry whose affinity includes that CPU in O(log N), -making it suitable for task schedulers where tasks are pinned to specific CPUs. - -## Features - -- **Affinity-aware scheduling** — every node stores the OR of all affinities in its subtree; - `pop_for_cpu` prunes ineligible subtrees in a single traversal -- **`no_std` + `alloc`** — works in embedded and OS-kernel environments -- **O(log N)** push / pop / remove via an arena-based augmented B-tree -- **Generic priority** — any `P: Ord`; smaller value = higher priority; - use `core::cmp::Reverse` for max-first order -- **Configurable CPU count** — `CpuMask` supports up to `N × 64` CPUs - (default `N = 1` → 64 CPUs, `N = 2` → 128 CPUs, …) -- **No unsafe code** — `#![forbid(unsafe_code)]` - -## Usage - -Add the crate to `Cargo.toml`: - -```toml -[dependencies] -affinity_btree_queue = { path = "..." } -``` - -### Basic example - -```rust -use affinity_btree_queue::{AffinityBTreeQueue, CpuMask}; - -// Create a queue for a 4-CPU system. -// Smaller priority value = higher priority. -let mut q: AffinityBTreeQueue = AffinityBTreeQueue::new(4); - -// Build CPU affinity masks. -let mut cpu0 = CpuMask::empty(); -cpu0.insert(0); - -let mut cpu12 = CpuMask::empty(); -cpu12.insert(1); -cpu12.insert(2); - -let mut all = CpuMask::empty(); -for c in 0..4 { all.insert(c); } - -// push returns a PriorityKey handle that can be passed to remove() later. -q.push(10, all, "low priority, any CPU").unwrap(); -q.push(5, cpu12, "medium, CPUs 1-2 only").unwrap(); -q.push(1, cpu0, "high priority, CPU 0 only").unwrap(); - -// pop_for_cpu returns the highest-priority entry runnable on the given CPU. -let (_, _, v) = q.pop_for_cpu(0).unwrap(); -assert_eq!(v, "high priority, CPU 0 only"); // priority=1 - -// The cpu12 entry (priority=5) is not runnable on CPU 0, so priority=10 is next. -let (_, _, v) = q.pop_for_cpu(0).unwrap(); -assert_eq!(v, "low priority, any CPU"); // priority=10 - -// CPU 1 can run the cpu12 entry (priority=5). -let (_, _, v) = q.pop_for_cpu(1).unwrap(); -assert_eq!(v, "medium, CPUs 1-2 only"); // priority=5 -``` - -### Max-first ordering with `Reverse` - -```rust -use affinity_btree_queue::{AffinityBTreeQueue, CpuMask}; -use core::cmp::Reverse; - -let mut q: AffinityBTreeQueue, u32> = AffinityBTreeQueue::new(1); -let mut m = CpuMask::empty(); -m.insert(0); -q.push(Reverse(1), m, 1).unwrap(); -q.push(Reverse(100), m, 100).unwrap(); - -// Reverse(100) < Reverse(1), so the numerically larger value dequeues first. -let (_, _, v) = q.pop_for_cpu(0).unwrap(); -assert_eq!(v, 100); -``` - -## API overview - -### `AffinityBTreeQueue` - -| Method | Description | -|--------|-------------| -| `new(num_cpus)` | Create an empty queue for CPUs `0..num_cpus` | -| `push(priority, affinity, value)` | Insert an entry; returns a `PriorityKey` handle | -| `pop_for_cpu(cpu)` | Remove and return the highest-priority entry runnable on `cpu` | -| `peek_key_for_cpu(cpu)` | Return the key of the best entry for `cpu` without removing it | -| `remove(&key)` | Remove an entry by its `PriorityKey` | -| `len()` / `is_empty()` | Entry count | -| `num_cpus()` | Configured CPU count | - -Type parameters: `P: Ord` (priority), `V` (value), `T` (B-tree minimum degree, default 16), -`N` (words in `CpuMask`, default 1). - -### `CpuMask` - -| Method | Description | -|--------|-------------| -| `empty()` | All-zero mask | -| `insert(cpu)` | Set the bit for `cpu` | -| `contains(cpu)` | Test whether `cpu` is set | -| `union(other)` | Bitwise OR of two masks | -| `from_bits_truncate(bits)` | Construct from a raw `u64` (word 0) | - -### `PushError` - -| Variant | Meaning | -|---------|---------| -| `EmptyAffinity` | `affinity` has no bits set | -| `InvalidCpu` | `affinity` has a bit set for a CPU ≥ `num_cpus` | -| `SequenceOverflow` | Internal sequence counter exhausted (2⁶⁴ pushes) | - -## Building and testing - -```bash -cargo build -cargo test # unit tests + randomized property tests + doctests -``` - -Tests compare every operation against a simple `Vec`-based reference model and verify -all B-tree structural invariants (key order, subtree affinity summaries, leaf depth -uniformity, entry count) after every mutation. - -## Implementation notes - -- The tree is an **arena-based augmented B-tree** (`Vec>` + free list). - Node IDs are stable indices; the arena is never compacted. -- Every node stores **`subtree_affinity`** = OR of all entry affinities in its subtree. - `pop_for_cpu` uses this summary to skip entire subtrees with no eligible entry. -- Ordering is by **`(priority, seq)`** — the monotonically increasing `seq` counter - makes every key unique, so no two entries ever compare equal. -- Priority / affinity updates are **remove + re-insert**; no in-place mutation of keys. -- Deletion is **top-down**: before descending into a child, the algorithm ensures it has - ≥ T entries (borrowing from a sibling or merging), so underflow never propagates upward. - -## License - -This project is licensed under either of - -- Apache License, Version 2.0 -- MIT license - -at your option. - -## Test Coverage - -Line-by-line coverage data can be generated locally with `cargo llvm-cov`: - -```bash -cargo llvm-cov --workspace --all-features --text --output-path coverage.txt -cargo llvm-cov --workspace --all-features --lcov --output-path lcov.info -``` diff --git a/affinity_btree_queue/src/lib.rs b/affinity_btree_queue/src/lib.rs deleted file mode 100644 index 39200abf1..000000000 --- a/affinity_btree_queue/src/lib.rs +++ /dev/null @@ -1,1325 +0,0 @@ -//! Affinity-aware priority queue backed by an augmented classical (CLRS-style) B-tree. -//! -//! Each entry carries `(priority, affinity: CpuMask, value)`. The tree is ordered -//! solely by `PriorityKey

= (priority, seq)`, where `seq` is a unique tie-breaker. -//! Every node stores `subtree_affinity`, the OR of all entry affinities in its -//! subtree, which lets `peek_key_for_cpu`/`pop_for_cpu` skip subtrees with no -//! entry runnable on the requested CPU. Smaller priority value = higher priority; -//! wrap with `core::cmp::Reverse` for the opposite order. -//! -//! # Quick start -//! -//! ``` -//! use affinity_btree_queue::{AffinityBTreeQueue, CpuMask}; -//! -//! // Create a queue for a 4-CPU system. -//! // Smaller priority value = higher priority. -//! // Use `core::cmp::Reverse` if you want larger values to dequeue first. -//! let mut q: AffinityBTreeQueue = AffinityBTreeQueue::new(4); -//! -//! // Build CPU affinity masks. -//! let mut cpu0 = CpuMask::empty(); -//! cpu0.insert(0); -//! -//! let mut cpu12 = CpuMask::empty(); -//! cpu12.insert(1); -//! cpu12.insert(2); -//! -//! let mut all = CpuMask::empty(); -//! for c in 0..4 { all.insert(c); } -//! -//! // push returns a PriorityKey handle that can be passed to remove() later. -//! q.push(10, all, "low priority, any CPU").unwrap(); -//! q.push(5, cpu12, "medium, CPUs 1-2 only").unwrap(); -//! q.push(1, cpu0, "high priority, CPU 0 only").unwrap(); -//! -//! // pop_for_cpu(cpu) returns the highest-priority (lowest value) entry -//! // whose affinity includes the given CPU. -//! let (_, _, v) = q.pop_for_cpu(0).unwrap(); -//! assert_eq!(v, "high priority, CPU 0 only"); // priority=1 -//! -//! // The cpu12 entry (priority=5) is not runnable on CPU 0, so priority=10 is next. -//! let (_, _, v) = q.pop_for_cpu(0).unwrap(); -//! assert_eq!(v, "low priority, any CPU"); // priority=10 -//! -//! // CPU 1 can run the cpu12 entry (priority=5). -//! let (_, _, v) = q.pop_for_cpu(1).unwrap(); -//! assert_eq!(v, "medium, CPUs 1-2 only"); // priority=5 -//! ``` -//! -//! # Max-first ordering with `Reverse` -//! -//! ``` -//! use affinity_btree_queue::{AffinityBTreeQueue, CpuMask}; -//! use core::cmp::Reverse; -//! -//! let mut q: AffinityBTreeQueue, u32> = AffinityBTreeQueue::new(1); -//! let mut m = CpuMask::empty(); -//! m.insert(0); -//! q.push(Reverse(1), m, 1).unwrap(); -//! q.push(Reverse(100), m, 100).unwrap(); -//! -//! // Reverse(100) < Reverse(1), so the numerically larger value dequeues first. -//! let (_, _, v) = q.pop_for_cpu(0).unwrap(); -//! assert_eq!(v, 100); -//! ``` - -#![no_std] -#![forbid(unsafe_code)] - -extern crate alloc; - -use alloc::vec::Vec; -use core::cmp::Ordering; -use core::mem; - -/// Default B-tree minimum degree for production builds. -pub const DEFAULT_MIN_DEGREE: usize = 16; - -/// Sequence number type used as a unique tie-breaker for entries with the same -/// priority. -pub type Seq = u64; - -/// CPU affinity bitmask backed by `N` 64-bit words, supporting up to `N * 64` -/// CPUs. `CpuMask<1>` (the default) supports up to 64 CPUs; `CpuMask<2>` -/// supports up to 128, etc. -/// -/// All constructors and bit operations are `const fn`, so masks can be built -/// in `const`/`static` initializers with the chainable [`with`](Self::with) -/// builder. -/// -/// # Examples -/// -/// ``` -/// use affinity_btree_queue::CpuMask; -/// -/// let mut m: CpuMask = CpuMask::empty(); -/// m.insert(0); -/// m.insert(3); -/// assert!(m.contains(0)); -/// assert!(m.contains(3)); -/// assert!(!m.contains(1)); -/// -/// // Construct directly from a bitmask. -/// let m2: CpuMask = CpuMask::from_bits_truncate(0b1001); // CPU 0 and 3 -/// assert_eq!(m, m2); -/// -/// // Construct in a const context. -/// const M3: CpuMask = CpuMask::empty().with(0).with(3); -/// assert_eq!(M3, m); -/// ``` -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct CpuMask([u64; N]); - -impl CpuMask { - pub const fn empty() -> Self { - CpuMask([0u64; N]) - } - - /// Creates a mask with all `N * 64` bits set. - /// - /// Combine with [`masked_below`](Self::masked_below) to build "all CPUs - /// of an `n`-CPU system" masks. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::CpuMask; - /// let m: CpuMask<2> = CpuMask::all(); - /// assert!(m.contains(0)); - /// assert!(m.contains(127)); - /// ``` - pub const fn all() -> Self { - CpuMask([u64::MAX; N]) - } - - pub const fn is_empty(self) -> bool { - let mut i = 0; - while i < N { - if self.0[i] != 0 { - return false; - } - i += 1; - } - true - } - - pub const fn contains(self, cpu: usize) -> bool { - let word = cpu / 64; - let bit = cpu % 64; - word < N && (self.0[word] >> bit) & 1 == 1 - } - - pub const fn union(self, other: Self) -> Self { - let mut result = [0u64; N]; - let mut i = 0; - while i < N { - result[i] = self.0[i] | other.0[i]; - i += 1; - } - CpuMask(result) - } - - /// Sets the bit for `cpu`. Panics if `cpu >= N * 64`. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::CpuMask; - /// let mut m: CpuMask = CpuMask::empty(); - /// m.insert(0); - /// m.insert(63); - /// assert!(m.contains(0)); - /// assert!(m.contains(63)); - /// assert!(!m.contains(1)); - /// ``` - pub const fn insert(&mut self, cpu: usize) { - assert!(cpu < N * 64, "CPU index out of CpuMask range"); - self.0[cpu / 64] |= 1u64 << (cpu % 64); - } - - /// Clears the bit for `cpu`. Panics if `cpu >= N * 64`. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::CpuMask; - /// let mut m: CpuMask = CpuMask::from_bits_truncate(0b11); - /// m.remove(0); - /// assert!(!m.contains(0)); - /// assert!(m.contains(1)); - /// ``` - pub const fn remove(&mut self, cpu: usize) { - assert!(cpu < N * 64, "CPU index out of CpuMask range"); - self.0[cpu / 64] &= !(1u64 << (cpu % 64)); - } - - /// Chainable version of [`insert`](Self::insert): returns a copy of the - /// mask with the bit for `cpu` set. Panics if `cpu >= N * 64`. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::CpuMask; - /// const M: CpuMask = CpuMask::empty().with(1).with(2); - /// assert!(!M.contains(0)); - /// assert!(M.contains(1)); - /// assert!(M.contains(2)); - /// ``` - pub const fn with(mut self, cpu: usize) -> Self { - self.insert(cpu); - self - } - - /// Chainable version of [`remove`](Self::remove): returns a copy of the - /// mask with the bit for `cpu` cleared. Panics if `cpu >= N * 64`. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::CpuMask; - /// const M: CpuMask = CpuMask::all().without(0); - /// assert!(!M.contains(0)); - /// assert!(M.contains(1)); - /// ``` - pub const fn without(mut self, cpu: usize) -> Self { - self.remove(cpu); - self - } - - /// Returns a copy of the mask with every bit at index `>= n` cleared, - /// keeping only CPUs `0..n`. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::CpuMask; - /// let m: CpuMask<2> = CpuMask::all().masked_below(65); - /// assert!(m.contains(64)); - /// assert!(!m.contains(65)); - /// assert_eq!(m.word(0), u64::MAX); - /// assert_eq!(m.word(1), 1); - /// ``` - pub const fn masked_below(mut self, n: usize) -> Self { - let mut i = 0; - while i < N { - let lo = i * 64; - if lo >= n { - self.0[i] = 0; - } else { - let valid = n - lo; // 1..=64 bits of this word are valid. - if valid < 64 { - self.0[i] &= (1u64 << valid) - 1; - } - } - i += 1; - } - self - } - - /// Creates a mask with word 0 set to `bits` and all other words zeroed. - /// For `N = 1`, this is equivalent to setting the full mask. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::CpuMask; - /// // Set bits for CPU 0 and CPU 2. - /// let m: CpuMask = CpuMask::from_bits_truncate(0b0101); - /// assert!(m.contains(0)); - /// assert!(!m.contains(1)); - /// assert!(m.contains(2)); - /// ``` - pub const fn from_bits_truncate(bits: u64) -> Self { - let mut result = [0u64; N]; - if N > 0 { - result[0] = bits; - } - CpuMask(result) - } - - /// Returns word 0 of the mask as a `u64`. - /// For `N = 1`, this is the full mask. - pub const fn bits(self) -> u64 { - if N > 0 { self.0[0] } else { 0 } - } - - /// Returns word `idx` of the backing array. Panics if `idx >= N`. - pub const fn word(self, idx: usize) -> u64 { - self.0[idx] - } - - /// Sets word `idx` of the backing array. Panics if `idx >= N`. - pub const fn set_word(&mut self, idx: usize, bits: u64) { - self.0[idx] = bits; - } - - /// Iterates over the CPU indices contained in the mask, in ascending - /// order. Cost is O(number of set bits): zero words are skipped in one - /// step and each set bit is located with `trailing_zeros`. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::CpuMask; - /// let m: CpuMask<2> = CpuMask::empty().with(1).with(63).with(64).with(100); - /// let cpus: Vec = m.iter().collect(); - /// assert_eq!(cpus, [1, 63, 64, 100]); - /// ``` - pub fn iter(self) -> impl Iterator { - (0..N).flat_map(move |w| { - let mut word = self.0[w]; - core::iter::from_fn(move || { - if word == 0 { - return None; - } - let bit = word.trailing_zeros() as usize; - word &= word - 1; // Clear the lowest set bit. - Some(w * 64 + bit) - }) - }) - } - - /// True if no bit at index >= `num_cpus` is set. - fn fits_within(self, num_cpus: usize) -> bool { - if num_cpus >= N * 64 { - return true; - } - let full_words = num_cpus / 64; - let extra_bits = num_cpus % 64; - let upper_start = if extra_bits > 0 { - full_words + 1 - } else { - full_words - }; - for i in upper_start..N { - if self.0[i] != 0 { - return false; - } - } - if extra_bits > 0 && full_words < N && self.0[full_words] >> extra_bits != 0 { - return false; - } - true - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum PushError { - EmptyAffinity, - SequenceOverflow, - InvalidCpu, -} - -/// Unique B-tree key: ordered by `(priority, seq)` lexicographically. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct PriorityKey

{ - pub priority: P, - pub seq: Seq, -} - -impl Ord for PriorityKey

{ - fn cmp(&self, other: &Self) -> Ordering { - self.priority - .cmp(&other.priority) - .then_with(|| self.seq.cmp(&other.seq)) - } -} - -impl PartialOrd for PriorityKey

{ - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -#[derive(Debug)] -struct Entry { - key: PriorityKey

, - affinity: CpuMask, - value: V, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -struct NodeId(usize); - -#[derive(Debug)] -struct Node { - leaf: bool, - entries: Vec>, - // Empty when leaf; otherwise children.len() == entries.len() + 1. - children: Vec, - // OR of all affinities in this node and its descendants. - subtree_affinity: CpuMask, -} - -/// First index in `entries` whose key is >= `key`. -fn lower_bound( - entries: &[Entry], - key: &PriorityKey

, -) -> usize { - let mut lo = 0; - let mut hi = entries.len(); - while lo < hi { - let mid = (lo + hi) / 2; - if entries[mid].key < *key { - lo = mid + 1; - } else { - hi = mid; - } - } - lo -} - -/// Affinity-aware priority queue over an augmented classical B-tree of -/// minimum degree `T` (`MIN_KEYS = T - 1`, `MAX_KEYS = 2 * T - 1`). -/// -/// `N` is the number of 64-bit words in the CPU mask; the default `N = 1` -/// supports up to 64 CPUs. Use `N = 2` for up to 128 CPUs, etc. -pub struct AffinityBTreeQueue { - root: Option, - // Arena with stable ids; never compacted. Freed slots go to `free_list`. - nodes: Vec>>, - free_list: Vec, - len: usize, - next_seq: Seq, - num_cpus: usize, -} - -impl AffinityBTreeQueue { - // Only referenced by the test invariant checker. - #[cfg_attr(not(test), allow(dead_code))] - const MIN_KEYS: usize = T - 1; - const MAX_KEYS: usize = 2 * T - 1; - - /// Creates an empty queue for CPUs `0..num_cpus`. Panics if `T < 2` or - /// `num_cpus > N * 64`. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::AffinityBTreeQueue; - /// - /// // Default degree T=16, supports up to 64 CPUs (N=1). - /// let q: AffinityBTreeQueue = AffinityBTreeQueue::new(8); - /// assert_eq!(q.num_cpus(), 8); - /// assert!(q.is_empty()); - /// ``` - pub fn new(num_cpus: usize) -> Self { - assert!(T >= 2, "B-tree minimum degree T must be at least 2"); - assert!(num_cpus <= N * 64, "num_cpus exceeds CpuMask capacity N*64"); - Self { - root: None, - nodes: Vec::new(), - free_list: Vec::new(), - len: 0, - next_seq: 0, - num_cpus, - } - } - - pub fn len(&self) -> usize { - self.len - } - - pub fn is_empty(&self) -> bool { - self.len == 0 - } - - pub fn num_cpus(&self) -> usize { - self.num_cpus - } - - /// Returns the OR of the affinities of all entries currently in the - /// queue, i.e. the set of CPUs for which at least one entry is eligible. - /// O(1): reads the root node's aggregated `subtree_affinity`. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::{AffinityBTreeQueue, CpuMask}; - /// - /// let mut q: AffinityBTreeQueue = AffinityBTreeQueue::new(4); - /// assert!(q.affinity_mask().is_empty()); - /// - /// q.push(1, CpuMask::empty().with(1), ()).unwrap(); - /// q.push(2, CpuMask::empty().with(3), ()).unwrap(); - /// assert_eq!(q.affinity_mask(), CpuMask::empty().with(1).with(3)); - /// - /// q.pop_for_cpu(1).unwrap(); - /// assert_eq!(q.affinity_mask(), CpuMask::empty().with(3)); - /// ``` - pub fn affinity_mask(&self) -> CpuMask { - match self.root { - Some(r) => self.node(r).subtree_affinity, - None => CpuMask::empty(), - } - } - - fn node(&self, id: NodeId) -> &Node { - self.nodes[id.0].as_ref().expect("dangling NodeId") - } - - fn node_mut(&mut self, id: NodeId) -> &mut Node { - self.nodes[id.0].as_mut().expect("dangling NodeId") - } - - fn alloc_node(&mut self, leaf: bool) -> NodeId { - let node = Node { - leaf, - entries: Vec::new(), - children: Vec::new(), - subtree_affinity: CpuMask::empty(), - }; - if let Some(id) = self.free_list.pop() { - debug_assert!(self.nodes[id.0].is_none()); - self.nodes[id.0] = Some(node); - id - } else { - let id = NodeId(self.nodes.len()); - self.nodes.push(Some(node)); - id - } - } - - fn free_node(&mut self, id: NodeId) { - let taken = self.nodes[id.0].take(); - debug_assert!(taken.is_some()); - self.free_list.push(id); - } - - /// Recomputes `subtree_affinity` of `id` from its entries and children. - /// Must be called after every mutation of a node's entries or children. - fn pull(&mut self, id: NodeId) { - let mut mask = CpuMask::empty(); - { - let x = self.node(id); - for e in &x.entries { - mask = mask.union(e.affinity); - } - for &c in &x.children { - mask = mask.union(self.node(c).subtree_affinity); - } - } - self.node_mut(id).subtree_affinity = mask; - } - - /// Inserts an entry with the given `priority`, CPU `affinity`, and `value`. - /// - /// Returns a [`PriorityKey`] handle that uniquely identifies the entry and - /// can be passed to [`remove`](Self::remove) later. - /// - /// # Errors - /// - /// - [`PushError::EmptyAffinity`] — `affinity` has no bits set. - /// - [`PushError::InvalidCpu`] — `affinity` has a bit set for a CPU ≥ `num_cpus`. - /// - [`PushError::SequenceOverflow`] — internal sequence counter wrapped - /// (requires 2⁶⁴ prior pushes; practically impossible). - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::{AffinityBTreeQueue, CpuMask, PushError}; - /// - /// let mut q: AffinityBTreeQueue = AffinityBTreeQueue::new(4); - /// - /// let mut m = CpuMask::empty(); - /// m.insert(0); - /// let _key = q.push(10, m, 42).unwrap(); - /// assert_eq!(q.len(), 1); - /// - /// // Empty affinity mask is rejected. - /// assert_eq!(q.push(1, CpuMask::empty(), 0), Err(PushError::EmptyAffinity)); - /// - /// // CPU 4 is out of range for a 4-CPU queue. - /// let mut bad = CpuMask::empty(); - /// bad.insert(4); - /// assert_eq!(q.push(1, bad, 0), Err(PushError::InvalidCpu)); - /// ``` - pub fn push( - &mut self, - priority: P, - affinity: CpuMask, - value: V, - ) -> Result, PushError> { - if affinity.is_empty() { - return Err(PushError::EmptyAffinity); - } - if !affinity.fits_within(self.num_cpus) { - return Err(PushError::InvalidCpu); - } - let seq = self.next_seq; - let next_seq = seq.checked_add(1).ok_or(PushError::SequenceOverflow)?; - self.next_seq = next_seq; - - let key = PriorityKey { priority, seq }; - let ret = key.clone(); - let entry = Entry { - key, - affinity, - value, - }; - - match self.root { - None => { - let r = self.alloc_node(true); - self.node_mut(r).entries.push(entry); - self.pull(r); - self.root = Some(r); - } - Some(r) => { - if self.node(r).entries.len() == Self::MAX_KEYS { - let s = self.alloc_node(false); - self.node_mut(s).children.push(r); - self.root = Some(s); - self.split_child(s, 0); - self.insert_non_full(s, entry); - } else { - self.insert_non_full(r, entry); - } - } - } - self.len += 1; - Ok(ret) - } - - /// Splits the full child `parent.children[index]` around its median entry. - fn split_child(&mut self, parent_id: NodeId, index: usize) { - let y_id = self.node(parent_id).children[index]; - let y_leaf = self.node(y_id).leaf; - let z_id = self.alloc_node(y_leaf); - - let (median, right_entries, right_children) = { - let y = self.node_mut(y_id); - debug_assert_eq!(y.entries.len(), Self::MAX_KEYS); - let right_entries = y.entries.split_off(T); - let median = y.entries.pop().expect("full node has a median"); - let right_children = if y_leaf { - Vec::new() - } else { - y.children.split_off(T) - }; - (median, right_entries, right_children) - }; - - { - let z = self.node_mut(z_id); - z.entries = right_entries; - z.children = right_children; - } - { - let p = self.node_mut(parent_id); - p.entries.insert(index, median); - p.children.insert(index + 1, z_id); - } - - self.pull(y_id); - self.pull(z_id); - self.pull(parent_id); - } - - fn insert_non_full(&mut self, node_id: NodeId, entry: Entry) { - debug_assert!(self.node(node_id).entries.len() < Self::MAX_KEYS); - - if self.node(node_id).leaf { - let i = lower_bound(&self.node(node_id).entries, &entry.key); - debug_assert!( - i >= self.node(node_id).entries.len() - || self.node(node_id).entries[i].key != entry.key, - "duplicate key" - ); - self.node_mut(node_id).entries.insert(i, entry); - self.pull(node_id); - return; - } - - let mut i = lower_bound(&self.node(node_id).entries, &entry.key); - let child_id = self.node(node_id).children[i]; - if self.node(child_id).entries.len() == Self::MAX_KEYS { - self.split_child(node_id, i); - if entry.key > self.node(node_id).entries[i].key { - i += 1; - } else { - debug_assert!( - entry.key != self.node(node_id).entries[i].key, - "duplicate key" - ); - } - } - let next_child = self.node(node_id).children[i]; - self.insert_non_full(next_child, entry); - self.pull(node_id); - } - - /// Smallest key whose affinity contains `cpu`, or `None`. - /// - /// Does not remove the entry; use [`pop_for_cpu`](Self::pop_for_cpu) to - /// remove it. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::{AffinityBTreeQueue, CpuMask}; - /// - /// let mut q: AffinityBTreeQueue = AffinityBTreeQueue::new(2); - /// assert_eq!(q.peek_key_for_cpu(0), None); // empty queue - /// - /// let mut m = CpuMask::empty(); - /// m.insert(1); // CPU 1 only - /// let key = q.push(5, m, 99).unwrap(); - /// - /// assert_eq!(q.peek_key_for_cpu(1), Some(key.clone())); // visible on CPU 1 - /// assert_eq!(q.peek_key_for_cpu(0), None); // not visible on CPU 0 - /// assert_eq!(q.len(), 1); // entry was not consumed - /// ``` - pub fn peek_key_for_cpu(&self, cpu: usize) -> Option> { - if cpu >= self.num_cpus { - return None; - } - let r = self.root?; - if !self.node(r).subtree_affinity.contains(cpu) { - return None; - } - Some(self.find_first_eligible_key(r, cpu)) - } - - /// Precondition: `node(node_id).subtree_affinity.contains(cpu)`. - /// Scans in logical order child[0], entry[0], child[1], ... so the first - /// eligible entry found is the minimum eligible key in the subtree. - fn find_first_eligible_key(&self, node_id: NodeId, cpu: usize) -> PriorityKey

{ - let mut cur = node_id; - - 'descend: loop { - let x = self.node(cur); - debug_assert!(x.subtree_affinity.contains(cpu)); - let n = x.entries.len(); - - for i in 0..n { - if !x.leaf { - let child_id = x.children[i]; - if self.node(child_id).subtree_affinity.contains(cpu) { - cur = child_id; - continue 'descend; - } - } - if x.entries[i].affinity.contains(cpu) { - return x.entries[i].key.clone(); - } - } - - if !x.leaf { - let child_id = x.children[n]; - if self.node(child_id).subtree_affinity.contains(cpu) { - cur = child_id; - continue 'descend; - } - } - - panic!("subtree_affinity invariant is broken"); - } - } - - /// Removes and returns the highest-priority entry runnable on `cpu`. - /// - /// Returns `None` if the queue is empty or no entry's affinity includes - /// `cpu`. - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::{AffinityBTreeQueue, CpuMask}; - /// - /// let mut q: AffinityBTreeQueue = AffinityBTreeQueue::new(2); - /// - /// let mut cpu0 = CpuMask::empty(); - /// cpu0.insert(0); - /// let mut cpu1 = CpuMask::empty(); - /// cpu1.insert(1); - /// - /// q.push(1, cpu0, "only cpu0").unwrap(); - /// q.push(2, cpu1, "only cpu1").unwrap(); - /// - /// // CPU 0 only dequeues its own entry. - /// let (_, affinity, value) = q.pop_for_cpu(0).unwrap(); - /// assert_eq!(value, "only cpu0"); - /// assert!(affinity.contains(0)); - /// - /// // No more entries runnable on CPU 0. - /// assert!(q.pop_for_cpu(0).is_none()); - /// ``` - pub fn pop_for_cpu(&mut self, cpu: usize) -> Option<(PriorityKey

, CpuMask, V)> { - let key = self.peek_key_for_cpu(cpu)?; - let removed = self.remove(&key); - debug_assert!(removed.is_some()); - debug_assert!(removed.as_ref().is_none_or(|(_, a, _)| a.contains(cpu))); - removed - } - - /// Removes the entry identified by `key` and returns `(key, affinity, value)`, - /// or `None` if no such entry exists. - /// - /// The key is obtained from [`push`](Self::push) or - /// [`peek_key_for_cpu`](Self::peek_key_for_cpu). - /// - /// # Examples - /// - /// ``` - /// use affinity_btree_queue::{AffinityBTreeQueue, CpuMask}; - /// - /// let mut q: AffinityBTreeQueue = AffinityBTreeQueue::new(1); - /// let mut m = CpuMask::empty(); - /// m.insert(0); - /// - /// let key = q.push(7, m, 42).unwrap(); - /// let (k, _, v) = q.remove(&key).unwrap(); - /// assert_eq!(k, key); - /// assert_eq!(v, 42); - /// - /// // Removing the same key a second time returns None. - /// assert!(q.remove(&key).is_none()); - /// ``` - pub fn remove(&mut self, key: &PriorityKey

) -> Option<(PriorityKey

, CpuMask, V)> { - let root = self.root?; - let removed = self.delete_from_node(root, key); - - if let Some(r) = self.root - && self.node(r).entries.is_empty() - { - if self.node(r).leaf { - self.free_node(r); - self.root = None; - } else { - let new_root = self.node(r).children[0]; - self.free_node(r); - self.root = Some(new_root); - } - } - - if removed.is_some() { - self.len -= 1; - } - removed.map(|e| (e.key, e.affinity, e.value)) - } - - fn delete_from_node( - &mut self, - node_id: NodeId, - key: &PriorityKey

, - ) -> Option> { - let i = lower_bound(&self.node(node_id).entries, key); - let leaf = self.node(node_id).leaf; - let found_here = - i < self.node(node_id).entries.len() && self.node(node_id).entries[i].key == *key; - - if found_here { - if leaf { - let entry = self.node_mut(node_id).entries.remove(i); - self.pull(node_id); - Some(entry) - } else { - let removed = self.delete_internal_entry(node_id, i); - self.pull(node_id); - removed - } - } else if leaf { - None - } else { - let child_index = self.ensure_child_has_at_least_t(node_id, i); - let child_id = self.node(node_id).children[child_index]; - let removed = self.delete_from_node(child_id, key); - self.pull(node_id); - removed - } - } - - /// Deletes `entries[index]` of the internal node `node_id` using the - /// standard predecessor / successor / merge cases. - fn delete_internal_entry(&mut self, node_id: NodeId, index: usize) -> Option> { - let left_id = self.node(node_id).children[index]; - let right_id = self.node(node_id).children[index + 1]; - - if self.node(left_id).entries.len() >= T { - let predecessor = self.remove_max_entry(left_id); - let old = mem::replace(&mut self.node_mut(node_id).entries[index], predecessor); - self.pull(node_id); - return Some(old); - } - - if self.node(right_id).entries.len() >= T { - let successor = self.remove_min_entry(right_id); - let old = mem::replace(&mut self.node_mut(node_id).entries[index], successor); - self.pull(node_id); - return Some(old); - } - - let target_key = self.node(node_id).entries[index].key.clone(); - let merged_id = self.merge_children_around_entry(node_id, index); - let removed = self.delete_from_node(merged_id, &target_key); - self.pull(node_id); - removed - } - - fn remove_min_entry(&mut self, node_id: NodeId) -> Entry { - if self.node(node_id).leaf { - let entry = self.node_mut(node_id).entries.remove(0); - self.pull(node_id); - return entry; - } - let child_index = self.ensure_child_has_at_least_t(node_id, 0); - debug_assert_eq!(child_index, 0); - let child_id = self.node(node_id).children[child_index]; - let entry = self.remove_min_entry(child_id); - self.pull(node_id); - entry - } - - fn remove_max_entry(&mut self, node_id: NodeId) -> Entry { - if self.node(node_id).leaf { - let entry = self - .node_mut(node_id) - .entries - .pop() - .expect("remove_max_entry on empty node"); - self.pull(node_id); - return entry; - } - let last_child_index = self.node(node_id).children.len() - 1; - // If the rightmost child gets merged into its left sibling, the - // returned index is last_child_index - 1; always use the return value. - let child_index = self.ensure_child_has_at_least_t(node_id, last_child_index); - let child_id = self.node(node_id).children[child_index]; - let entry = self.remove_max_entry(child_id); - self.pull(node_id); - entry - } - - /// Ensures `parent.children[index]` has at least `T` entries before a - /// descent, borrowing from a sibling or merging. Returns the (possibly - /// shifted) index of the child to descend into. - fn ensure_child_has_at_least_t(&mut self, parent_id: NodeId, index: usize) -> usize { - let child_id = self.node(parent_id).children[index]; - if self.node(child_id).entries.len() >= T { - return index; - } - - if index > 0 { - let left_id = self.node(parent_id).children[index - 1]; - if self.node(left_id).entries.len() >= T { - self.borrow_from_left(parent_id, index); - return index; - } - } - - if index + 1 < self.node(parent_id).children.len() { - let right_id = self.node(parent_id).children[index + 1]; - if self.node(right_id).entries.len() >= T { - self.borrow_from_right(parent_id, index); - return index; - } - } - - if index > 0 { - self.merge_children_around_entry(parent_id, index - 1); - index - 1 - } else { - self.merge_children_around_entry(parent_id, index); - index - } - } - - /// Rotates the separator down into `children[index]` and the left - /// sibling's max entry up into the parent. - fn borrow_from_left(&mut self, parent_id: NodeId, index: usize) { - let left_id = self.node(parent_id).children[index - 1]; - let child_id = self.node(parent_id).children[index]; - - let up_entry = self - .node_mut(left_id) - .entries - .pop() - .expect("borrow_from_left from empty sibling"); - let down_entry = mem::replace(&mut self.node_mut(parent_id).entries[index - 1], up_entry); - self.node_mut(child_id).entries.insert(0, down_entry); - - if !self.node(child_id).leaf { - let moved_child = self - .node_mut(left_id) - .children - .pop() - .expect("internal sibling has children"); - self.node_mut(child_id).children.insert(0, moved_child); - } - - self.pull(left_id); - self.pull(child_id); - self.pull(parent_id); - } - - /// Rotates the separator down into `children[index]` and the right - /// sibling's min entry up into the parent. - fn borrow_from_right(&mut self, parent_id: NodeId, index: usize) { - let child_id = self.node(parent_id).children[index]; - let right_id = self.node(parent_id).children[index + 1]; - - let up_entry = self.node_mut(right_id).entries.remove(0); - let down_entry = mem::replace(&mut self.node_mut(parent_id).entries[index], up_entry); - self.node_mut(child_id).entries.push(down_entry); - - if !self.node(child_id).leaf { - let moved_child = self.node_mut(right_id).children.remove(0); - self.node_mut(child_id).children.push(moved_child); - } - - self.pull(child_id); - self.pull(right_id); - self.pull(parent_id); - } - - /// Merges `children[index]`, `entries[index]`, and `children[index + 1]` - /// into the left child. Frees the right child and returns the merged id. - fn merge_children_around_entry(&mut self, parent_id: NodeId, index: usize) -> NodeId { - let left_id = self.node(parent_id).children[index]; - let right_id = self.node(parent_id).children[index + 1]; - - let separator = self.node_mut(parent_id).entries.remove(index); - self.node_mut(parent_id).children.remove(index + 1); - - let right = self.nodes[right_id.0].take().expect("dangling NodeId"); - self.free_list.push(right_id); - - let left = self.node_mut(left_id); - left.entries.push(separator); - left.entries.extend(right.entries); - left.children.extend(right.children); - - self.pull(left_id); - self.pull(parent_id); - left_id - } -} - -#[cfg(test)] -impl AffinityBTreeQueue { - pub(crate) fn debug_set_next_seq(&mut self, seq: Seq) { - self.next_seq = seq; - } - - fn for_each_entry_in_order)>(&self, mut f: F) { - enum Item { - Visit(NodeId), - Emit(NodeId, usize), - } - - let Some(root) = self.root else { return }; - let mut stack = Vec::new(); - stack.push(Item::Visit(root)); - - while let Some(item) = stack.pop() { - match item { - Item::Visit(id) => { - let x = self.node(id); - if x.leaf { - for e in &x.entries { - f(e); - } - } else { - let n = x.entries.len(); - // Push in reverse so child[0] is processed first: - // child[0], entry[0], child[1], ..., entry[n-1], child[n]. - stack.push(Item::Visit(x.children[n])); - let mut i = n; - while i > 0 { - i -= 1; - stack.push(Item::Emit(id, i)); - stack.push(Item::Visit(x.children[i])); - } - } - } - Item::Emit(id, i) => f(&self.node(id).entries[i]), - } - } - } - - pub(crate) fn debug_count_entries(&self) -> usize { - let mut count = 0; - self.for_each_entry_in_order(|_| count += 1); - count - } - - pub(crate) fn debug_all_keys_in_order(&self) -> Vec> { - let mut out = Vec::new(); - self.for_each_entry_in_order(|e| out.push(e.key.clone())); - out - } - - /// Asserts every structural, key-order, augmentation, and arena invariant. - pub(crate) fn check_invariants(&self) { - for id in &self.free_list { - assert!( - self.nodes[id.0].is_none(), - "free_list entry points to a live node" - ); - } - - let Some(root) = self.root else { - assert_eq!(self.len, 0, "empty tree must have len 0"); - assert!( - self.nodes.iter().all(|n| n.is_none()), - "empty tree must have no live nodes" - ); - assert_eq!(self.free_list.len(), self.nodes.len()); - return; - }; - assert!(self.len > 0, "non-empty root with len 0"); - - struct WorkItem

{ - id: NodeId, - is_root: bool, - min_key: Option>, - max_key: Option>, - depth: usize, - } - - let mut total_entries = 0usize; - let mut total_nodes = 0usize; - let mut leaf_depth: Option = None; - - let mut stack: Vec> = Vec::new(); - stack.push(WorkItem { - id: root, - is_root: true, - min_key: None, - max_key: None, - depth: 0, - }); - - while let Some(item) = stack.pop() { - let x = self.node(item.id); - total_nodes += 1; - total_entries += x.entries.len(); - - assert!(x.entries.len() <= Self::MAX_KEYS, "node overflow"); - if item.is_root { - assert!(!x.entries.is_empty(), "root must be non-empty"); - } else { - assert!(x.entries.len() >= Self::MIN_KEYS, "non-root underflow"); - } - - for w in x.entries.windows(2) { - assert!(w[0].key < w[1].key, "node entries not strictly sorted"); - } - for e in &x.entries { - if let Some(m) = &item.min_key { - assert!(*m < e.key, "entry below child key range"); - } - if let Some(m) = &item.max_key { - assert!(e.key < *m, "entry above child key range"); - } - assert!(!e.affinity.is_empty(), "entry with empty affinity"); - assert!( - e.affinity.fits_within(self.num_cpus), - "entry affinity outside num_cpus" - ); - } - - let mut mask = CpuMask::empty(); - for e in &x.entries { - mask = mask.union(e.affinity); - } - - if x.leaf { - assert!(x.children.is_empty(), "leaf with children"); - match leaf_depth { - None => leaf_depth = Some(item.depth), - Some(d) => assert_eq!(d, item.depth, "leaves at different depths"), - } - } else { - assert_eq!( - x.children.len(), - x.entries.len() + 1, - "internal node child count" - ); - for ci in 0..x.children.len() { - let child_id = x.children[ci]; - mask = mask.union(self.node(child_id).subtree_affinity); - stack.push(WorkItem { - id: child_id, - is_root: false, - min_key: if ci == 0 { - item.min_key.clone() - } else { - Some(x.entries[ci - 1].key.clone()) - }, - max_key: if ci == x.entries.len() { - item.max_key.clone() - } else { - Some(x.entries[ci].key.clone()) - }, - depth: item.depth + 1, - }); - } - } - - assert_eq!(mask, x.subtree_affinity, "stale subtree_affinity"); - } - - assert_eq!(total_entries, self.len, "entry count != len"); - let live = self.nodes.iter().filter(|n| n.is_some()).count(); - assert_eq!(live, total_nodes, "unreachable live nodes in arena"); - assert_eq!(live + self.free_list.len(), self.nodes.len()); - - { - let keys = self.debug_all_keys_in_order(); - for w in keys.windows(2) { - assert!(w[0] < w[1], "in-order keys not strictly sorted"); - } - } - } -} - -#[cfg(test)] -impl AffinityBTreeQueue { - pub(crate) fn debug_all_entries_in_order(&self) -> Vec<(PriorityKey

, CpuMask, V)> { - let mut out = Vec::new(); - self.for_each_entry_in_order(|e| out.push((e.key.clone(), e.affinity, e.value.clone()))); - out - } -} - -/// Deliberately simple Vec-based reference model used by tests to check the -/// B-tree's observable behavior. -#[cfg(test)] -#[derive(Clone, Debug, PartialEq, Eq)] -struct RefEntry { - key: PriorityKey

, - affinity: CpuMask, - value: V, -} - -#[cfg(test)] -#[derive(Clone, Debug)] -struct RefModel { - entries: Vec>, - next_seq: Seq, - num_cpus: usize, -} - -#[cfg(test)] -impl RefModel { - fn new(num_cpus: usize) -> Self { - Self { - entries: Vec::new(), - next_seq: 0, - num_cpus, - } - } - - fn push( - &mut self, - priority: P, - affinity: CpuMask, - value: V, - ) -> Result, PushError> { - if affinity.is_empty() { - return Err(PushError::EmptyAffinity); - } - if !affinity.fits_within(self.num_cpus) { - return Err(PushError::InvalidCpu); - } - let seq = self.next_seq; - let next_seq = seq.checked_add(1).ok_or(PushError::SequenceOverflow)?; - self.next_seq = next_seq; - - let key = PriorityKey { priority, seq }; - self.entries.push(RefEntry { - key: key.clone(), - affinity, - value, - }); - Ok(key) - } - - fn peek_key_for_cpu(&self, cpu: usize) -> Option> { - if cpu >= self.num_cpus { - return None; - } - self.entries - .iter() - .filter(|e| e.affinity.contains(cpu)) - .map(|e| e.key.clone()) - .min() - } - - fn pop_for_cpu(&mut self, cpu: usize) -> Option<(PriorityKey

, CpuMask, V)> { - let key = self.peek_key_for_cpu(cpu)?; - let index = self.entries.iter().position(|e| e.key == key).unwrap(); - let entry = self.entries.remove(index); - Some((entry.key, entry.affinity, entry.value)) - } - - fn remove(&mut self, key: &PriorityKey

) -> Option<(PriorityKey

, CpuMask, V)> { - let index = self.entries.iter().position(|e| e.key == *key)?; - let entry = self.entries.remove(index); - Some((entry.key, entry.affinity, entry.value)) - } - - fn len(&self) -> usize { - self.entries.len() - } - - fn is_empty(&self) -> bool { - self.entries.is_empty() - } - - fn entries_in_order(&self) -> Vec<(PriorityKey

, CpuMask, V)> { - let mut out: Vec<_> = self - .entries - .iter() - .map(|e| (e.key.clone(), e.affinity, e.value.clone())) - .collect(); - out.sort_by(|a, b| a.0.cmp(&b.0)); - out - } -} - -#[cfg(test)] -mod tests; diff --git a/affinity_btree_queue/src/tests.rs b/affinity_btree_queue/src/tests.rs deleted file mode 100644 index e316f5ec2..000000000 --- a/affinity_btree_queue/src/tests.rs +++ /dev/null @@ -1,1080 +0,0 @@ -use super::*; -use alloc::vec::Vec; -use core::cmp::Reverse; - -/// Degree-2 queue: every structural case (split, borrow, merge, root shrink) -/// is reachable with a handful of entries. -type Q2 = AffinityBTreeQueue; -/// Production-degree queue. -type Q16 = AffinityBTreeQueue; - -fn mask(cpus: &[usize]) -> CpuMask { - let mut m = CpuMask::empty(); - for &c in cpus { - m.insert(c); - } - m -} - -fn all_cpus(num_cpus: usize) -> CpuMask { - let mut m = CpuMask::empty(); - for cpu in 0..num_cpus { - m.insert(cpu); - } - m -} - -struct XorShift64 { - state: u64, -} - -impl XorShift64 { - fn new(seed: u64) -> Self { - Self { - state: seed.wrapping_mul(0x9E37_79B9_7F4A_7C15).max(1), - } - } - - fn next(&mut self) -> u64 { - let mut x = self.state; - x ^= x << 13; - x ^= x >> 7; - x ^= x << 17; - self.state = x; - x - } -} - -// --------------------------------------------------------------------------- -// CpuMask -// --------------------------------------------------------------------------- - -#[test] -fn cpu_mask_basics() { - let mut m = CpuMask::empty(); - assert!(m.is_empty()); - assert!(!m.contains(0)); - m.insert(0); - m.insert(63); - assert!(m.contains(0)); - assert!(m.contains(63)); - assert!(!m.contains(1)); - assert!(!m.contains(64)); - assert!(!m.contains(1000)); - assert_eq!(m.union(mask(&[1])), mask(&[0, 1, 63])); - assert!(mask(&[3]).fits_within(4)); - assert!(!mask(&[4]).fits_within(4)); - assert!(mask(&[63]).fits_within(64)); -} - -#[test] -fn cpu_mask_wide() { - let mut m = CpuMask::<2>::empty(); - assert!(m.is_empty()); - m.insert(0); - m.insert(63); - m.insert(64); - m.insert(127); - assert!(m.contains(0)); - assert!(m.contains(63)); - assert!(m.contains(64)); - assert!(m.contains(127)); - assert!(!m.contains(1)); - assert!(!m.contains(128)); - assert!(!m.contains(10000)); - - assert!(m.fits_within(128)); - assert!(!m.fits_within(127)); // bit 127 is set - assert!(!m.fits_within(64)); // bits 64+ are set - assert!(!m.fits_within(63)); // bits 63+ are set - - let mut a = CpuMask::<2>::empty(); - a.insert(65); - let mut b = CpuMask::<2>::empty(); - b.insert(70); - let c = a.union(b); - assert!(c.contains(65)); - assert!(c.contains(70)); - assert!(!c.contains(66)); - - let mut w = CpuMask::<2>::empty(); - w.set_word(0, 0xAB); - w.set_word(1, 0xCD); - assert_eq!(w.word(0), 0xAB); - assert_eq!(w.word(1), 0xCD); -} - -#[test] -fn cpu_mask_from_bits_truncate_and_bits() { - let m: CpuMask = CpuMask::from_bits_truncate(0b0101); - assert!(m.contains(0)); - assert!(!m.contains(1)); - assert!(m.contains(2)); - assert_eq!(m.bits(), 0b0101); - - let zero: CpuMask = CpuMask::from_bits_truncate(0); - assert!(zero.is_empty()); - assert_eq!(zero.bits(), 0); -} - -#[test] -fn cpu_mask_const_builders() { - // All constructors must be usable in const contexts. - const M: CpuMask<2> = CpuMask::empty().with(1).with(64).without(1); - assert!(!M.contains(1)); - assert!(M.contains(64)); - - const ALL: CpuMask<2> = CpuMask::all(); - assert!(ALL.contains(0)); - assert!(ALL.contains(127)); - - const WORKERS: CpuMask<2> = CpuMask::all().masked_below(4).without(0); - assert_eq!(WORKERS, CpuMask::empty().with(1).with(2).with(3)); -} - -#[test] -fn cpu_mask_remove() { - let mut m = CpuMask::<2>::empty(); - m.insert(0); - m.insert(64); - m.remove(0); - assert!(!m.contains(0)); - assert!(m.contains(64)); - m.remove(64); - assert!(m.is_empty()); - // Removing an already-clear bit is a no-op. - m.remove(1); - assert!(m.is_empty()); -} - -#[test] -fn cpu_mask_masked_below_word_boundaries() { - let all = CpuMask::<2>::all(); - - assert!(all.masked_below(0).is_empty()); - - assert_eq!(all.masked_below(1), CpuMask::empty().with(0)); - - let below63 = all.masked_below(63); - assert!(below63.contains(62)); - assert!(!below63.contains(63)); - - let below64 = all.masked_below(64); - assert_eq!(below64.word(0), u64::MAX); - assert_eq!(below64.word(1), 0); - - let below65 = all.masked_below(65); - assert_eq!(below65.word(0), u64::MAX); - assert_eq!(below65.word(1), 1); - - // n >= N * 64 keeps every bit. - assert_eq!(all.masked_below(128), all); - assert_eq!(all.masked_below(1000), all); - - // Already-clear bits stay clear. - let sparse = CpuMask::<2>::empty().with(1).with(70); - assert_eq!(sparse.masked_below(70), CpuMask::empty().with(1)); -} - -#[test] -fn cpu_mask_iter() { - // Empty mask yields nothing. - assert_eq!(CpuMask::<2>::empty().iter().count(), 0); - - // Ascending order across word boundaries; zero words are skipped. - let m = CpuMask::<2>::empty().with(1).with(63).with(64).with(127); - let cpus: Vec = m.iter().collect(); - assert_eq!(cpus, [1, 63, 64, 127]); - - // A mask whose first word is zero. - let hi = CpuMask::<2>::empty().with(100); - assert_eq!(hi.iter().collect::>(), [100]); - - // Full mask yields every index. - let all: Vec = CpuMask::<2>::all().iter().collect(); - assert_eq!(all.len(), 128); - assert_eq!(all[0], 0); - assert_eq!(all[127], 127); -} - -#[test] -fn fits_within_upper_word_all_zero() { - // word 0 has bits; word 1 is entirely zero. - // fits_within(64): full_words=1, extra_bits=0, upper_start=1 - // Loop for i in 1..2: self.0[1]==0 → no early return, falls through line 199. - let mut m = CpuMask::<2>::empty(); - m.insert(0); - m.insert(63); - assert!(m.fits_within(64)); -} - -#[test] -fn priority_key_ordering() { - let a = PriorityKey { - priority: 1u32, - seq: 5, - }; - let b = PriorityKey { - priority: 2u32, - seq: 0, - }; - let c = PriorityKey { - priority: 1u32, - seq: 6, - }; - assert!(a < b); - assert!(a < c); - assert!(c < b); - assert_eq!(a, a.clone()); -} - -// --------------------------------------------------------------------------- -// Basic operations -// --------------------------------------------------------------------------- - -#[test] -fn queue_num_cpus_getter() { - let q: Q16 = AffinityBTreeQueue::new(8); - assert_eq!(q.num_cpus(), 8); - let q2: Q2 = AffinityBTreeQueue::new(2); - assert_eq!(q2.num_cpus(), 2); -} - -#[test] -fn queue_affinity_mask_tracks_entries() { - let mut q = Q16::new(4); - assert!(q.affinity_mask().is_empty()); - - let k1 = q.push(10, mask(&[1]), 100).unwrap(); - assert_eq!(q.affinity_mask(), mask(&[1])); - - q.push(20, mask(&[2, 3]), 200).unwrap(); - assert_eq!(q.affinity_mask(), mask(&[1, 2, 3])); - - // Two entries share CPU 1; removing one must keep the bit set. - q.push(30, mask(&[1]), 300).unwrap(); - q.remove(&k1).unwrap(); - assert_eq!(q.affinity_mask(), mask(&[1, 2, 3])); - - assert!(q.pop_for_cpu(1).is_some()); - assert_eq!(q.affinity_mask(), mask(&[2, 3])); - - assert!(q.pop_for_cpu(3).is_some()); - assert!(q.affinity_mask().is_empty()); - q.check_invariants(); -} - -#[test] -fn queue_affinity_mask_with_deep_tree() { - // Force splits (T=2) so the mask is aggregated across multiple levels. - let mut q = Q2::new(2); - for i in 0..64u32 { - let cpu = (i % 2) as usize; - q.push(i, mask(&[cpu]), i).unwrap(); - } - assert_eq!(q.affinity_mask(), mask(&[0, 1])); - - // Drain every CPU-0 entry; only CPU-1 entries remain. - while q.pop_for_cpu(0).is_some() {} - assert_eq!(q.affinity_mask(), mask(&[1])); - - while q.pop_for_cpu(1).is_some() {} - assert!(q.affinity_mask().is_empty()); - q.check_invariants(); -} - -#[test] -fn empty_tree_pop_for_cpu() { - let mut q = Q16::new(4); - assert_eq!(q.len(), 0); - assert!(q.is_empty()); - for cpu in 0..4 { - assert_eq!(q.peek_key_for_cpu(cpu), None); - assert_eq!(q.pop_for_cpu(cpu), None); - } - assert_eq!(q.pop_for_cpu(99), None); - q.check_invariants(); -} - -#[test] -fn single_insert_and_pop() { - let mut q = Q16::new(4); - let key = q.push(10, mask(&[1, 2]), 100).unwrap(); - q.check_invariants(); - assert_eq!(q.len(), 1); - assert_eq!(q.peek_key_for_cpu(0), None); - assert_eq!(q.peek_key_for_cpu(1), Some(key.clone())); - assert_eq!(q.pop_for_cpu(2), Some((key, mask(&[1, 2]), 100))); - assert!(q.is_empty()); - q.check_invariants(); -} - -#[test] -fn single_insert_and_remove() { - let mut q = Q16::new(4); - let key = q.push(10, mask(&[0]), 7).unwrap(); - assert_eq!(q.remove(&key), Some((key.clone(), mask(&[0]), 7))); - assert_eq!(q.remove(&key), None); - assert!(q.is_empty()); - q.check_invariants(); -} - -#[test] -fn same_priority_multiple_sequence_numbers() { - let mut q = Q2::new(2); - let mut keys = Vec::new(); - for i in 0..8u32 { - let key = q.push(7, mask(&[0]), i).unwrap(); - keys.push(key); - q.check_invariants(); - } - // All same priority: dequeue order must follow seq (FIFO). - for (i, expected) in keys.into_iter().enumerate() { - let (key, _, value) = q.pop_for_cpu(0).unwrap(); - assert_eq!(key, expected); - assert_eq!(value, i as u32); - q.check_invariants(); - } - assert!(q.is_empty()); -} - -#[test] -fn sequence_number_uniqueness() { - let mut q = Q16::new(1); - let mut seqs = Vec::new(); - for _ in 0..50 { - let key = q.push(3, mask(&[0]), 0).unwrap(); - assert!(!seqs.contains(&key.seq)); - seqs.push(key.seq); - } - for w in seqs.windows(2) { - assert!(w[0] < w[1]); - } -} - -#[test] -fn ascending_priority_insertion() { - let mut q = Q16::new(4); - for p in 0..300u32 { - q.push(p, all_cpus(4), p).unwrap(); - q.check_invariants(); - } - for p in 0..300u32 { - let (key, _, value) = q.pop_for_cpu((p as usize) % 4).unwrap(); - assert_eq!(key.priority, p); - assert_eq!(value, p); - q.check_invariants(); - } - assert!(q.is_empty()); -} - -#[test] -fn descending_priority_insertion() { - let mut q = Q16::new(4); - for p in (0..300u32).rev() { - q.push(p, all_cpus(4), p).unwrap(); - q.check_invariants(); - } - for p in 0..300u32 { - let (key, _, value) = q.pop_for_cpu(0).unwrap(); - assert_eq!(key.priority, p); - assert_eq!(value, p); - q.check_invariants(); - } -} - -#[test] -fn random_priority_insertion() { - let mut rng = XorShift64::new(42); - let mut q = Q16::new(4); - let mut model = RefModel::::new(4); - for i in 0..500u32 { - let p = (rng.next() % 100) as u32; - assert_eq!(q.push(p, all_cpus(4), i), model.push(p, all_cpus(4), i)); - q.check_invariants(); - } - while !model.is_empty() { - assert_eq!(q.pop_for_cpu(1), model.pop_for_cpu(1)); - q.check_invariants(); - } - assert!(q.is_empty()); -} - -// --------------------------------------------------------------------------- -// Structural cases (degree 2 makes shapes deterministic) -// --------------------------------------------------------------------------- - -/// Pushes priorities 10, 20, 30, 40 into a T=2 queue, producing the shape -/// root=[20], left=[10], right=[30,40]. Returns the key of priority 20. -fn build_split_tree(q: &mut Q2) -> PriorityKey { - let mut key20 = None; - for p in [10u32, 20, 30, 40] { - let key = q.push(p, mask(&[0]), p).unwrap(); - if p == 20 { - key20 = Some(key); - } - q.check_invariants(); - } - key20.unwrap() -} - -fn pop_all_priorities(q: &mut Q2, cpu: usize) -> Vec { - let mut out = Vec::new(); - while let Some((key, _, _)) = q.pop_for_cpu(cpu) { - out.push(key.priority); - q.check_invariants(); - } - out -} - -#[test] -fn root_split() { - let mut q = Q2::new(1); - build_split_tree(&mut q); - assert_eq!(q.len(), 4); - assert_eq!(pop_all_priorities(&mut q, 0), [10, 20, 30, 40]); -} - -#[test] -fn multiple_level_split() { - let mut q = Q2::new(2); - for p in 0..60u32 { - q.push(p, mask(&[(p % 2) as usize]), p).unwrap(); - q.check_invariants(); - } - assert_eq!(q.len(), 60); - let evens = pop_all_priorities(&mut q, 0); - assert_eq!(evens, (0..60).filter(|p| p % 2 == 0).collect::>()); - q.check_invariants(); - - let odds = pop_all_priorities(&mut q, 1); - assert_eq!(odds, (0..60).filter(|p| p % 2 == 1).collect::>()); - assert!(q.is_empty()); -} - -#[test] -fn multiple_level_split_production_degree() { - let mut q = Q16::new(4); - // > MAX_KEYS + (MAX_KEYS + 1) * MAX_KEYS = 1023 entries forces height >= 3. - for p in 0..1200u32 { - q.push(p, all_cpus(4), p).unwrap(); - if p % 50 == 0 { - q.check_invariants(); - } - } - q.check_invariants(); - assert_eq!(q.len(), 1200); - for p in 0..1200u32 { - let (key, _, _) = q.pop_for_cpu(3).unwrap(); - assert_eq!(key.priority, p); - if p % 50 == 0 { - q.check_invariants(); - } - } - q.check_invariants(); -} - -#[test] -fn root_shrink_after_deletion() { - let mut q = Q2::new(1); - let mut keys = Vec::new(); - for p in 0..20u32 { - keys.push(q.push(p, mask(&[0]), p).unwrap()); - } - q.check_invariants(); - for key in keys { - assert!(q.remove(&key).is_some()); - q.check_invariants(); - } - assert!(q.is_empty()); - assert_eq!(q.peek_key_for_cpu(0), None); -} - -#[test] -fn delete_from_leaf() { - let mut q = Q2::new(1); - let k10 = q.push(10, mask(&[0]), 10).unwrap(); - let k20 = q.push(20, mask(&[0]), 20).unwrap(); - let k30 = q.push(30, mask(&[0]), 30).unwrap(); - // Root is still a single leaf [10, 20, 30]. - assert_eq!(q.remove(&k20), Some((k20, mask(&[0]), 20))); - q.check_invariants(); - assert_eq!(q.remove(&k10), Some((k10, mask(&[0]), 10))); - q.check_invariants(); - assert_eq!(q.remove(&k30), Some((k30, mask(&[0]), 30))); - q.check_invariants(); - assert!(q.is_empty()); -} - -#[test] -fn delete_internal_entry_using_predecessor() { - let mut q = Q2::new(1); - let key20 = build_split_tree(&mut q); - // Push 5 so the left child [5,10] has T=2 entries: removing the internal - // key 20 must take the predecessor (10) from the left child. - q.push(5, mask(&[0]), 5).unwrap(); - q.check_invariants(); - assert!(q.remove(&key20).is_some()); - q.check_invariants(); - assert_eq!(pop_all_priorities(&mut q, 0), [5, 10, 30, 40]); -} - -#[test] -fn delete_internal_entry_using_successor() { - let mut q = Q2::new(1); - let key20 = build_split_tree(&mut q); - // Shape: root=[20], left=[10], right=[30,40]. Left child is minimal but - // the right child has T entries: must take the successor (30). - assert!(q.remove(&key20).is_some()); - q.check_invariants(); - assert_eq!(pop_all_priorities(&mut q, 0), [10, 30, 40]); -} - -#[test] -fn delete_internal_entry_using_merge() { - let mut q = Q2::new(1); - let key20 = build_split_tree(&mut q); - // Trim the right child to [30] so both children are minimal: removing the - // internal key 20 must merge left + separator + right, then root-shrink. - let (k40, _, _) = { - let mut found = None; - let keys = q.debug_all_keys_in_order(); - for key in keys { - if key.priority == 40 { - found = Some(key); - } - } - (found.unwrap(), (), ()) - }; - assert!(q.remove(&k40).is_some()); - q.check_invariants(); - assert!(q.remove(&key20).is_some()); - q.check_invariants(); - assert_eq!(pop_all_priorities(&mut q, 0), [10, 30]); -} - -#[test] -fn borrow_from_left_sibling() { - let mut q = Q2::new(1); - build_split_tree(&mut q); - q.push(5, mask(&[0]), 5).unwrap(); - // Shape: root=[20], left=[5,10], right=[30,40]. Trim right to [30]. - let k40 = q - .debug_all_keys_in_order() - .into_iter() - .find(|k| k.priority == 40) - .unwrap(); - assert!(q.remove(&k40).is_some()); - q.check_invariants(); - // Removing 30 descends into the minimal right child; the left sibling has - // T entries, so the fix-up must borrow from the left. - let k30 = q - .debug_all_keys_in_order() - .into_iter() - .find(|k| k.priority == 30) - .unwrap(); - assert!(q.remove(&k30).is_some()); - q.check_invariants(); - assert_eq!(pop_all_priorities(&mut q, 0), [5, 10, 20]); -} - -#[test] -fn borrow_from_right_sibling() { - let mut q = Q2::new(1); - build_split_tree(&mut q); - // Shape: root=[20], left=[10], right=[30,40]. Removing 10 descends into - // the minimal left child; the right sibling has T entries, so the fix-up - // must borrow from the right. - let k10 = q - .debug_all_keys_in_order() - .into_iter() - .find(|k| k.priority == 10) - .unwrap(); - assert!(q.remove(&k10).is_some()); - q.check_invariants(); - assert_eq!(pop_all_priorities(&mut q, 0), [20, 30, 40]); -} - -#[test] -fn merge_children() { - let mut q = Q2::new(1); - build_split_tree(&mut q); - // Trim right child to [30]: both children minimal. - let k40 = q - .debug_all_keys_in_order() - .into_iter() - .find(|k| k.priority == 40) - .unwrap(); - assert!(q.remove(&k40).is_some()); - q.check_invariants(); - // Removing 10 descends into the minimal left child with a minimal right - // sibling: the fix-up must merge children around the separator. - let k10 = q - .debug_all_keys_in_order() - .into_iter() - .find(|k| k.priority == 10) - .unwrap(); - assert!(q.remove(&k10).is_some()); - q.check_invariants(); - assert_eq!(pop_all_priorities(&mut q, 0), [20, 30]); -} - -// --------------------------------------------------------------------------- -// Affinity behavior -// --------------------------------------------------------------------------- - -#[test] -fn pop_for_cpu_when_no_runnable_entry_exists() { - let mut q = Q16::new(4); - q.push(1, mask(&[1]), 1).unwrap(); - q.push(2, mask(&[1, 3]), 2).unwrap(); - assert_eq!(q.peek_key_for_cpu(0), None); - assert_eq!(q.pop_for_cpu(0), None); - assert_eq!(q.pop_for_cpu(2), None); - assert_eq!(q.len(), 2); - q.check_invariants(); -} - -#[test] -fn pop_for_cpu_when_only_one_cpu_matches() { - let mut q = Q2::new(4); - // Entry with priority 5 is only runnable on CPU 2; entries runnable on - // other CPUs have worse priority. - q.push(5, mask(&[2]), 5).unwrap(); - q.push(10, mask(&[0, 1, 3]), 10).unwrap(); - let (key, affinity, _) = q.pop_for_cpu(2).unwrap(); - assert_eq!(key.priority, 5); - assert_eq!(affinity, mask(&[2])); - assert_eq!(q.pop_for_cpu(2), None); - q.check_invariants(); -} - -#[test] -fn pop_for_cpu_with_all_cpu_affinity() { - let mut q = Q2::new(4); - for p in [30u32, 10, 20] { - q.push(p, all_cpus(4), p).unwrap(); - } - for cpu in 0..4 { - assert_eq!(q.peek_key_for_cpu(cpu).unwrap().priority, 10); - } - let (key, _, _) = q.pop_for_cpu(3).unwrap(); - assert_eq!(key.priority, 10); - q.check_invariants(); -} - -#[test] -fn affinity_pruning_skips_subtrees() { - // Many CPU-0 entries with low priorities, one CPU-1 entry with the worst - // priority: peek for CPU 1 must descend past subtrees full of CPU-0 work. - let mut q = Q2::new(2); - for p in 0..40u32 { - q.push(p, mask(&[0]), p).unwrap(); - } - let key = q.push(1000, mask(&[1]), 1000).unwrap(); - q.check_invariants(); - assert_eq!(q.peek_key_for_cpu(1), Some(key.clone())); - assert_eq!(q.pop_for_cpu(1), Some((key, mask(&[1]), 1000))); - assert_eq!(q.peek_key_for_cpu(1), None); - assert_eq!(q.peek_key_for_cpu(0).unwrap().priority, 0); - q.check_invariants(); -} - -#[test] -fn example_behavior_from_spec() { - // Spec section 32. - let mut q = Q16::new(2); - let k10 = q.push(10, mask(&[1]), 0).unwrap(); - let k20 = q.push(20, mask(&[0, 1]), 1).unwrap(); - let _k30 = q.push(30, mask(&[0]), 2).unwrap(); - - assert_eq!(q.peek_key_for_cpu(0), Some(k20.clone())); - assert_eq!(q.peek_key_for_cpu(1), Some(k10.clone())); - - let (popped, _, _) = q.pop_for_cpu(1).unwrap(); - assert_eq!(popped, k10); - - assert_eq!(q.peek_key_for_cpu(1), Some(k20.clone())); - assert_eq!(q.peek_key_for_cpu(0), Some(k20)); -} - -#[test] -fn num_cpus_64_uses_high_bits() { - let mut q = AffinityBTreeQueue::::new(64); - let key = q.push(1, mask(&[63]), 1).unwrap(); - assert_eq!(q.peek_key_for_cpu(63), Some(key)); - assert_eq!(q.peek_key_for_cpu(62), None); - assert_eq!(q.peek_key_for_cpu(64), None); - q.check_invariants(); -} - -#[test] -fn num_cpus_wide_uses_high_bits() { - // N=2: CPUs 0..128. Exercise CPU 64 (first bit of word 1) and CPU 127. - let mut q = AffinityBTreeQueue::::new(128); - let mut m64 = CpuMask::<2>::empty(); - m64.insert(64); - let mut m127 = CpuMask::<2>::empty(); - m127.insert(127); - - let k64 = q.push(1, m64, 64).unwrap(); - let k127 = q.push(2, m127, 127).unwrap(); - q.check_invariants(); - - assert_eq!(q.peek_key_for_cpu(64), Some(k64.clone())); - assert_eq!(q.peek_key_for_cpu(127), Some(k127.clone())); - assert_eq!(q.peek_key_for_cpu(63), None); - assert_eq!(q.peek_key_for_cpu(128), None); - - assert_eq!(q.pop_for_cpu(64), Some((k64, m64, 64))); - assert_eq!(q.pop_for_cpu(127), Some((k127, m127, 127))); - assert!(q.is_empty()); - q.check_invariants(); -} - -// --------------------------------------------------------------------------- -// Error handling -// --------------------------------------------------------------------------- - -#[test] -fn remove_nonexistent_key() { - let mut q = Q16::new(2); - q.push(10, mask(&[0]), 0).unwrap(); - let missing = PriorityKey { - priority: 10u32, - seq: 999, - }; - assert_eq!(q.remove(&missing), None); - assert_eq!(q.len(), 1); - q.check_invariants(); -} - -#[test] -fn invalid_empty_affinity() { - let mut q = Q16::new(4); - assert_eq!( - q.push(1, CpuMask::empty(), 0), - Err(PushError::EmptyAffinity) - ); - assert!(q.is_empty()); - q.check_invariants(); -} - -#[test] -fn invalid_cpu_in_affinity() { - let mut q = Q16::new(4); - assert_eq!(q.push(1, mask(&[4]), 0), Err(PushError::InvalidCpu)); - assert_eq!(q.push(1, mask(&[0, 5]), 0), Err(PushError::InvalidCpu)); - assert!(q.is_empty()); - q.check_invariants(); -} - -#[test] -fn sequence_overflow_is_rejected() { - let mut q = Q16::new(2); - q.debug_set_next_seq(u64::MAX - 1); - let key = q.push(1, mask(&[0]), 0).unwrap(); - assert_eq!(key.seq, u64::MAX - 1); - assert_eq!(q.push(2, mask(&[0]), 1), Err(PushError::SequenceOverflow)); - assert_eq!(q.len(), 1); - q.check_invariants(); -} - -// --------------------------------------------------------------------------- -// Generic priority types -// --------------------------------------------------------------------------- - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -struct ClassPriority { - class: u8, - sub: u8, -} - -impl Ord for ClassPriority { - fn cmp(&self, other: &Self) -> core::cmp::Ordering { - self.class - .cmp(&other.class) - .then_with(|| self.sub.cmp(&other.sub)) - } -} - -impl PartialOrd for ClassPriority { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -#[test] -fn custom_priority_type_implementing_ord() { - let mut q = AffinityBTreeQueue::::new(2); - let inputs = [ - ClassPriority { class: 2, sub: 0 }, - ClassPriority { class: 0, sub: 9 }, - ClassPriority { class: 1, sub: 1 }, - ClassPriority { class: 0, sub: 3 }, - ClassPriority { class: 1, sub: 0 }, - ]; - for (i, p) in inputs.iter().enumerate() { - q.push(*p, mask(&[0]), i as u32).unwrap(); - q.check_invariants(); - } - let mut sorted = inputs; - sorted.sort(); - for expected in sorted { - let (key, _, _) = q.pop_for_cpu(0).unwrap(); - assert_eq!(key.priority, expected); - q.check_invariants(); - } -} - -#[test] -fn reverse_priority_ordering() { - let mut q = AffinityBTreeQueue::, u32, 2>::new(2); - for p in [10u32, 30, 20] { - q.push(Reverse(p), mask(&[0]), p).unwrap(); - q.check_invariants(); - } - // With Reverse, larger raw values dequeue first. - let mut popped = Vec::new(); - while let Some((key, _, _)) = q.pop_for_cpu(0) { - popped.push(key.priority.0); - q.check_invariants(); - } - assert_eq!(popped, [30, 20, 10]); -} - -// --------------------------------------------------------------------------- -// Randomized property tests against the reference model (spec section 28) -// --------------------------------------------------------------------------- - -fn random_mask(rng: &mut XorShift64, num_cpus: usize) -> CpuMask { - let mut m = CpuMask::empty(); - let num_words = (num_cpus + 63).div_ceil(64); - for w in 0..num_words.min(N) { - let first_cpu = w * 64; - let last_cpu = ((w + 1) * 64).min(num_cpus); - let bits_count = last_cpu - first_cpu; - let allowed = if bits_count >= 64 { - u64::MAX - } else { - (1u64 << bits_count) - 1 - }; - let bits = rng.next() & allowed; - m.set_word(w, bits); - } - if m.is_empty() { - let cpu = (rng.next() as usize) % num_cpus; - m.insert(cpu); - } - m -} - -fn run_random_ops(num_cpus: usize, seed: u64, num_ops: usize) { - let mut rng = XorShift64::new(seed); - let mut q = AffinityBTreeQueue::::new(num_cpus); - let mut model = RefModel::::new(num_cpus); - - for op_index in 0..num_ops { - match rng.next() % 6 { - // Push half the time so trees grow into multiple levels. - 0..=2 => { - let priority = (rng.next() % 64) as u32; - let affinity = random_mask(&mut rng, num_cpus); - let value = op_index as u32; - assert_eq!( - q.push(priority, affinity, value), - model.push(priority, affinity, value) - ); - } - 3 => { - let cpu = (rng.next() as usize) % num_cpus; - assert_eq!(q.pop_for_cpu(cpu), model.pop_for_cpu(cpu)); - } - 4 => { - if !model.entries.is_empty() { - let idx = (rng.next() as usize) % model.entries.len(); - let key = model.entries[idx].key.clone(); - assert_eq!(q.remove(&key), model.remove(&key)); - } - } - _ => { - let key = PriorityKey { - priority: (rng.next() % 64) as u32, - seq: rng.next() % 64, - }; - assert_eq!(q.remove(&key), model.remove(&key)); - } - } - - q.check_invariants(); - assert_eq!(q.len(), model.len()); - assert_eq!(q.debug_count_entries(), model.len()); - for cpu in 0..num_cpus { - assert_eq!( - q.peek_key_for_cpu(cpu), - model.peek_key_for_cpu(cpu), - "peek mismatch for cpu {cpu}" - ); - } - assert_eq!(q.debug_all_entries_in_order(), model.entries_in_order()); - } -} - -#[test] -fn random_ops_1_cpu() { - for seed in 0..30 { - run_random_ops::<16, 1>(1, seed, 1000); - } -} - -#[test] -fn random_ops_2_cpus() { - for seed in 100..125 { - run_random_ops::<16, 1>(2, seed, 1000); - } -} - -#[test] -fn random_ops_4_cpus() { - for seed in 200..220 { - run_random_ops::<16, 1>(4, seed, 1000); - } -} - -#[test] -fn random_ops_8_cpus() { - for seed in 300..315 { - run_random_ops::<16, 1>(8, seed, 1000); - } -} - -#[test] -fn random_ops_16_cpus() { - for seed in 400..407 { - run_random_ops::<16, 1>(16, seed, 1000); - } -} - -#[test] -fn random_ops_64_cpus() { - for seed in 500..503 { - run_random_ops::<16, 1>(64, seed, 1000); - } -} - -#[test] -fn random_ops_degree_2() { - // T=2 reaches splits, merges, borrows, and root shrink constantly. - for seed in 600..620 { - run_random_ops::<2, 1>(4, seed, 600); - } -} - -#[test] -fn random_ops_degree_3() { - for seed in 700..710 { - run_random_ops::<3, 1>(4, seed, 600); - } -} - -#[test] -fn random_ops_narrow_cpu_count_degree_2() { - // T=2, N=1, 8 CPUs: exercises low CPU count with many structural changes. - for seed in 800..815 { - run_random_ops::<2, 1>(8, seed, 600); - } -} - -#[test] -fn random_ops_narrow_cpu_count_degree_16() { - for seed in 900..905 { - run_random_ops::<16, 1>(16, seed, 800); - } -} - -#[test] -fn random_ops_n2_65_cpus() { - // N=2: first test that crosses the 64-CPU boundary. - for seed in 1000..1010 { - run_random_ops::<2, 2>(65, seed, 600); - } -} - -#[test] -fn random_ops_n2_128_cpus() { - // N=2: full 128-CPU mask exercises both words of CpuMask<2>. - for seed in 1100..1105 { - run_random_ops::<16, 2>(128, seed, 800); - } -} - -// --------------------------------------------------------------------------- -// Invariant violation (defensive panic paths) -// --------------------------------------------------------------------------- - -#[test] -#[should_panic(expected = "subtree_affinity invariant is broken")] -fn broken_subtree_affinity_internal_node_panics() { - let mut q: Q2 = AffinityBTreeQueue::new(2); - // Insert 4 entries to force a root split (root becomes internal with T=2). - for i in 1u32..=4 { - q.push(i, mask(&[0]), i).unwrap(); - } - // Corrupt the root's subtree_affinity to claim CPU 1 is present. - let root_id = q.root.unwrap(); - q.nodes[root_id.0] - .as_mut() - .unwrap() - .subtree_affinity - .insert(1); - // pop_for_cpu(1): at the internal root, all children and entries are CPU 0 - // only, so the rightmost-child block (lines 584-590) is entered but the - // inner condition is false → falls through to panic. - let _ = q.pop_for_cpu(1); -} - -#[test] -#[should_panic(expected = "subtree_affinity invariant is broken")] -fn broken_subtree_affinity_leaf_node_panics() { - let mut q: Q2 = AffinityBTreeQueue::new(2); - // Single entry: root is a leaf. - q.push(1, mask(&[0]), 1).unwrap(); - // Corrupt the leaf root's subtree_affinity to claim CPU 1 is present. - let root_id = q.root.unwrap(); - q.nodes[root_id.0] - .as_mut() - .unwrap() - .subtree_affinity - .insert(1); - // pop_for_cpu(1): the for loop exhausts all leaf entries, the outer - // `if !x.leaf` is skipped (leaf), falls straight to panic. - let _ = q.pop_for_cpu(1); -} - -// --------------------------------------------------------------------------- -// RefModel error paths -// --------------------------------------------------------------------------- - -#[test] -fn ref_model_push_empty_affinity() { - let mut model: RefModel = RefModel::new(4); - assert_eq!( - model.push(1, CpuMask::empty(), 0), - Err(PushError::EmptyAffinity) - ); -} - -#[test] -fn ref_model_push_invalid_cpu() { - let mut model: RefModel = RefModel::new(4); - // CPU 4 >= num_cpus(4) → InvalidCpu - assert_eq!(model.push(1, mask(&[4]), 0), Err(PushError::InvalidCpu)); -} - -#[test] -fn ref_model_pop_out_of_range_cpu_returns_none() { - let mut model: RefModel = RefModel::new(4); - model.push(1, mask(&[0]), 0).unwrap(); - // CPU 4 >= num_cpus(4): peek_key_for_cpu returns None - assert_eq!(model.pop_for_cpu(4), None); -} diff --git a/awkernel_async_lib/Cargo.toml b/awkernel_async_lib/Cargo.toml index 34f665c36..411e2c41a 100644 --- a/awkernel_async_lib/Cargo.toml +++ b/awkernel_async_lib/Cargo.toml @@ -12,7 +12,7 @@ array-macro = "2.1" fixedbitset = "0.5.7" async-trait = "0.1" async-recursion = "1.1" -affinity_btree_queue = { path = "../affinity_btree_queue" } +affinity_btree_queue = { workspace = true } [dependencies.futures] version = "0.3" diff --git a/awkernel_lib/Cargo.toml b/awkernel_lib/Cargo.toml index c746e744a..085e7f995 100644 --- a/awkernel_lib/Cargo.toml +++ b/awkernel_lib/Cargo.toml @@ -6,7 +6,7 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] -affinity_btree_queue = { path = "../affinity_btree_queue" } +affinity_btree_queue = { workspace = true } log = "0.4" rlsf = "0.2" bitflags = "2.3"