    1|       |//! Affinity-aware priority queue backed by an augmented classical (CLRS-style) B-tree.
    2|       |//!
    3|       |//! Each entry carries `(priority, affinity: CpuMask, value)`. The tree is ordered
    4|       |//! solely by `PriorityKey<P> = (priority, seq)`, where `seq` is a unique tie-breaker.
    5|       |//! Every node stores `subtree_affinity`, the OR of all entry affinities in its
    6|       |//! subtree, which lets `peek_key_for_cpu`/`pop_for_cpu` skip subtrees with no
    7|       |//! entry runnable on the requested CPU. Smaller priority value = higher priority;
    8|       |//! wrap with `core::cmp::Reverse` for the opposite order.
    9|       |//!
   10|       |//! # Quick start
   11|       |//!
   12|       |//! ```
   13|       |//! use affinity_btree_queue::{AffinityBTreeQueue, CpuMask};
   14|       |//!
   15|       |//! // Create a queue for a 4-CPU system.
   16|       |//! // Smaller priority value = higher priority.
   17|       |//! // Use `core::cmp::Reverse` if you want larger values to dequeue first.
   18|       |//! let mut q: AffinityBTreeQueue<u32, &str> = AffinityBTreeQueue::new(4);
   19|       |//!
   20|       |//! // Build CPU affinity masks.
   21|       |//! let mut cpu0 = CpuMask::empty();
   22|       |//! cpu0.insert(0);
   23|       |//!
   24|       |//! let mut cpu12 = CpuMask::empty();
   25|       |//! cpu12.insert(1);
   26|       |//! cpu12.insert(2);
   27|       |//!
   28|       |//! let mut all = CpuMask::empty();
   29|       |//! for c in 0..4 { all.insert(c); }
   30|       |//!
   31|       |//! // push returns a PriorityKey handle that can be passed to remove() later.
   32|       |//! q.push(10, all,   "low priority, any CPU").unwrap();
   33|       |//! q.push(5,  cpu12, "medium, CPUs 1-2 only").unwrap();
   34|       |//! q.push(1,  cpu0,  "high priority, CPU 0 only").unwrap();
   35|       |//!
   36|       |//! // pop_for_cpu(cpu) returns the highest-priority (lowest value) entry
   37|       |//! // whose affinity includes the given CPU.
   38|       |//! let (_, _, v) = q.pop_for_cpu(0).unwrap();
   39|       |//! assert_eq!(v, "high priority, CPU 0 only"); // priority=1
   40|       |//!
   41|       |//! // The cpu12 entry (priority=5) is not runnable on CPU 0, so priority=10 is next.
   42|       |//! let (_, _, v) = q.pop_for_cpu(0).unwrap();
   43|       |//! assert_eq!(v, "low priority, any CPU");     // priority=10
   44|       |//!
   45|       |//! // CPU 1 can run the cpu12 entry (priority=5).
   46|       |//! let (_, _, v) = q.pop_for_cpu(1).unwrap();
   47|       |//! assert_eq!(v, "medium, CPUs 1-2 only");     // priority=5
   48|       |//! ```
   49|       |//!
   50|       |//! # Max-first ordering with `Reverse`
   51|       |//!
   52|       |//! ```
   53|       |//! use affinity_btree_queue::{AffinityBTreeQueue, CpuMask};
   54|       |//! use core::cmp::Reverse;
   55|       |//!
   56|       |//! let mut q: AffinityBTreeQueue<Reverse<u32>, u32> = AffinityBTreeQueue::new(1);
   57|       |//! let mut m = CpuMask::empty();
   58|       |//! m.insert(0);
   59|       |//! q.push(Reverse(1),   m, 1).unwrap();
   60|       |//! q.push(Reverse(100), m, 100).unwrap();
   61|       |//!
   62|       |//! // Reverse(100) < Reverse(1), so the numerically larger value dequeues first.
   63|       |//! let (_, _, v) = q.pop_for_cpu(0).unwrap();
   64|       |//! assert_eq!(v, 100);
   65|       |//! ```
   66|       |
   67|       |#![no_std]
   68|       |#![forbid(unsafe_code)]
   69|       |
   70|       |extern crate alloc;
   71|       |
   72|       |use alloc::vec::Vec;
   73|       |use core::cmp::Ordering;
   74|       |use core::mem;
   75|       |
   76|       |/// Default B-tree minimum degree for production builds.
   77|       |pub const DEFAULT_MIN_DEGREE: usize = 16;
   78|       |
   79|       |/// Sequence number type used as a unique tie-breaker for entries with the same
   80|       |/// priority.
   81|       |pub type Seq = u64;
   82|       |
   83|       |/// CPU affinity bitmask backed by `N` 64-bit words, supporting up to `N * 64`
   84|       |/// CPUs. `CpuMask<1>` (the default) supports up to 64 CPUs; `CpuMask<2>`
   85|       |/// supports up to 128, etc.
   86|       |///
   87|       |/// # Examples
   88|       |///
   89|       |/// ```
   90|       |/// use affinity_btree_queue::CpuMask;
   91|       |///
   92|       |/// let mut m: CpuMask = CpuMask::empty();
   93|       |/// m.insert(0);
   94|       |/// m.insert(3);
   95|       |/// assert!(m.contains(0));
   96|       |/// assert!(m.contains(3));
   97|       |/// assert!(!m.contains(1));
   98|       |///
   99|       |/// // Construct directly from a bitmask.
  100|       |/// let m2: CpuMask = CpuMask::from_bits_truncate(0b1001); // CPU 0 and 3
  101|       |/// assert_eq!(m, m2);
  102|       |/// ```
  103|       |#[derive(Clone, Copy, Debug, PartialEq, Eq)]
  104|       |pub struct CpuMask<const N: usize = 1>([u64; N]);
  105|       |
  106|       |impl<const N: usize> CpuMask<N> {
  107|  1.87M|    pub fn empty() -> Self {
  108|  1.87M|        CpuMask([0u64; N])
  109|  1.87M|    }
  110|       |
  111|  11.3M|    pub fn is_empty(self) -> bool {
  112|  11.3M|        self.0.iter().all(|&w| w == 0)
                      ^11.3M        ^11.3M
  113|  11.3M|    }
  114|       |
  115|   131M|    pub fn contains(self, cpu: usize) -> bool {
  116|   131M|        let word = cpu / 64;
  117|   131M|        let bit = cpu % 64;
  118|   131M|        word < N && (self.0[word] >> bit) & 1 == 1
                                  ^131M
  119|   131M|    }
  120|       |
  121|  16.0M|    pub fn union(self, other: Self) -> Self {
  122|  16.0M|        let mut result = [0u64; N];
  123|  17.0M|        for (i, item) in result.iter_mut().enumerate().take(N) {
                                       ^16.0M ^16.0M     ^16.0M      ^16.0M
  124|  17.0M|            *item = self.0[i] | other.0[i];
  125|  17.0M|        }
  126|  16.0M|        CpuMask(result)
  127|  16.0M|    }
  128|       |
  129|       |    /// Sets the bit for `cpu`. Panics if `cpu >= N * 64`.
  130|       |    ///
  131|       |    /// # Examples
  132|       |    ///
  133|       |    /// ```
  134|       |    /// use affinity_btree_queue::CpuMask;
  135|       |    /// let mut m: CpuMask = CpuMask::empty();
  136|       |    /// m.insert(0);
  137|       |    /// m.insert(63);
  138|       |    /// assert!(m.contains(0));
  139|       |    /// assert!(m.contains(63));
  140|       |    /// assert!(!m.contains(1));
  141|       |    /// ```
  142|  23.3k|    pub fn insert(&mut self, cpu: usize) {
  143|  23.3k|        assert!(cpu < N * 64, "CPU index out of CpuMask range");
  144|  23.3k|        self.0[cpu / 64] |= 1u64 << (cpu % 64);
  145|  23.3k|    }
  146|       |
  147|       |    /// Creates a mask with word 0 set to `bits` and all other words zeroed.
  148|       |    /// For `N = 1`, this is equivalent to setting the full mask.
  149|       |    ///
  150|       |    /// # Examples
  151|       |    ///
  152|       |    /// ```
  153|       |    /// use affinity_btree_queue::CpuMask;
  154|       |    /// // Set bits for CPU 0 and CPU 2.
  155|       |    /// let m: CpuMask = CpuMask::from_bits_truncate(0b0101);
  156|       |    /// assert!(m.contains(0));
  157|       |    /// assert!(!m.contains(1));
  158|       |    /// assert!(m.contains(2));
  159|       |    /// ```
  160|      2|    pub fn from_bits_truncate(bits: u64) -> Self {
  161|      2|        let mut result = [0u64; N];
  162|      2|        if N > 0 {
  163|      2|            result[0] = bits;
  164|      2|        }
                      ^0
  165|      2|        CpuMask(result)
  166|      2|    }
  167|       |
  168|       |    /// Returns word 0 of the mask as a `u64`.
  169|       |    /// For `N = 1`, this is the full mask.
  170|      2|    pub fn bits(self) -> u64 {
  171|      2|        if N > 0 { self.0[0] } else { 0 }
                                                    ^0
  172|      2|    }
  173|       |
  174|       |    /// Returns word `idx` of the backing array. Panics if `idx >= N`.
  175|      2|    pub fn word(self, idx: usize) -> u64 {
  176|      2|        self.0[idx]
  177|      2|    }
  178|       |
  179|       |    /// Sets word `idx` of the backing array. Panics if `idx >= N`.
  180|  75.0k|    pub fn set_word(&mut self, idx: usize, bits: u64) {
  181|  75.0k|        self.0[idx] = bits;
  182|  75.0k|    }
  183|       |
  184|       |    /// True if no bit at index >= `num_cpus` is set.
  185|  11.2M|    fn fits_within(self, num_cpus: usize) -> bool {
  186|  11.2M|        if num_cpus >= N * 64 {
  187|   583k|            return true;
  188|  10.6M|        }
  189|  10.6M|        let full_words = num_cpus / 64;
  190|  10.6M|        let extra_bits = num_cpus % 64;
  191|  10.6M|        let upper_start = if extra_bits > 0 {
  192|  10.6M|            full_words + 1
  193|       |        } else {
  194|      2|            full_words
  195|       |        };
  196|  10.6M|        for i in upper_start..N {
                          ^3
  197|      3|            if self.0[i] != 0 {
  198|      2|                return false;
  199|      1|            }
  200|       |        }
  201|  10.6M|        if extra_bits > 0 && full_words < N && self.0[full_words] >> extra_bits != 0 {
                                           ^10.6M            ^10.6M
  202|      5|            return false;
  203|  10.6M|        }
  204|  10.6M|        true
  205|  11.2M|    }
  206|       |}
  207|       |
  208|       |#[derive(Clone, Copy, Debug, PartialEq, Eq)]
  209|       |pub enum PushError {
  210|       |    EmptyAffinity,
  211|       |    SequenceOverflow,
  212|       |    InvalidCpu,
  213|       |}
  214|       |
  215|       |/// Unique B-tree key: ordered by `(priority, seq)` lexicographically.
  216|       |#[derive(Clone, Debug, PartialEq, Eq)]
  217|       |pub struct PriorityKey<P> {
  218|       |    pub priority: P,
  219|       |    pub seq: Seq,
  220|       |}
  221|       |
  222|       |impl<P: Ord> Ord for PriorityKey<P> {
  223|   172M|    fn cmp(&self, other: &Self) -> Ordering {
  224|   172M|        self.priority
  225|   172M|            .cmp(&other.priority)
  226|   172M|            .then_with(|| self.seq.cmp(&other.seq))
                                        ^22.7M   ^22.7M^22.7M
  227|   172M|    }
  228|       |}
  229|       |
  230|       |impl<P: Ord> PartialOrd for PriorityKey<P> {
  231|  38.0M|    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
  232|  38.0M|        Some(self.cmp(other))
  233|  38.0M|    }
  234|       |}
  235|       |
  236|       |#[derive(Debug)]
  237|       |struct Entry<P, V, const N: usize> {
  238|       |    key: PriorityKey<P>,
  239|       |    affinity: CpuMask<N>,
  240|       |    value: V,
  241|       |}
  242|       |
  243|       |#[derive(Clone, Copy, Debug, PartialEq, Eq)]
  244|       |struct NodeId(usize);
  245|       |
  246|       |#[derive(Debug)]
  247|       |struct Node<P, V, const N: usize> {
  248|       |    leaf: bool,
  249|       |    entries: Vec<Entry<P, V, N>>,
  250|       |    // Empty when leaf; otherwise children.len() == entries.len() + 1.
  251|       |    children: Vec<NodeId>,
  252|       |    // OR of all affinities in this node and its descendants.
  253|       |    subtree_affinity: CpuMask<N>,
  254|       |}
  255|       |
  256|       |/// First index in `entries` whose key is >= `key`.
  257|   322k|fn lower_bound<P: Ord, V, const N: usize>(
  258|   322k|    entries: &[Entry<P, V, N>],
  259|   322k|    key: &PriorityKey<P>,
  260|   322k|) -> usize {
  261|   322k|    let mut lo = 0;
  262|   322k|    let mut hi = entries.len();
  263|  1.21M|    while lo < hi {
  264|   896k|        let mid = (lo + hi) / 2;
  265|   896k|        if entries[mid].key < *key {
  266|   287k|            lo = mid + 1;
  267|   608k|        } else {
  268|   608k|            hi = mid;
  269|   608k|        }
  270|       |    }
  271|   322k|    lo
  272|   322k|}
  273|       |
  274|       |/// Affinity-aware priority queue over an augmented classical B-tree of
  275|       |/// minimum degree `T` (`MIN_KEYS = T - 1`, `MAX_KEYS = 2 * T - 1`).
  276|       |///
  277|       |/// `N` is the number of 64-bit words in the CPU mask; the default `N = 1`
  278|       |/// supports up to 64 CPUs. Use `N = 2` for up to 128 CPUs, etc.
  279|       |pub struct AffinityBTreeQueue<P, V, const T: usize = DEFAULT_MIN_DEGREE, const N: usize = 1> {
  280|       |    root: Option<NodeId>,
  281|       |    // Arena with stable ids; never compacted. Freed slots go to `free_list`.
  282|       |    nodes: Vec<Option<Node<P, V, N>>>,
  283|       |    free_list: Vec<NodeId>,
  284|       |    len: usize,
  285|       |    next_seq: Seq,
  286|       |    num_cpus: usize,
  287|       |}
  288|       |
  289|       |impl<P: Ord + Clone, V, const T: usize, const N: usize> AffinityBTreeQueue<P, V, T, N> {
  290|       |    // Only referenced by the test invariant checker.
  291|       |    #[cfg_attr(not(test), allow(dead_code))]
  292|       |    const MIN_KEYS: usize = T - 1;
  293|       |    const MAX_KEYS: usize = 2 * T - 1;
  294|       |
  295|       |    /// Creates an empty queue for CPUs `0..num_cpus`. Panics if `T < 2` or
  296|       |    /// `num_cpus > N * 64`.
  297|       |    ///
  298|       |    /// # Examples
  299|       |    ///
  300|       |    /// ```
  301|       |    /// use affinity_btree_queue::AffinityBTreeQueue;
  302|       |    ///
  303|       |    /// // Default degree T=16, supports up to 64 CPUs (N=1).
  304|       |    /// let q: AffinityBTreeQueue<u32, ()> = AffinityBTreeQueue::new(8);
  305|       |    /// assert_eq!(q.num_cpus(), 8);
  306|       |    /// assert!(q.is_empty());
  307|       |    /// ```
  308|    201|    pub fn new(num_cpus: usize) -> Self {
  309|    201|        assert!(T >= 2, "B-tree minimum degree T must be at least 2");
  310|    201|        assert!(num_cpus <= N * 64, "num_cpus exceeds CpuMask capacity N*64");
  311|    201|        Self {
  312|    201|            root: None,
  313|    201|            nodes: Vec::new(),
  314|    201|            free_list: Vec::new(),
  315|    201|            len: 0,
  316|    201|            next_seq: 0,
  317|    201|            num_cpus,
  318|    201|        }
  319|    201|    }
  320|       |
  321|   141k|    pub fn len(&self) -> usize {
  322|   141k|        self.len
  323|   141k|    }
  324|       |
  325|     12|    pub fn is_empty(&self) -> bool {
  326|     12|        self.len == 0
  327|     12|    }
  328|       |
  329|      2|    pub fn num_cpus(&self) -> usize {
  330|      2|        self.num_cpus
  331|      2|    }
  332|       |
  333|  20.8M|    fn node(&self, id: NodeId) -> &Node<P, V, N> {
  334|  20.8M|        self.nodes[id.0].as_ref().expect("dangling NodeId")
  335|  20.8M|    }
  336|       |
  337|   598k|    fn node_mut(&mut self, id: NodeId) -> &mut Node<P, V, N> {
  338|   598k|        self.nodes[id.0].as_mut().expect("dangling NodeId")
  339|   598k|    }
  340|       |
  341|  13.6k|    fn alloc_node(&mut self, leaf: bool) -> NodeId {
  342|  13.6k|        let node = Node {
  343|  13.6k|            leaf,
  344|  13.6k|            entries: Vec::new(),
  345|  13.6k|            children: Vec::new(),
  346|  13.6k|            subtree_affinity: CpuMask::empty(),
  347|  13.6k|        };
  348|  13.6k|        if let Some(id) = self.free_list.pop() {
                                  ^9.45k
  349|  9.45k|            debug_assert!(self.nodes[id.0].is_none());
  350|  9.45k|            self.nodes[id.0] = Some(node);
  351|  9.45k|            id
  352|       |        } else {
  353|  4.22k|            let id = NodeId(self.nodes.len());
  354|  4.22k|            self.nodes.push(Some(node));
  355|  4.22k|            id
  356|       |        }
  357|  13.6k|    }
  358|       |
  359|  2.42k|    fn free_node(&mut self, id: NodeId) {
  360|  2.42k|        let taken = self.nodes[id.0].take();
  361|  2.42k|        debug_assert!(taken.is_some());
  362|  2.42k|        self.free_list.push(id);
  363|  2.42k|    }
  364|       |
  365|       |    /// Recomputes `subtree_affinity` of `id` from its entries and children.
  366|       |    /// Must be called after every mutation of a node's entries or children.
  367|   380k|    fn pull(&mut self, id: NodeId) {
  368|   380k|        let mut mask = CpuMask::empty();
  369|       |        {
  370|   380k|            let x = self.node(id);
  371|  2.82M|            for e in &x.entries {
                                   ^380k
  372|  2.82M|                mask = mask.union(e.affinity);
  373|  2.82M|            }
  374|   889k|            for &c in &x.children {
                                    ^380k
  375|   889k|                mask = mask.union(self.node(c).subtree_affinity);
  376|   889k|            }
  377|       |        }
  378|   380k|        self.node_mut(id).subtree_affinity = mask;
  379|   380k|    }
  380|       |
  381|       |    /// Inserts an entry with the given `priority`, CPU `affinity`, and `value`.
  382|       |    ///
  383|       |    /// Returns a [`PriorityKey`] handle that uniquely identifies the entry and
  384|       |    /// can be passed to [`remove`](Self::remove) later.
  385|       |    ///
  386|       |    /// # Errors
  387|       |    ///
  388|       |    /// - [`PushError::EmptyAffinity`] — `affinity` has no bits set.
  389|       |    /// - [`PushError::InvalidCpu`] — `affinity` has a bit set for a CPU ≥ `num_cpus`.
  390|       |    /// - [`PushError::SequenceOverflow`] — internal sequence counter wrapped
  391|       |    ///   (requires 2⁶⁴ prior pushes; practically impossible).
  392|       |    ///
  393|       |    /// # Examples
  394|       |    ///
  395|       |    /// ```
  396|       |    /// use affinity_btree_queue::{AffinityBTreeQueue, CpuMask, PushError};
  397|       |    ///
  398|       |    /// let mut q: AffinityBTreeQueue<u32, u32> = AffinityBTreeQueue::new(4);
  399|       |    ///
  400|       |    /// let mut m = CpuMask::empty();
  401|       |    /// m.insert(0);
  402|       |    /// let _key = q.push(10, m, 42).unwrap();
  403|       |    /// assert_eq!(q.len(), 1);
  404|       |    ///
  405|       |    /// // Empty affinity mask is rejected.
  406|       |    /// assert_eq!(q.push(1, CpuMask::empty(), 0), Err(PushError::EmptyAffinity));
  407|       |    ///
  408|       |    /// // CPU 4 is out of range for a 4-CPU queue.
  409|       |    /// let mut bad = CpuMask::empty();
  410|       |    /// bad.insert(4);
  411|       |    /// assert_eq!(q.push(1, bad, 0), Err(PushError::InvalidCpu));
  412|       |    /// ```
  413|  72.6k|    pub fn push(
  414|  72.6k|        &mut self,
  415|  72.6k|        priority: P,
  416|  72.6k|        affinity: CpuMask<N>,
  417|  72.6k|        value: V,
  418|  72.6k|    ) -> Result<PriorityKey<P>, PushError> {
  419|  72.6k|        if affinity.is_empty() {
  420|      1|            return Err(PushError::EmptyAffinity);
  421|  72.6k|        }
  422|  72.6k|        if !affinity.fits_within(self.num_cpus) {
  423|      2|            return Err(PushError::InvalidCpu);
  424|  72.6k|        }
  425|  72.6k|        let seq = self.next_seq;
  426|  72.6k|        let next_seq = seq.checked_add(1).ok_or(PushError::SequenceOverflow)?;
                          ^72.6k                                                          ^1
  427|  72.6k|        self.next_seq = next_seq;
  428|       |
  429|  72.6k|        let key = PriorityKey { priority, seq };
  430|  72.6k|        let ret = key.clone();
  431|  72.6k|        let entry = Entry {
  432|  72.6k|            key,
  433|  72.6k|            affinity,
  434|  72.6k|            value,
  435|  72.6k|        };
  436|       |
  437|  72.6k|        match self.root {
  438|    464|            None => {
  439|    464|                let r = self.alloc_node(true);
  440|    464|                self.node_mut(r).entries.push(entry);
  441|    464|                self.pull(r);
  442|    464|                self.root = Some(r);
  443|    464|            }
  444|  72.1k|            Some(r) => {
  445|  72.1k|                if self.node(r).entries.len() == Self::MAX_KEYS {
  446|  2.45k|                    let s = self.alloc_node(false);
  447|  2.45k|                    self.node_mut(s).children.push(r);
  448|  2.45k|                    self.root = Some(s);
  449|  2.45k|                    self.split_child(s, 0);
  450|  2.45k|                    self.insert_non_full(s, entry);
  451|  69.7k|                } else {
  452|  69.7k|                    self.insert_non_full(r, entry);
  453|  69.7k|                }
  454|       |            }
  455|       |        }
  456|  72.6k|        self.len += 1;
  457|  72.6k|        Ok(ret)
  458|  72.6k|    }
  459|       |
  460|       |    /// Splits the full child `parent.children[index]` around its median entry.
  461|  10.7k|    fn split_child(&mut self, parent_id: NodeId, index: usize) {
  462|  10.7k|        let y_id = self.node(parent_id).children[index];
  463|  10.7k|        let y_leaf = self.node(y_id).leaf;
  464|  10.7k|        let z_id = self.alloc_node(y_leaf);
  465|       |
  466|  10.7k|        let (median, right_entries, right_children) = {
  467|  10.7k|            let y = self.node_mut(y_id);
  468|  10.7k|            debug_assert_eq!(y.entries.len(), Self::MAX_KEYS);
  469|  10.7k|            let right_entries = y.entries.split_off(T);
  470|  10.7k|            let median = y.entries.pop().expect("full node has a median");
  471|  10.7k|            let right_children = if y_leaf {
  472|  5.07k|                Vec::new()
  473|       |            } else {
  474|  5.69k|                y.children.split_off(T)
  475|       |            };
  476|  10.7k|            (median, right_entries, right_children)
  477|       |        };
  478|       |
  479|  10.7k|        {
  480|  10.7k|            let z = self.node_mut(z_id);
  481|  10.7k|            z.entries = right_entries;
  482|  10.7k|            z.children = right_children;
  483|  10.7k|        }
  484|  10.7k|        {
  485|  10.7k|            let p = self.node_mut(parent_id);
  486|  10.7k|            p.entries.insert(index, median);
  487|  10.7k|            p.children.insert(index + 1, z_id);
  488|  10.7k|        }
  489|       |
  490|  10.7k|        self.pull(y_id);
  491|  10.7k|        self.pull(z_id);
  492|  10.7k|        self.pull(parent_id);
  493|  10.7k|    }
  494|       |
  495|   163k|    fn insert_non_full(&mut self, node_id: NodeId, entry: Entry<P, V, N>) {
  496|   163k|        debug_assert!(self.node(node_id).entries.len() < Self::MAX_KEYS);
  497|       |
  498|   163k|        if self.node(node_id).leaf {
  499|  72.1k|            let i = lower_bound(&self.node(node_id).entries, &entry.key);
  500|  72.1k|            debug_assert!(
  501|  72.1k|                i >= self.node(node_id).entries.len()
  502|  60.5k|                    || self.node(node_id).entries[i].key != entry.key,
  503|       |                "duplicate key"
  504|       |            );
  505|  72.1k|            self.node_mut(node_id).entries.insert(i, entry);
  506|  72.1k|            self.pull(node_id);
  507|  72.1k|            return;
  508|  91.1k|        }
  509|       |
  510|  91.1k|        let mut i = lower_bound(&self.node(node_id).entries, &entry.key);
  511|  91.1k|        let child_id = self.node(node_id).children[i];
  512|  91.1k|        if self.node(child_id).entries.len() == Self::MAX_KEYS {
  513|  8.31k|            self.split_child(node_id, i);
  514|  8.31k|            if entry.key > self.node(node_id).entries[i].key {
  515|  3.73k|                i += 1;
  516|  3.73k|            } else {
  517|  4.58k|                debug_assert!(
  518|  4.58k|                    entry.key != self.node(node_id).entries[i].key,
  519|       |                    "duplicate key"
  520|       |                );
  521|       |            }
  522|  82.8k|        }
  523|  91.1k|        let next_child = self.node(node_id).children[i];
  524|  91.1k|        self.insert_non_full(next_child, entry);
  525|  91.1k|        self.pull(node_id);
  526|   163k|    }
  527|       |
  528|       |    /// Smallest key whose affinity contains `cpu`, or `None`.
  529|       |    ///
  530|       |    /// Does not remove the entry; use [`pop_for_cpu`](Self::pop_for_cpu) to
  531|       |    /// remove it.
  532|       |    ///
  533|       |    /// # Examples
  534|       |    ///
  535|       |    /// ```
  536|       |    /// use affinity_btree_queue::{AffinityBTreeQueue, CpuMask};
  537|       |    ///
  538|       |    /// let mut q: AffinityBTreeQueue<u32, u32> = AffinityBTreeQueue::new(2);
  539|       |    /// assert_eq!(q.peek_key_for_cpu(0), None); // empty queue
  540|       |    ///
  541|       |    /// let mut m = CpuMask::empty();
  542|       |    /// m.insert(1); // CPU 1 only
  543|       |    /// let key = q.push(5, m, 99).unwrap();
  544|       |    ///
  545|       |    /// assert_eq!(q.peek_key_for_cpu(1), Some(key.clone())); // visible on CPU 1
  546|       |    /// assert_eq!(q.peek_key_for_cpu(0), None);              // not visible on CPU 0
  547|       |    /// assert_eq!(q.len(), 1);                               // entry was not consumed
  548|       |    /// ```
  549|  1.72M|    pub fn peek_key_for_cpu(&self, cpu: usize) -> Option<PriorityKey<P>> {
  550|  1.72M|        if cpu >= self.num_cpus {
  551|      3|            return None;
  552|  1.72M|        }
  553|  1.72M|        let r = self.root?;
                          ^1.71M       ^9.13k
  554|  1.71M|        if !self.node(r).subtree_affinity.contains(cpu) {
  555|  13.8k|            return None;
  556|  1.69M|        }
  557|  1.69M|        Some(self.find_first_eligible_key(r, cpu))
  558|  1.72M|    }
  559|       |
  560|       |    /// Precondition: `node(node_id).subtree_affinity.contains(cpu)`.
  561|       |    /// Scans in logical order child[0], entry[0], child[1], ... so the first
  562|       |    /// eligible entry found is the minimum eligible key in the subtree.
  563|  1.69M|    fn find_first_eligible_key(&self, node_id: NodeId, cpu: usize) -> PriorityKey<P> {
  564|  1.69M|        let mut cur = node_id;
  565|       |
  566|       |        'descend: loop {
  567|  3.94M|            let x = self.node(cur);
  568|  3.94M|            debug_assert!(x.subtree_affinity.contains(cpu));
  569|  3.94M|            let n = x.entries.len();
  570|       |
  571|  5.41M|            for i in 0..n {
                                   ^3.94M
  572|  5.41M|                if !x.leaf {
  573|  2.42M|                    let child_id = x.children[i];
  574|  2.42M|                    if self.node(child_id).subtree_affinity.contains(cpu) {
  575|  2.22M|                        cur = child_id;
  576|  2.22M|                        continue 'descend;
  577|   207k|                    }
  578|  2.98M|                }
  579|  3.18M|                if x.entries[i].affinity.contains(cpu) {
  580|  1.69M|                    return x.entries[i].key.clone();
  581|  1.49M|                }
  582|       |            }
  583|       |
  584|  23.7k|            if !x.leaf {
  585|  23.7k|                let child_id = x.children[n];
  586|  23.7k|                if self.node(child_id).subtree_affinity.contains(cpu) {
  587|  23.7k|                    cur = child_id;
  588|  23.7k|                    continue 'descend;
  589|      1|                }
  590|      1|            }
  591|       |
  592|      2|            panic!("subtree_affinity invariant is broken");
  593|       |        }
  594|  1.69M|    }
  595|       |
  596|       |    /// Removes and returns the highest-priority entry runnable on `cpu`.
  597|       |    ///
  598|       |    /// Returns `None` if the queue is empty or no entry's affinity includes
  599|       |    /// `cpu`.
  600|       |    ///
  601|       |    /// # Examples
  602|       |    ///
  603|       |    /// ```
  604|       |    /// use affinity_btree_queue::{AffinityBTreeQueue, CpuMask};
  605|       |    ///
  606|       |    /// let mut q: AffinityBTreeQueue<u32, &str> = AffinityBTreeQueue::new(2);
  607|       |    ///
  608|       |    /// let mut cpu0 = CpuMask::empty();
  609|       |    /// cpu0.insert(0);
  610|       |    /// let mut cpu1 = CpuMask::empty();
  611|       |    /// cpu1.insert(1);
  612|       |    ///
  613|       |    /// q.push(1, cpu0, "only cpu0").unwrap();
  614|       |    /// q.push(2, cpu1, "only cpu1").unwrap();
  615|       |    ///
  616|       |    /// // CPU 0 only dequeues its own entry.
  617|       |    /// let (_, affinity, value) = q.pop_for_cpu(0).unwrap();
  618|       |    /// assert_eq!(value, "only cpu0");
  619|       |    /// assert!(affinity.contains(0));
  620|       |    ///
  621|       |    /// // No more entries runnable on CPU 0.
  622|       |    /// assert!(q.pop_for_cpu(0).is_none());
  623|       |    /// ```
  624|  26.0k|    pub fn pop_for_cpu(&mut self, cpu: usize) -> Option<(PriorityKey<P>, CpuMask<N>, V)> {
  625|  26.0k|        let key = self.peek_key_for_cpu(cpu)?;
                          ^25.7k                          ^281
  626|  25.7k|        let removed = self.remove(&key);
  627|  25.7k|        debug_assert!(removed.is_some());
  628|  25.7k|        debug_assert!(removed.as_ref().is_none_or(|(_, a, _)| a.contains(cpu)));
  629|  25.7k|        removed
  630|  26.0k|    }
  631|       |
  632|       |    /// Removes the entry identified by `key` and returns `(key, affinity, value)`,
  633|       |    /// or `None` if no such entry exists.
  634|       |    ///
  635|       |    /// The key is obtained from [`push`](Self::push) or
  636|       |    /// [`peek_key_for_cpu`](Self::peek_key_for_cpu).
  637|       |    ///
  638|       |    /// # Examples
  639|       |    ///
  640|       |    /// ```
  641|       |    /// use affinity_btree_queue::{AffinityBTreeQueue, CpuMask};
  642|       |    ///
  643|       |    /// let mut q: AffinityBTreeQueue<u32, u32> = AffinityBTreeQueue::new(1);
  644|       |    /// let mut m = CpuMask::empty();
  645|       |    /// m.insert(0);
  646|       |    ///
  647|       |    /// let key = q.push(7, m, 42).unwrap();
  648|       |    /// let (k, _, v) = q.remove(&key).unwrap();
  649|       |    /// assert_eq!(k, key);
  650|       |    /// assert_eq!(v, 42);
  651|       |    ///
  652|       |    /// // Removing the same key a second time returns None.
  653|       |    /// assert!(q.remove(&key).is_none());
  654|       |    /// ```
  655|  72.9k|    pub fn remove(&mut self, key: &PriorityKey<P>) -> Option<(PriorityKey<P>, CpuMask<N>, V)> {
  656|  72.9k|        let root = self.root?;
                          ^72.8k          ^147
  657|  72.8k|        let removed = self.delete_from_node(root, key);
  658|       |
  659|  72.8k|        if let Some(r) = self.root
  660|  72.8k|            && self.node(r).entries.is_empty()
  661|       |        {
  662|  2.42k|            if self.node(r).leaf {
  663|    288|                self.free_node(r);
  664|    288|                self.root = None;
  665|  2.13k|            } else {
  666|  2.13k|                let new_root = self.node(r).children[0];
  667|  2.13k|                self.free_node(r);
  668|  2.13k|                self.root = Some(new_root);
  669|  2.13k|            }
  670|  70.3k|        }
  671|       |
  672|  72.8k|        if removed.is_some() {
  673|  49.2k|            self.len -= 1;
  674|  49.2k|        }
                      ^23.5k
  675|  72.8k|        removed.map(|e| (e.key, e.affinity, e.value))
                                       ^49.2k ^49.2k      ^49.2k
  676|  72.9k|    }
  677|       |
  678|   159k|    fn delete_from_node(
  679|   159k|        &mut self,
  680|   159k|        node_id: NodeId,
  681|   159k|        key: &PriorityKey<P>,
  682|   159k|    ) -> Option<Entry<P, V, N>> {
  683|   159k|        let i = lower_bound(&self.node(node_id).entries, key);
  684|   159k|        let leaf = self.node(node_id).leaf;
  685|   159k|        let found_here =
  686|   159k|            i < self.node(node_id).entries.len() && self.node(node_id).entries[i].key == *key;
                                                                  ^142k
  687|       |
  688|   159k|        if found_here {
  689|  49.8k|            if leaf {
  690|  46.5k|                let entry = self.node_mut(node_id).entries.remove(i);
  691|  46.5k|                self.pull(node_id);
  692|  46.5k|                Some(entry)
  693|       |            } else {
  694|  3.31k|                let removed = self.delete_internal_entry(node_id, i);
  695|  3.31k|                self.pull(node_id);
  696|  3.31k|                removed
  697|       |            }
  698|   109k|        } else if leaf {
  699|  23.5k|            None
  700|       |        } else {
  701|  85.8k|            let child_index = self.ensure_child_has_at_least_t(node_id, i);
  702|  85.8k|            let child_id = self.node(node_id).children[child_index];
  703|  85.8k|            let removed = self.delete_from_node(child_id, key);
  704|  85.8k|            self.pull(node_id);
  705|  85.8k|            removed
  706|       |        }
  707|   159k|    }
  708|       |
  709|       |    /// Deletes `entries[index]` of the internal node `node_id` using the
  710|       |    /// standard predecessor / successor / merge cases.
  711|  3.31k|    fn delete_internal_entry(&mut self, node_id: NodeId, index: usize) -> Option<Entry<P, V, N>> {
  712|  3.31k|        let left_id = self.node(node_id).children[index];
  713|  3.31k|        let right_id = self.node(node_id).children[index + 1];
  714|       |
  715|  3.31k|        if self.node(left_id).entries.len() >= T {
  716|  1.86k|            let predecessor = self.remove_max_entry(left_id);
  717|  1.86k|            let old = mem::replace(&mut self.node_mut(node_id).entries[index], predecessor);
  718|  1.86k|            self.pull(node_id);
  719|  1.86k|            return Some(old);
  720|  1.44k|        }
  721|       |
  722|  1.44k|        if self.node(right_id).entries.len() >= T {
  723|    811|            let successor = self.remove_min_entry(right_id);
  724|    811|            let old = mem::replace(&mut self.node_mut(node_id).entries[index], successor);
  725|    811|            self.pull(node_id);
  726|    811|            return Some(old);
  727|    635|        }
  728|       |
  729|    635|        let target_key = self.node(node_id).entries[index].key.clone();
  730|    635|        let merged_id = self.merge_children_around_entry(node_id, index);
  731|    635|        let removed = self.delete_from_node(merged_id, &target_key);
  732|    635|        self.pull(node_id);
  733|    635|        removed
  734|  3.31k|    }
  735|       |
  736|    958|    fn remove_min_entry(&mut self, node_id: NodeId) -> Entry<P, V, N> {
  737|    958|        if self.node(node_id).leaf {
  738|    811|            let entry = self.node_mut(node_id).entries.remove(0);
  739|    811|            self.pull(node_id);
  740|    811|            return entry;
  741|    147|        }
  742|    147|        let child_index = self.ensure_child_has_at_least_t(node_id, 0);
  743|    147|        debug_assert_eq!(child_index, 0);
  744|    147|        let child_id = self.node(node_id).children[child_index];
  745|    147|        let entry = self.remove_min_entry(child_id);
  746|    147|        self.pull(node_id);
  747|    147|        entry
  748|    958|    }
  749|       |
  750|  2.23k|    fn remove_max_entry(&mut self, node_id: NodeId) -> Entry<P, V, N> {
  751|  2.23k|        if self.node(node_id).leaf {
  752|  1.86k|            let entry = self
  753|  1.86k|                .node_mut(node_id)
  754|  1.86k|                .entries
  755|  1.86k|                .pop()
  756|  1.86k|                .expect("remove_max_entry on empty node");
  757|  1.86k|            self.pull(node_id);
  758|  1.86k|            return entry;
  759|    368|        }
  760|    368|        let last_child_index = self.node(node_id).children.len() - 1;
  761|       |        // If the rightmost child gets merged into its left sibling, the
  762|       |        // returned index is last_child_index - 1; always use the return value.
  763|    368|        let child_index = self.ensure_child_has_at_least_t(node_id, last_child_index);
  764|    368|        let child_id = self.node(node_id).children[child_index];
  765|    368|        let entry = self.remove_max_entry(child_id);
  766|    368|        self.pull(node_id);
  767|    368|        entry
  768|  2.23k|    }
  769|       |
  770|       |    /// Ensures `parent.children[index]` has at least `T` entries before a
  771|       |    /// descent, borrowing from a sibling or merging. Returns the (possibly
  772|       |    /// shifted) index of the child to descend into.
  773|  86.3k|    fn ensure_child_has_at_least_t(&mut self, parent_id: NodeId, index: usize) -> usize {
  774|  86.3k|        let child_id = self.node(parent_id).children[index];
  775|  86.3k|        if self.node(child_id).entries.len() >= T {
  776|  70.4k|            return index;
  777|  15.9k|        }
  778|       |
  779|  15.9k|        if index > 0 {
  780|  7.08k|            let left_id = self.node(parent_id).children[index - 1];
  781|  7.08k|            if self.node(left_id).entries.len() >= T {
  782|  4.15k|                self.borrow_from_left(parent_id, index);
  783|  4.15k|                return index;
  784|  2.93k|            }
  785|  8.85k|        }
  786|       |
  787|  11.7k|        if index + 1 < self.node(parent_id).children.len() {
  788|  10.3k|            let right_id = self.node(parent_id).children[index + 1];
  789|  10.3k|            if self.node(right_id).entries.len() >= T {
  790|  5.02k|                self.borrow_from_right(parent_id, index);
  791|  5.02k|                return index;
  792|  5.31k|            }
  793|  1.45k|        }
  794|       |
  795|  6.76k|        if index > 0 {
  796|  2.02k|            self.merge_children_around_entry(parent_id, index - 1);
  797|  2.02k|            index - 1
  798|       |        } else {
  799|  4.73k|            self.merge_children_around_entry(parent_id, index);
  800|  4.73k|            index
  801|       |        }
  802|  86.3k|    }
  803|       |
  804|       |    /// Rotates the separator down into `children[index]` and the left
  805|       |    /// sibling's max entry up into the parent.
  806|  4.15k|    fn borrow_from_left(&mut self, parent_id: NodeId, index: usize) {
  807|  4.15k|        let left_id = self.node(parent_id).children[index - 1];
  808|  4.15k|        let child_id = self.node(parent_id).children[index];
  809|       |
  810|  4.15k|        let up_entry = self
  811|  4.15k|            .node_mut(left_id)
  812|  4.15k|            .entries
  813|  4.15k|            .pop()
  814|  4.15k|            .expect("borrow_from_left from empty sibling");
  815|  4.15k|        let down_entry = mem::replace(&mut self.node_mut(parent_id).entries[index - 1], up_entry);
  816|  4.15k|        self.node_mut(child_id).entries.insert(0, down_entry);
  817|       |
  818|  4.15k|        if !self.node(child_id).leaf {
  819|  2.05k|            let moved_child = self
  820|  2.05k|                .node_mut(left_id)
  821|  2.05k|                .children
  822|  2.05k|                .pop()
  823|  2.05k|                .expect("internal sibling has children");
  824|  2.05k|            self.node_mut(child_id).children.insert(0, moved_child);
  825|  2.10k|        }
  826|       |
  827|  4.15k|        self.pull(left_id);
  828|  4.15k|        self.pull(child_id);
  829|  4.15k|        self.pull(parent_id);
  830|  4.15k|    }
  831|       |
  832|       |    /// Rotates the separator down into `children[index]` and the right
  833|       |    /// sibling's min entry up into the parent.
  834|  5.02k|    fn borrow_from_right(&mut self, parent_id: NodeId, index: usize) {
  835|  5.02k|        let child_id = self.node(parent_id).children[index];
  836|  5.02k|        let right_id = self.node(parent_id).children[index + 1];
  837|       |
  838|  5.02k|        let up_entry = self.node_mut(right_id).entries.remove(0);
  839|  5.02k|        let down_entry = mem::replace(&mut self.node_mut(parent_id).entries[index], up_entry);
  840|  5.02k|        self.node_mut(child_id).entries.push(down_entry);
  841|       |
  842|  5.02k|        if !self.node(child_id).leaf {
  843|  2.17k|            let moved_child = self.node_mut(right_id).children.remove(0);
  844|  2.17k|            self.node_mut(child_id).children.push(moved_child);
  845|  2.84k|        }
  846|       |
  847|  5.02k|        self.pull(child_id);
  848|  5.02k|        self.pull(right_id);
  849|  5.02k|        self.pull(parent_id);
  850|  5.02k|    }
  851|       |
  852|       |    /// Merges `children[index]`, `entries[index]`, and `children[index + 1]`
  853|       |    /// into the left child. Frees the right child and returns the merged id.
  854|  7.39k|    fn merge_children_around_entry(&mut self, parent_id: NodeId, index: usize) -> NodeId {
  855|  7.39k|        let left_id = self.node(parent_id).children[index];
  856|  7.39k|        let right_id = self.node(parent_id).children[index + 1];
  857|       |
  858|  7.39k|        let separator = self.node_mut(parent_id).entries.remove(index);
  859|  7.39k|        self.node_mut(parent_id).children.remove(index + 1);
  860|       |
  861|  7.39k|        let right = self.nodes[right_id.0].take().expect("dangling NodeId");
  862|  7.39k|        self.free_list.push(right_id);
  863|       |
  864|  7.39k|        let left = self.node_mut(left_id);
  865|  7.39k|        left.entries.push(separator);
  866|  7.39k|        left.entries.extend(right.entries);
  867|  7.39k|        left.children.extend(right.children);
  868|       |
  869|  7.39k|        self.pull(left_id);
  870|  7.39k|        self.pull(parent_id);
  871|  7.39k|        left_id
  872|  7.39k|    }
  873|       |}
  874|       |
  875|       |#[cfg(test)]
  876|       |impl<P: Ord + Clone, V, const T: usize, const N: usize> AffinityBTreeQueue<P, V, T, N> {
  877|      1|    pub(crate) fn debug_set_next_seq(&mut self, seq: Seq) {
  878|      1|        self.next_seq = seq;
  879|      1|    }
  880|       |
  881|   424k|    fn for_each_entry_in_order<F: FnMut(&Entry<P, V, N>)>(&self, mut f: F) {
  882|       |        enum Item {
  883|       |            Visit(NodeId),
  884|       |            Emit(NodeId, usize),
  885|       |        }
  886|       |
  887|   424k|        let Some(root) = self.root else { return };
                               ^423k                    ^1.37k
  888|   423k|        let mut stack = Vec::new();
  889|   423k|        stack.push(Item::Visit(root));
  890|       |
  891|  7.15M|        while let Some(item) = stack.pop() {
                                     ^6.73M
  892|  6.73M|            match item {
  893|  4.16M|                Item::Visit(id) => {
  894|  4.16M|                    let x = self.node(id);
  895|  4.16M|                    if x.leaf {
  896|  29.8M|                        for e in &x.entries {
                                               ^2.99M
  897|  29.8M|                            f(e);
  898|  29.8M|                        }
  899|       |                    } else {
  900|  1.17M|                        let n = x.entries.len();
  901|       |                        // Push in reverse so child[0] is processed first:
  902|       |                        // child[0], entry[0], child[1], ..., entry[n-1], child[n].
  903|  1.17M|                        stack.push(Item::Visit(x.children[n]));
  904|  1.17M|                        let mut i = n;
  905|  3.74M|                        while i > 0 {
  906|  2.56M|                            i -= 1;
  907|  2.56M|                            stack.push(Item::Emit(id, i));
  908|  2.56M|                            stack.push(Item::Visit(x.children[i]));
  909|  2.56M|                        }
  910|       |                    }
  911|       |                }
  912|  2.56M|                Item::Emit(id, i) => f(&self.node(id).entries[i]),
  913|       |            }
  914|       |        }
  915|   424k|    }
  916|       |
  917|   141k|    pub(crate) fn debug_count_entries(&self) -> usize {
  918|   141k|        let mut count = 0;
  919|   141k|        self.for_each_entry_in_order(|_| count += 1);
  920|   141k|        count
  921|   141k|    }
  922|       |
  923|   142k|    pub(crate) fn debug_all_keys_in_order(&self) -> Vec<PriorityKey<P>> {
  924|   142k|        let mut out = Vec::new();
  925|  11.1M|        self.for_each_entry_in_order(|e| out.push(e.key.clone()));
                      ^142k^142k
  926|   142k|        out
  927|   142k|    }
  928|       |
  929|       |    /// Asserts every structural, key-order, augmentation, and arena invariant.
  930|   143k|    pub(crate) fn check_invariants(&self) {
  931|   143k|        for id in &self.free_list {
                          ^82.1k
  932|  82.1k|            assert!(
  933|  82.1k|                self.nodes[id.0].is_none(),
  934|       |                "free_list entry points to a live node"
  935|       |            );
  936|       |        }
  937|       |
  938|   143k|        let Some(root) = self.root else {
                               ^142k
  939|    711|            assert_eq!(self.len, 0, "empty tree must have len 0");
  940|    711|            assert!(
  941|    798|                self.nodes.iter().all(|n| n.is_none()),
                              ^711              ^711
  942|       |                "empty tree must have no live nodes"
  943|       |            );
  944|    711|            assert_eq!(self.free_list.len(), self.nodes.len());
  945|    711|            return;
  946|       |        };
  947|   142k|        assert!(self.len > 0, "non-empty root with len 0");
  948|       |
  949|       |        struct WorkItem<P> {
  950|       |            id: NodeId,
  951|       |            is_root: bool,
  952|       |            min_key: Option<PriorityKey<P>>,
  953|       |            max_key: Option<PriorityKey<P>>,
  954|       |            depth: usize,
  955|       |        }
  956|       |
  957|   142k|        let mut total_entries = 0usize;
  958|   142k|        let mut total_nodes = 0usize;
  959|   142k|        let mut leaf_depth: Option<usize> = None;
  960|       |
  961|   142k|        let mut stack: Vec<WorkItem<P>> = Vec::new();
  962|   142k|        stack.push(WorkItem {
  963|   142k|            id: root,
  964|   142k|            is_root: true,
  965|   142k|            min_key: None,
  966|   142k|            max_key: None,
  967|   142k|            depth: 0,
  968|   142k|        });
  969|       |
  970|  1.55M|        while let Some(item) = stack.pop() {
                                     ^1.40M
  971|  1.40M|            let x = self.node(item.id);
  972|  1.40M|            total_nodes += 1;
  973|  1.40M|            total_entries += x.entries.len();
  974|       |
  975|  1.40M|            assert!(x.entries.len() <= Self::MAX_KEYS, "node overflow");
  976|  1.40M|            if item.is_root {
  977|   142k|                assert!(!x.entries.is_empty(), "root must be non-empty");
  978|       |            } else {
  979|  1.26M|                assert!(x.entries.len() >= Self::MIN_KEYS, "non-root underflow");
  980|       |            }
  981|       |
  982|  9.69M|            for w in x.entries.windows(2) {
                                   ^1.40M
  983|  9.69M|                assert!(w[0].key < w[1].key, "node entries not strictly sorted");
  984|       |            }
  985|  11.1M|            for e in &x.entries {
                                   ^1.40M
  986|  11.1M|                if let Some(m) = &item.min_key {
                                          ^8.27M
  987|  8.27M|                    assert!(*m < e.key, "entry below child key range");
  988|  2.82M|                }
  989|  11.1M|                if let Some(m) = &item.max_key {
                                          ^8.26M
  990|  8.26M|                    assert!(e.key < *m, "entry above child key range");
  991|  2.83M|                }
  992|  11.1M|                assert!(!e.affinity.is_empty(), "entry with empty affinity");
  993|  11.1M|                assert!(
  994|  11.1M|                    e.affinity.fits_within(self.num_cpus),
  995|       |                    "entry affinity outside num_cpus"
  996|       |                );
  997|       |            }
  998|       |
  999|  1.40M|            let mut mask = CpuMask::empty();
 1000|  11.1M|            for e in &x.entries {
                                   ^1.40M
 1001|  11.1M|                mask = mask.union(e.affinity);
 1002|  11.1M|            }
 1003|       |
 1004|  1.40M|            if x.leaf {
 1005|  1.01M|                assert!(x.children.is_empty(), "leaf with children");
 1006|  1.01M|                match leaf_depth {
 1007|   142k|                    None => leaf_depth = Some(item.depth),
 1008|   870k|                    Some(d) => assert_eq!(d, item.depth, "leaves at different depths"),
 1009|       |                }
 1010|       |            } else {
 1011|   394k|                assert_eq!(
 1012|   394k|                    x.children.len(),
 1013|   394k|                    x.entries.len() + 1,
 1014|       |                    "internal node child count"
 1015|       |                );
 1016|  1.26M|                for ci in 0..x.children.len() {
                                        ^394k
 1017|  1.26M|                    let child_id = x.children[ci];
 1018|  1.26M|                    mask = mask.union(self.node(child_id).subtree_affinity);
 1019|  1.26M|                    stack.push(WorkItem {
 1020|  1.26M|                        id: child_id,
 1021|       |                        is_root: false,
 1022|  1.26M|                        min_key: if ci == 0 {
 1023|   394k|                            item.min_key.clone()
 1024|       |                        } else {
 1025|   870k|                            Some(x.entries[ci - 1].key.clone())
 1026|       |                        },
 1027|  1.26M|                        max_key: if ci == x.entries.len() {
 1028|   394k|                            item.max_key.clone()
 1029|       |                        } else {
 1030|   870k|                            Some(x.entries[ci].key.clone())
 1031|       |                        },
 1032|  1.26M|                        depth: item.depth + 1,
 1033|       |                    });
 1034|       |                }
 1035|       |            }
 1036|       |
 1037|  1.40M|            assert_eq!(mask, x.subtree_affinity, "stale subtree_affinity");
 1038|       |        }
 1039|       |
 1040|   142k|        assert_eq!(total_entries, self.len, "entry count != len");
 1041|  1.48M|        let live = self.nodes.iter().filter(|n| n.is_some()).count();
                          ^142k  ^142k             ^142k                   ^142k
 1042|   142k|        assert_eq!(live, total_nodes, "unreachable live nodes in arena");
 1043|   142k|        assert_eq!(live + self.free_list.len(), self.nodes.len());
 1044|       |
 1045|       |        {
 1046|   142k|            let keys = self.debug_all_keys_in_order();
 1047|  10.9M|            for w in keys.windows(2) {
                                   ^142k
 1048|  10.9M|                assert!(w[0] < w[1], "in-order keys not strictly sorted");
 1049|       |            }
 1050|       |        }
 1051|   143k|    }
 1052|       |}
 1053|       |
 1054|       |#[cfg(test)]
 1055|       |impl<P: Ord + Clone, V: Clone, const T: usize, const N: usize> AffinityBTreeQueue<P, V, T, N> {
 1056|   141k|    pub(crate) fn debug_all_entries_in_order(&self) -> Vec<(PriorityKey<P>, CpuMask<N>, V)> {
 1057|   141k|        let mut out = Vec::new();
 1058|  10.6M|        self.for_each_entry_in_order(|e| out.push((e.key.clone(), e.affinity, e.value.clone())));
                      ^141k^141k
 1059|   141k|        out
 1060|   141k|    }
 1061|       |}
 1062|       |
 1063|       |/// Deliberately simple Vec-based reference model used by tests to check the
 1064|       |/// B-tree's observable behavior.
 1065|       |#[cfg(test)]
 1066|       |#[derive(Clone, Debug, PartialEq, Eq)]
 1067|       |struct RefEntry<P, V, const N: usize> {
 1068|       |    key: PriorityKey<P>,
 1069|       |    affinity: CpuMask<N>,
 1070|       |    value: V,
 1071|       |}
 1072|       |
 1073|       |#[cfg(test)]
 1074|       |#[derive(Clone, Debug)]
 1075|       |struct RefModel<P, V, const N: usize = 1> {
 1076|       |    entries: Vec<RefEntry<P, V, N>>,
 1077|       |    next_seq: Seq,
 1078|       |    num_cpus: usize,
 1079|       |}
 1080|       |
 1081|       |#[cfg(test)]
 1082|       |impl<P: Ord + Clone, V: Clone, const N: usize> RefModel<P, V, N> {
 1083|    169|    fn new(num_cpus: usize) -> Self {
 1084|    169|        Self {
 1085|    169|            entries: Vec::new(),
 1086|    169|            next_seq: 0,
 1087|    169|            num_cpus,
 1088|    169|        }
 1089|    169|    }
 1090|       |
 1091|  70.5k|    fn push(
 1092|  70.5k|        &mut self,
 1093|  70.5k|        priority: P,
 1094|  70.5k|        affinity: CpuMask<N>,
 1095|  70.5k|        value: V,
 1096|  70.5k|    ) -> Result<PriorityKey<P>, PushError> {
 1097|  70.5k|        if affinity.is_empty() {
 1098|      1|            return Err(PushError::EmptyAffinity);
 1099|  70.5k|        }
 1100|  70.5k|        if !affinity.fits_within(self.num_cpus) {
 1101|      1|            return Err(PushError::InvalidCpu);
 1102|  70.5k|        }
 1103|  70.5k|        let seq = self.next_seq;
 1104|  70.5k|        let next_seq = seq.checked_add(1).ok_or(PushError::SequenceOverflow)?;
                                                                                          ^0
 1105|  70.5k|        self.next_seq = next_seq;
 1106|       |
 1107|  70.5k|        let key = PriorityKey { priority, seq };
 1108|  70.5k|        self.entries.push(RefEntry {
 1109|  70.5k|            key: key.clone(),
 1110|  70.5k|            affinity,
 1111|  70.5k|            value,
 1112|  70.5k|        });
 1113|  70.5k|        Ok(key)
 1114|  70.5k|    }
 1115|       |
 1116|  1.71M|    fn peek_key_for_cpu(&self, cpu: usize) -> Option<PriorityKey<P>> {
 1117|  1.71M|        if cpu >= self.num_cpus {
 1118|      1|            return None;
 1119|  1.71M|        }
 1120|  1.71M|        self.entries
 1121|  1.71M|            .iter()
 1122|   120M|            .filter(|e| e.affinity.contains(cpu))
                           ^1.71M
 1123|  62.1M|            .map(|e| e.key.clone())
                           ^1.71M
 1124|  1.71M|            .min()
 1125|  1.71M|    }
 1126|       |
 1127|  24.1k|    fn pop_for_cpu(&mut self, cpu: usize) -> Option<(PriorityKey<P>, CpuMask<N>, V)> {
 1128|  24.1k|        let key = self.peek_key_for_cpu(cpu)?;
                          ^23.8k                          ^264
 1129|  1.56M|        let index = self.entries.iter().position(|e| e.key == key).unwrap();
                          ^23.8k  ^23.8k              ^23.8k                     ^23.8k
 1130|  23.8k|        let entry = self.entries.remove(index);
 1131|  23.8k|        Some((entry.key, entry.affinity, entry.value))
 1132|  24.1k|    }
 1133|       |
 1134|  47.1k|    fn remove(&mut self, key: &PriorityKey<P>) -> Option<(PriorityKey<P>, CpuMask<N>, V)> {
 1135|  2.67M|        let index = self.entries.iter().position(|e| e.key == *key)?;
                          ^23.4k  ^47.1k              ^47.1k                     ^23.7k
 1136|  23.4k|        let entry = self.entries.remove(index);
 1137|  23.4k|        Some((entry.key, entry.affinity, entry.value))
 1138|  47.1k|    }
 1139|       |
 1140|   282k|    fn len(&self) -> usize {
 1141|   282k|        self.entries.len()
 1142|   282k|    }
 1143|       |
 1144|    501|    fn is_empty(&self) -> bool {
 1145|    501|        self.entries.is_empty()
 1146|    501|    }
 1147|       |
 1148|   141k|    fn entries_in_order(&self) -> Vec<(PriorityKey<P>, CpuMask<N>, V)> {
 1149|   141k|        let mut out: Vec<_> = self
 1150|   141k|            .entries
 1151|   141k|            .iter()
 1152|  10.6M|            .map(|e| (e.key.clone(), e.affinity, e.value.clone()))
                           ^141k
 1153|   141k|            .collect();
 1154|  73.5M|        out.sort_by(|a, b| a.0.cmp(&b.0));
                      ^141k^141k
 1155|   141k|        out
 1156|   141k|    }
 1157|       |}
 1158|       |
 1159|       |#[cfg(test)]
 1160|       |mod tests;