// Copyright (c) Meta Platforms, Inc. and affiliates.

// This software may be used and distributed according to the terms of the
// GNU General Public License version 2.
pub mod alloc;
mod config;
pub mod layer_core_growth;

pub mod bpf_intf;

use plain::Plain;

unsafe impl Plain for bpf_intf::cpu_ctx {}
unsafe impl Plain for bpf_intf::llc_ctx {}
unsafe impl Plain for bpf_intf::node_ctx {}
unsafe impl Plain for bpf_intf::refresh_node_ctx_arg {}

use std::collections::BTreeMap;

use anyhow::bail;
use anyhow::Result;
pub use config::LayerCommon;
pub use config::LayerConfig;
pub use config::LayerKind;
pub use config::LayerMatch;
pub use config::LayerPlacement;
pub use config::LayerSpec;
pub use layer_core_growth::LayerGrowthAlgo;
use scx_utils::Core;
use scx_utils::Cpumask;
use scx_utils::Topology;
use scx_utils::NR_CPUS_POSSIBLE;
use scx_utils::NR_CPU_IDS;
use std::sync::Arc;
use tracing::info;

const MAX_CPUS: usize = bpf_intf::consts_MAX_CPUS as usize;

/// Divide `total` into parts proportional to `quotas`, returning exact
/// integers that sum to `total`. Uses the largest-remainder method
/// (Hamilton's method) for fair rounding.
pub fn largest_remainder(total: usize, quotas: &[f64]) -> Vec<usize> {
    if quotas.is_empty() {
        return vec![];
    }

    let sum: f64 = quotas.iter().sum();
    if sum == 0.0 {
        // No demand — distribute nothing.
        return vec![0; quotas.len()];
    }

    // Scale quotas so they sum to `total`.
    let scaled: Vec<f64> = quotas.iter().map(|q| q / sum * total as f64).collect();

    // Floor each scaled value.
    let floors: Vec<usize> = scaled.iter().map(|s| *s as usize).collect();
    let floor_sum: usize = floors.iter().sum();
    let mut remainder = total.saturating_sub(floor_sum);

    // Sort indices by descending fractional part.
    let mut indices: Vec<usize> = (0..quotas.len()).collect();
    indices.sort_by(|&a, &b| {
        let frac_a = scaled[a] - floors[a] as f64;
        let frac_b = scaled[b] - floors[b] as f64;
        frac_b
            .partial_cmp(&frac_a)
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    let mut result = floors;
    for &i in &indices {
        if remainder == 0 {
            break;
        }
        result[i] += 1;
        remainder -= 1;
    }

    result
}

/// Round CPU targets up to alloc-unit multiples, capped at total_cpus.
pub fn round_targets_to_alloc_units(
    targets: &[(usize, usize)],
    alloc_unit: usize,
    total_cpus: usize,
) -> Vec<(usize, usize)> {
    if alloc_unit <= 1 {
        return targets.to_vec();
    }

    targets
        .iter()
        .map(|&(target, min)| {
            // Round target up to next alloc_unit multiple.
            let aligned = target.div_ceil(alloc_unit) * alloc_unit;
            // Cap at total_cpus.
            let aligned = aligned.min(total_cpus);
            // Round min up similarly.
            let min_aligned = min.div_ceil(alloc_unit) * alloc_unit;
            let min_aligned = min_aligned.min(total_cpus);
            (aligned, min_aligned)
        })
        .collect()
}

#[derive(Debug)]
/// `CpuPool` represents the CPU core and logical CPU topology within the system.
/// It manages the mapping and availability of physical and logical cores, including
/// how resources are allocated for tasks across the available CPUs.
pub struct CpuPool {
    pub topo: Arc<Topology>,

    /// Per-node free LLC pool: node_id → [(llc_id, consumed_cores)].
    /// StickyDynamic's LLC trading operates within each node's list,
    /// preferring same-node LLCs to avoid cross-node traffic.
    pub free_llcs: BTreeMap<usize, Vec<(usize, usize)>>,

    /// A mask for available CPUs (SMT hyperthreads).
    /// Use for sub-core allocations when turned on.
    available_cpus: Cpumask,

    /// The ID of the first physical CPU in the system.
    /// Used as a global default when no per-node CPU is suitable.
    #[allow(dead_code)]
    first_cpu: usize,

    /// Per-node fallback CPU: node_id → cpu_id. Each node has its own
    /// fallback CPU for draining empty-layer DSQs, avoiding cross-node traffic.
    pub fallback_cpus: BTreeMap<usize, usize>,

    /// Dense sequential core index (0, 1, 2, ...) assigned by walking
    /// topo.nodes → node.llcs → llc.cores in BTreeMap order. Hardware
    /// core IDs can have gaps (e.g. Ryzen); this provides a contiguous
    /// index space that preserves topological locality — all cores in
    /// node 0 get the lowest indices, then node 1, etc. Growth
    /// algorithms use these indices to define core allocation order.
    core_seq: BTreeMap<(usize, usize, usize), usize>,

    allow_partial: bool,
}

impl CpuPool {
    pub fn new(topo: Arc<Topology>, allow_partial: bool) -> Result<Self> {
        if *NR_CPU_IDS > MAX_CPUS {
            bail!("NR_CPU_IDS {} > MAX_CPUS {}", *NR_CPU_IDS, MAX_CPUS);
        }

        // Build core_seq
        let mut core_seq = BTreeMap::new();
        let mut next_seq: usize = 0;
        for node in topo.nodes.values() {
            for llc in node.llcs.values() {
                for core in llc.cores.values() {
                    core_seq.insert((core.node_id, core.llc_id, core.id), next_seq);
                    next_seq += 1;
                }
            }
        }

        info!(
            "CPUs: online/possible={}/{} nr_cores={}",
            topo.all_cpus.len(),
            *NR_CPUS_POSSIBLE,
            topo.all_cores.len(),
        );

        let first_cpu = *topo.all_cpus.keys().next().unwrap();

        let mut free_llcs: BTreeMap<usize, Vec<(usize, usize)>> = BTreeMap::new();
        for llc in topo.all_llcs.values() {
            free_llcs.entry(llc.node_id).or_default().push((llc.id, 0));
        }

        let mut available_cpus = Cpumask::new();
        available_cpus.set_all();

        let mut cpu_pool = Self {
            free_llcs,
            available_cpus,
            first_cpu,
            fallback_cpus: BTreeMap::new(),
            core_seq,
            topo,
            allow_partial,
        };
        cpu_pool.update_fallback_cpus();
        Ok(cpu_pool)
    }

    fn update_fallback_cpus(&mut self) {
        for node in self.topo.nodes.values() {
            let Some(fb) = self
                .available_cpus
                .and(&node.span)
                .iter()
                .next()
                .or_else(|| node.span.iter().next())
            else {
                continue;
            };
            self.fallback_cpus.insert(node.id, fb);
        }
    }

    fn core_cpu_available(&self, core: &Cpumask) -> usize {
        core.iter()
            .map(|cpu| self.available_cpus.test_cpu(cpu) as usize)
            .sum()
    }

    fn check_partial(&self) -> Result<()> {
        if self.allow_partial {
            return Ok(());
        }

        // Go through CPU to cores and check if any core has
        // a span of 1. If it does, we bail.
        for i in 0..self.topo.all_cores.len() {
            let core_cpus = &self.topo.all_cores[&i].span;

            let core_cpu_available = self.core_cpu_available(core_cpus);
            if core_cpu_available > 0 && core_cpu_available != core_cpus.weight() {
                bail!("Some cores only partially allocated");
            }
        }

        Ok(())
    }

    pub fn alloc_cpus(
        &mut self,
        allowed_cpus: &Cpumask,
        core_alloc_order: &[usize],
        mut max_to_alloc: usize,
    ) -> Option<Cpumask> {
        let mut allocated_cpus = Cpumask::new();

        for alloc_core in core_alloc_order {
            // Constrain CPUs by NUMA node or LLC. Since allowed_cpus is NUMA/LLC aligned,
            // this operation is guaranteed to produce either the core mask or an empty mask.
            let core_cpus = &self.topo.all_cores[alloc_core].span.and(allowed_cpus);
            if core_cpus.is_empty() {
                continue;
            }

            let available_core_cpus = core_cpus.and(&self.available_cpus);
            let core_num_available = available_core_cpus.weight();
            if core_num_available == 0 {
                continue;
            }

            if !self.allow_partial || max_to_alloc >= core_num_available {
                allocated_cpus = allocated_cpus.or(&available_core_cpus);
                self.available_cpus = self.available_cpus.and(&available_core_cpus.not());
            } else {
                let cpus = available_core_cpus.iter().take(max_to_alloc);
                for cpu in cpus {
                    allocated_cpus.set_cpu(cpu).ok()?;
                    self.available_cpus.clear_cpu(cpu).ok()?;
                }
            }

            // Are we done allocating CPUs?
            if max_to_alloc <= core_num_available {
                break;
            }

            max_to_alloc -= core_num_available;
        }

        self.update_fallback_cpus();
        self.check_partial().unwrap();

        if !allocated_cpus.is_empty() {
            Some(allocated_cpus)
        } else {
            None
        }
    }

    pub fn free(&mut self, cpus_to_free: &Cpumask) -> Result<()> {
        // Whether we allow partial CPUs or not does not matter because the cpumask
        // provided is create in next_to_free() below. If we do not allow partial
        // partial allocations, the cpumask provided will be core-aligned.

        if !self.available_cpus.and(cpus_to_free).is_empty() {
            bail!("Some of CPUs {} are already free", cpus_to_free);
        }

        self.available_cpus = self.available_cpus.or(cpus_to_free);
        self.update_fallback_cpus();

        self.check_partial()?;

        Ok(())
    }

    pub fn mark_allocated(&mut self, cpus_to_alloc: &Cpumask) -> Result<()> {
        if cpus_to_alloc.and(&self.available_cpus.not()).weight() > 0 {
            bail!(
                "Some of CPUs {} are not available to allocate",
                cpus_to_alloc
            );
        }

        self.available_cpus &= &cpus_to_alloc.not();
        self.update_fallback_cpus();

        self.check_partial()?;

        Ok(())
    }

    pub fn next_to_free<'a>(
        &'a self,
        cands: &Cpumask,
        core_order: impl Iterator<Item = &'a usize>,
    ) -> Result<Option<Cpumask>> {
        for pref_core in core_order.map(|i| &self.topo.all_cores[i]) {
            let pref_cpus = pref_core.span.and(cands);
            if pref_cpus.weight() > 0 {
                return Ok(Some(pref_cpus));
            }
        }
        Ok(None)
    }

    pub fn available_cpus(&self) -> Cpumask {
        self.available_cpus.clone()
    }

    fn core_seq(&self, core: &Core) -> usize {
        *self
            .core_seq
            .get(&(core.node_id, core.llc_id, core.id))
            .expect("unrecognised core")
    }

    /// Allocation unit size: cpus_per_core when `!allow_partial`, 1 otherwise.
    /// Budget targets should be multiples of this to avoid double-rounding.
    pub fn alloc_unit(&self) -> usize {
        if self.allow_partial {
            1
        } else {
            self.topo
                .all_cores
                .first_key_value()
                .map_or(1, |(_, c)| c.cpus.len())
        }
    }

    /// Pop a free LLC from the given node. Returns the LLC ID if available.
    pub fn take_llc_from_node(&mut self, node_id: usize) -> Option<usize> {
        self.free_llcs
            .get_mut(&node_id)
            .and_then(|v| v.pop())
            .map(|(llc_id, _)| llc_id)
    }

    /// Pop a free LLC, trying nodes in the given order.
    pub fn take_llc(&mut self, node_order: &[usize]) -> Option<usize> {
        for &node in node_order {
            if let Some(llc) = self.take_llc_from_node(node) {
                return Some(llc);
            }
        }
        None
    }

    /// Return an LLC to the correct node's free list.
    pub fn return_llc(&mut self, llc_id: usize) {
        let node_id = self.topo.all_llcs[&llc_id].node_id;
        self.free_llcs.get_mut(&node_id).unwrap().push((llc_id, 0));
    }

    /// Total free LLCs across all nodes.
    pub fn total_free_llcs(&self) -> usize {
        self.free_llcs.values().map(|v| v.len()).sum()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use scx_utils::testutils::{make_het_test_topo, make_test_topo, mask_from_bits};
    use scx_utils::Llc;

    // 1N: 1 node, 2 LLCs, 4 cores/LLC, 2 HTs/core = 16 CPUs
    //   LLC0: cores 0-3 (cpus 0-7), LLC1: cores 4-7 (cpus 8-15)
    fn topo_1n() -> (Arc<Topology>, usize) {
        let (topo, total) = make_test_topo(1, 2, 4, 2);
        (Arc::new(topo), total)
    }

    // 2N: 2 nodes, 2 LLCs/node, 4 cores/LLC, 2 HTs/core = 32 CPUs
    //   Node0: LLC0 (cores 0-3, cpus 0-7), LLC1 (cores 4-7, cpus 8-15)
    //   Node1: LLC2 (cores 8-11, cpus 16-23), LLC3 (cores 12-15, cpus 24-31)
    fn topo_2n() -> (Arc<Topology>, usize) {
        let (topo, total) = make_test_topo(2, 2, 4, 2);
        (Arc::new(topo), total)
    }

    // 4N: 4 nodes, 2 LLCs/node, 2 cores/LLC, 2 HTs/core = 32 CPUs
    //   Node0: LLC0 (cores 0-1, cpus 0-3), LLC1 (cores 2-3, cpus 4-7)
    //   Node1: LLC2 (cores 4-5, cpus 8-11), LLC3 (cores 6-7, cpus 12-15)
    //   Node2: LLC4 (cores 8-9, cpus 16-19), LLC5 (cores 10-11, cpus 20-23)
    //   Node3: LLC6 (cores 12-13, cpus 24-27), LLC7 (cores 14-15, cpus 28-31)
    fn topo_4n() -> (Arc<Topology>, usize) {
        let (topo, total) = make_test_topo(4, 2, 2, 2);
        (Arc::new(topo), total)
    }

    fn all_cpus_mask(total: usize) -> Cpumask {
        mask_from_bits(total, &(0..total).collect::<Vec<_>>())
    }

    fn core_order_sequential(topo: &Topology) -> Vec<usize> {
        topo.all_cores.keys().copied().collect()
    }

    fn core_order_per_node(topo: &Topology) -> Vec<Vec<usize>> {
        let nr_nodes = topo.nodes.len();
        let mut result = vec![Vec::new(); nr_nodes];
        for (&core_id, core) in &topo.all_cores {
            result[core.node_id].push(core_id);
        }
        result
    }

    // --- new() ---

    #[test]
    fn test_new_1n_all_available() {
        let (topo, total) = topo_1n();
        let pool = CpuPool::new(topo, false).unwrap();
        // All CPUs in the topology should be available.
        for cpu in 0..total {
            assert!(
                pool.available_cpus.test_cpu(cpu),
                "cpu {} not available",
                cpu
            );
        }
    }

    #[test]
    fn test_new_1n_fallback_cpus() {
        let (topo, _total) = topo_1n();
        let pool = CpuPool::new(topo, false).unwrap();
        // Single node: fallback_cpus[0] should be the first CPU (0).
        assert_eq!(pool.fallback_cpus.len(), 1);
        assert_eq!(pool.fallback_cpus[&0], 0);
    }

    #[test]
    fn test_new_1n_core_topo_ids() {
        let (topo, _total) = topo_1n();
        let pool = CpuPool::new(topo.clone(), false).unwrap();
        // All cores should have unique topology IDs assigned sequentially.
        for (i, core_id) in topo.all_cores.keys().enumerate() {
            let core = &topo.all_cores[core_id];
            assert_eq!(pool.core_seq(core), i);
        }
    }

    #[test]
    fn test_new_1n_free_llcs() {
        let (topo, _total) = topo_1n();
        let pool = CpuPool::new(topo, false).unwrap();
        // Should have 2 LLCs on node 0, each with consumed_count=0.
        assert_eq!(pool.total_free_llcs(), 2);
        let node0 = &pool.free_llcs[&0];
        assert_eq!(node0.len(), 2);
        assert_eq!(node0[0], (0, 0)); // (llc_id, consumed)
        assert_eq!(node0[1], (1, 0));
    }

    #[test]
    fn test_new_2n_all_available() {
        let (topo, total) = topo_2n();
        let pool = CpuPool::new(topo, false).unwrap();
        for cpu in 0..total {
            assert!(
                pool.available_cpus.test_cpu(cpu),
                "cpu {} not available",
                cpu
            );
        }
    }

    #[test]
    fn test_new_2n_free_llcs() {
        let (topo, _total) = topo_2n();
        let pool = CpuPool::new(topo, false).unwrap();
        // 4 LLCs across 2 nodes, grouped by node.
        assert_eq!(pool.total_free_llcs(), 4);
        assert_eq!(pool.free_llcs.len(), 2); // 2 nodes
        let node0 = &pool.free_llcs[&0];
        let node1 = &pool.free_llcs[&1];
        assert_eq!(node0.len(), 2);
        assert_eq!(node1.len(), 2);
    }

    // --- alloc_cpus() ---

    #[test]
    fn test_alloc_one_core() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Allocate 2 CPUs (one core with 2 HTs).
        let result = pool.alloc_cpus(&allowed, &order, 2);
        assert!(result.is_some());
        let alloc = result.unwrap();
        assert_eq!(alloc.weight(), 2);
        // Should be core 0: cpus 0,1.
        assert!(alloc.test_cpu(0));
        assert!(alloc.test_cpu(1));
        // Those CPUs should no longer be available.
        assert!(!pool.available_cpus.test_cpu(0));
        assert!(!pool.available_cpus.test_cpu(1));
    }

    #[test]
    fn test_alloc_respects_core_order() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        // Reverse core order: allocate from core 7 first.
        let order: Vec<usize> = core_order_sequential(&topo).into_iter().rev().collect();
        let mut pool = CpuPool::new(topo, false).unwrap();

        let result = pool.alloc_cpus(&allowed, &order, 2);
        assert!(result.is_some());
        let alloc = result.unwrap();
        // Core 7: cpus 14,15.
        assert!(alloc.test_cpu(14));
        assert!(alloc.test_cpu(15));
    }

    #[test]
    fn test_alloc_respects_allowed_cpus() {
        let (topo, total) = topo_1n();
        // Only allow node0 LLC1 (cpus 8-15).
        let allowed = mask_from_bits(total, &(8..16).collect::<Vec<_>>());
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        let result = pool.alloc_cpus(&allowed, &order, 2);
        assert!(result.is_some());
        let alloc = result.unwrap();
        // Should skip cores 0-3 (LLC0) and allocate from core 4 (first in LLC1).
        assert!(alloc.test_cpu(8));
        assert!(alloc.test_cpu(9));
    }

    #[test]
    fn test_alloc_multiple_cores() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Allocate 6 CPUs = 3 full cores.
        let result = pool.alloc_cpus(&allowed, &order, 6);
        assert!(result.is_some());
        let alloc = result.unwrap();
        assert_eq!(alloc.weight(), 6);
        // Cores 0,1,2 → cpus 0-5.
        for cpu in 0..6 {
            assert!(alloc.test_cpu(cpu));
        }
    }

    #[test]
    fn test_alloc_exhausts_returns_none() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Allocate all 16 CPUs.
        let result = pool.alloc_cpus(&allowed, &order, 16);
        assert!(result.is_some());
        assert_eq!(result.unwrap().weight(), 16);

        // Try to allocate more — should return None.
        let result = pool.alloc_cpus(&allowed, &order, 2);
        assert!(result.is_none());
    }

    #[test]
    fn test_alloc_more_than_available() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Ask for 100 but only 16 exist. Should allocate all 16.
        let result = pool.alloc_cpus(&allowed, &order, 100);
        assert!(result.is_some());
        assert_eq!(result.unwrap().weight(), 16);
    }

    #[test]
    fn test_alloc_partial_core() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, true).unwrap(); // allow_partial=true

        // Allocate 3 CPUs: 1 full core (2) + 1 partial.
        let result = pool.alloc_cpus(&allowed, &order, 3);
        assert!(result.is_some());
        let alloc = result.unwrap();
        assert_eq!(alloc.weight(), 3);
    }

    #[test]
    fn test_alloc_no_partial_rounds_to_core() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap(); // allow_partial=false

        // Allocate 3 CPUs with allow_partial=false. Since core has 2 HTs,
        // it should allocate the first core (2 CPUs) then stop at the
        // boundary of 3 which is ≤ the second core's 2, so it takes the
        // second core too, giving 4.
        let result = pool.alloc_cpus(&allowed, &order, 3);
        assert!(result.is_some());
        let alloc = result.unwrap();
        // With !allow_partial, it always takes full cores. Asking for 3
        // means: first core (2 CPUs, 3-2=1 remaining), second core has 2
        // available which is ≥ 1 remaining, so it takes the full second
        // core and breaks. Total: 4.
        assert_eq!(alloc.weight(), 4);
    }

    // --- free() ---

    #[test]
    fn test_free_returns_cpus() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        let alloc = pool.alloc_cpus(&allowed, &order, 2).unwrap();
        assert!(!pool.available_cpus.test_cpu(0));

        pool.free(&alloc).unwrap();
        assert!(pool.available_cpus.test_cpu(0));
        assert!(pool.available_cpus.test_cpu(1));
    }

    #[test]
    fn test_free_double_free_errors() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        let alloc = pool.alloc_cpus(&allowed, &order, 2).unwrap();
        pool.free(&alloc).unwrap();
        // Freeing again should error.
        assert!(pool.free(&alloc).is_err());
    }

    // --- mark_allocated() ---

    #[test]
    fn test_mark_allocated() {
        let (topo, total) = topo_1n();
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Mark core 0 (cpus 0,1) as allocated.
        let to_mark = mask_from_bits(total, &[0, 1]);
        pool.mark_allocated(&to_mark).unwrap();
        assert!(!pool.available_cpus.test_cpu(0));
        assert!(!pool.available_cpus.test_cpu(1));
        assert!(pool.available_cpus.test_cpu(2));
    }

    #[test]
    fn test_mark_allocated_already_taken_errors() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Allocate core 0.
        pool.alloc_cpus(&allowed, &order, 2).unwrap();
        // Trying to mark_allocated the same CPUs should error.
        let to_mark = mask_from_bits(total, &[0, 1]);
        assert!(pool.mark_allocated(&to_mark).is_err());
    }

    // --- next_to_free() ---

    #[test]
    fn test_next_to_free_finds_core() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Allocate cores 0,1 (4 CPUs).
        let alloc = pool.alloc_cpus(&allowed, &order, 4).unwrap();

        // next_to_free with reverse order should find core 1 first.
        let rev_order: Vec<usize> = order.iter().copied().rev().collect();
        let result = pool.next_to_free(&alloc, rev_order.iter()).unwrap();
        assert!(result.is_some());
        let to_free = result.unwrap();
        // Core 1 = cpus 2,3.
        assert!(to_free.test_cpu(2));
        assert!(to_free.test_cpu(3));
        assert_eq!(to_free.weight(), 2);
    }

    #[test]
    fn test_next_to_free_respects_cands() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Allocate cores 0,1,2 (6 CPUs).
        let _alloc = pool.alloc_cpus(&allowed, &order, 6).unwrap();

        // Restrict candidates to only core 1's CPUs (2,3).
        let cands = mask_from_bits(total, &[2, 3]);
        let result = pool.next_to_free(&cands, order.iter()).unwrap();
        assert!(result.is_some());
        let to_free = result.unwrap();
        assert!(to_free.test_cpu(2));
        assert!(to_free.test_cpu(3));
    }

    #[test]
    fn test_next_to_free_nothing_to_free() {
        let (topo, total) = topo_1n();
        let pool = CpuPool::new(topo.clone(), false).unwrap();
        let order = core_order_sequential(&topo);

        // No CPUs allocated, empty cands.
        let cands = mask_from_bits(total, &[]);
        let result = pool.next_to_free(&cands, order.iter()).unwrap();
        assert!(result.is_none());
    }

    // --- Round-trip ---

    #[test]
    fn test_alloc_free_roundtrip() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        let before_weight = pool.available_cpus().weight();

        let alloc = pool.alloc_cpus(&allowed, &order, 4).unwrap();
        assert_eq!(pool.available_cpus().weight(), before_weight - 4);

        pool.free(&alloc).unwrap();
        assert_eq!(pool.available_cpus().weight(), before_weight);
        // All topology CPUs should be available again.
        for cpu in 0..total {
            assert!(
                pool.available_cpus().test_cpu(cpu),
                "cpu {} not restored",
                cpu
            );
        }
    }

    #[test]
    fn test_fallback_cpus_updates() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        assert_eq!(pool.fallback_cpus[&0], 0);

        // Allocate core 0 (cpus 0,1). Fallback should move to cpu 2.
        let alloc = pool.alloc_cpus(&allowed, &order, 2).unwrap();
        assert_eq!(pool.fallback_cpus[&0], 2);

        // Free it back. Fallback should return to 0.
        pool.free(&alloc).unwrap();
        assert_eq!(pool.fallback_cpus[&0], 0);
    }

    // --- 2N: node-aware allocation ---

    #[test]
    fn test_2n_alloc_node_restricted() {
        let (topo, total) = topo_2n();
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Restrict to node 1 only (cpus 16-31).
        let node1_mask = mask_from_bits(total, &(16..32).collect::<Vec<_>>());
        let result = pool.alloc_cpus(&node1_mask, &order, 4);
        assert!(result.is_some());
        let alloc = result.unwrap();
        assert_eq!(alloc.weight(), 4);
        // All allocated CPUs should be on node 1.
        for cpu in alloc.iter() {
            assert!(cpu >= 16, "cpu {} is not on node 1", cpu);
        }
    }

    #[test]
    fn test_2n_alloc_full_picks_sequential() {
        let (topo, total) = topo_2n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // With sequential order and full allowed mask, allocation starts
        // from node 0. Document this as baseline behavior.
        let result = pool.alloc_cpus(&allowed, &order, 2);
        assert!(result.is_some());
        let alloc = result.unwrap();
        // Current behavior: picks core 0 (cpus 0,1) from node 0.
        assert!(alloc.test_cpu(0));
        assert!(alloc.test_cpu(1));
    }

    #[test]
    fn test_2n_alloc_across_nodes() {
        let (topo, total) = topo_2n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Allocate all of node 0 (16 CPUs) then more.
        let _alloc1 = pool.alloc_cpus(&allowed, &order, 16).unwrap();

        // Next allocation should come from node 1.
        let alloc2 = pool.alloc_cpus(&allowed, &order, 2).unwrap();
        for cpu in alloc2.iter() {
            assert!(cpu >= 16, "cpu {} should be on node 1", cpu);
        }
    }

    #[test]
    fn test_alloc_free_realloc_same_cpus() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Alloc 4 CPUs (cores 0-1), free, re-alloc — should get same CPUs.
        let alloc1 = pool.alloc_cpus(&allowed, &order, 4).unwrap();
        let cpus1: Vec<usize> = alloc1.iter().collect();
        pool.free(&alloc1).unwrap();
        let alloc2 = pool.alloc_cpus(&allowed, &order, 4).unwrap();
        let cpus2: Vec<usize> = alloc2.iter().collect();
        assert_eq!(cpus1, cpus2, "re-alloc should yield same CPUs");
    }

    #[test]
    fn test_alloc_free_partial_realloc() {
        let (topo, total) = topo_1n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();
        let initial = pool.available_cpus().weight();

        // Alloc 8 + 4 CPUs, then free the first 8.
        let alloc1 = pool.alloc_cpus(&allowed, &order, 8).unwrap();
        let _alloc2 = pool.alloc_cpus(&allowed, &order, 4).unwrap();
        assert_eq!(pool.available_cpus().weight(), initial - 12);

        pool.free(&alloc1).unwrap();
        assert_eq!(pool.available_cpus().weight(), initial - _alloc2.weight());

        // Re-alloc 4 — should come from the freed alloc1 CPUs.
        let alloc3 = pool.alloc_cpus(&allowed, &order, 4).unwrap();
        for cpu in alloc3.iter() {
            assert!(
                alloc1.test_cpu(cpu),
                "cpu {} should come from freed alloc1",
                cpu
            );
        }
    }

    #[test]
    fn test_alloc_free_realloc_2n_stays_on_node() {
        let (topo, total) = topo_2n();
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Allocate all CPUs, then free only node 1 CPUs.
        let allowed = all_cpus_mask(total);
        let _alloc_all = pool.alloc_cpus(&allowed, &order, total).unwrap();
        let node1_cpus = mask_from_bits(total, &(16..32).collect::<Vec<_>>());
        pool.free(&node1_cpus).unwrap();

        // Re-alloc with full allowed mask — should get node 1 CPUs (only ones free).
        let realloc = pool.alloc_cpus(&allowed, &order, 4).unwrap();
        for cpu in realloc.iter() {
            assert!(
                cpu >= 16,
                "cpu {} should be on node 1 (only free node)",
                cpu
            );
        }
    }

    #[test]
    fn test_2n_free_llcs() {
        let (topo, _total) = topo_2n();
        let pool = CpuPool::new(topo, false).unwrap();
        // 4 LLCs across 2 nodes, grouped by node.
        assert_eq!(pool.total_free_llcs(), 4);
        assert_eq!(pool.free_llcs.len(), 2); // 2 nodes
        let node0 = &pool.free_llcs[&0];
        let node1 = &pool.free_llcs[&1];
        assert_eq!(node0.len(), 2);
        assert_eq!(node1.len(), 2);
        // LLC 0,1 on node 0; LLC 2,3 on node 1.
        assert_eq!(node0[0].0, 0);
        assert_eq!(node0[1].0, 1);
        assert_eq!(node1[0].0, 2);
        assert_eq!(node1[1].0, 3);
    }

    // --- 4N: four-node topology ---

    #[test]
    fn test_new_4n_all_available() {
        let (topo, total) = topo_4n();
        let pool = CpuPool::new(topo, false).unwrap();
        for cpu in 0..total {
            assert!(
                pool.available_cpus.test_cpu(cpu),
                "cpu {} not available",
                cpu
            );
        }
    }

    #[test]
    fn test_new_4n_free_llcs() {
        let (topo, _total) = topo_4n();
        let pool = CpuPool::new(topo, false).unwrap();
        // 8 LLCs across 4 nodes, 2 per node.
        assert_eq!(pool.total_free_llcs(), 8);
        assert_eq!(pool.free_llcs.len(), 4); // 4 nodes
        for node_id in 0..4 {
            let node_llcs = &pool.free_llcs[&node_id];
            assert_eq!(node_llcs.len(), 2);
        }
    }

    #[test]
    fn test_4n_alloc_across_all_nodes() {
        let (topo, total) = topo_4n();
        let allowed = all_cpus_mask(total);
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Allocate 24 CPUs (fills nodes 0-2 completely, 8 cpus each).
        let _alloc1 = pool.alloc_cpus(&allowed, &order, 24).unwrap();

        // Next allocation should come from node 3 (cpus 24-31).
        let alloc2 = pool.alloc_cpus(&allowed, &order, 2).unwrap();
        for cpu in alloc2.iter() {
            assert!(cpu >= 24, "cpu {} should be on node 3", cpu);
        }
    }

    #[test]
    fn test_4n_alloc_node_restricted() {
        let (topo, total) = topo_4n();
        let order = core_order_sequential(&topo);
        let mut pool = CpuPool::new(topo, false).unwrap();

        // Restrict to node 2 only (cpus 16-23).
        let node2_mask = mask_from_bits(total, &(16..24).collect::<Vec<_>>());
        let result = pool.alloc_cpus(&node2_mask, &order, 4);
        assert!(result.is_some());
        let alloc = result.unwrap();
        assert_eq!(alloc.weight(), 4);
        for cpu in alloc.iter() {
            assert!(
                cpu >= 16 && cpu < 24,
                "cpu {} should be on node 2 (16-23)",
                cpu
            );
        }
    }

    // =========================================================================
    // Growth algorithm tests
    // =========================================================================

    fn test_layer_spec(name: &str, algo: LayerGrowthAlgo) -> LayerSpec {
        let json = r#"{"name":"_","matches":[],"kind":{"Confined":{"util_range":[0.0,1.0]}}}"#;
        let mut spec: LayerSpec = serde_json::from_str(json).unwrap();
        spec.name = name.to_string();
        spec.kind.common_mut().growth_algo = algo;
        spec
    }

    fn test_layer_spec_with_nodes(
        name: &str,
        algo: LayerGrowthAlgo,
        nodes: Vec<usize>,
    ) -> LayerSpec {
        let mut spec = test_layer_spec(name, algo);
        *spec.nodes_mut() = nodes;
        spec
    }

    fn test_layer_spec_with_llcs(name: &str, algo: LayerGrowthAlgo, llcs: Vec<usize>) -> LayerSpec {
        let mut spec = test_layer_spec(name, algo);
        *spec.llcs_mut() = llcs;
        spec
    }

    /// Verify a core_order contains all expected cores with no duplicates.
    fn assert_valid_core_order(order: &[usize], nr_cores: usize) {
        assert_eq!(
            order.len(),
            nr_cores,
            "core_order length {} != expected {}",
            order.len(),
            nr_cores
        );
        let mut seen = std::collections::HashSet::new();
        for &core in order {
            assert!(
                core < nr_cores,
                "core {} out of range [0, {})",
                core,
                nr_cores
            );
            assert!(seen.insert(core), "duplicate core {} in order", core);
        }
    }

    fn get_core_order(topo: &Arc<Topology>, specs: &[LayerSpec], layer_idx: usize) -> Vec<usize> {
        let pool = CpuPool::new(topo.clone(), false).unwrap();
        let orders = LayerGrowthAlgo::layer_core_orders(&pool, specs, topo).unwrap();
        let per_node = &orders[&layer_idx];
        // Flatten per-node vecs in node_order sequence to reconstruct the
        // same flat ordering that the old code produced.
        let all_nodes: Vec<&[usize]> = specs.iter().map(|s| s.nodes().as_slice()).collect();
        let norder =
            layer_core_growth::node_order(specs[layer_idx].nodes(), topo, layer_idx, &all_nodes);
        let mut flat = Vec::new();
        for nid in norder {
            flat.extend_from_slice(&per_node[nid]);
        }
        flat
    }

    fn get_core_order_per_node(
        topo: &Arc<Topology>,
        specs: &[LayerSpec],
        layer_idx: usize,
    ) -> Vec<Vec<usize>> {
        let pool = CpuPool::new(topo.clone(), false).unwrap();
        let orders = LayerGrowthAlgo::layer_core_orders(&pool, specs, topo).unwrap();
        orders[&layer_idx].clone()
    }

    // --- Sticky ---

    #[test]
    fn test_growth_sticky_1n() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::Sticky)];
        let order = get_core_order(&topo, &specs, 0);
        // Single layer, idx=0: per-LLC rotation cycles back to identity.
        assert_eq!(order, vec![0, 1, 2, 3, 4, 5, 6, 7]);
    }

    #[test]
    fn test_growth_sticky_2n() {
        let (topo, _total) = topo_2n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::Sticky)];
        let order = get_core_order(&topo, &specs, 0);
        assert_eq!(
            order,
            vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
        );
    }

    #[test]
    fn test_growth_sticky_multi_layer_offsets() {
        let (topo, _total) = topo_1n();
        let specs = vec![
            test_layer_spec("L0", LayerGrowthAlgo::Sticky),
            test_layer_spec("L1", LayerGrowthAlgo::Sticky),
            test_layer_spec("L2", LayerGrowthAlgo::Sticky),
        ];
        let o0 = get_core_order(&topo, &specs, 0);
        let o1 = get_core_order(&topo, &specs, 1);
        let o2 = get_core_order(&topo, &specs, 2);
        // Each layer should have a different starting point.
        assert_ne!(o0[0], o1[0], "L0 and L1 should start at different cores");
        assert_ne!(o1[0], o2[0], "L1 and L2 should start at different cores");
        // All should be valid.
        assert_valid_core_order(&o0, 8);
        assert_valid_core_order(&o1, 8);
        assert_valid_core_order(&o2, 8);
    }

    // --- Linear ---

    #[test]
    fn test_growth_linear_1n() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::Linear)];
        let order = get_core_order(&topo, &specs, 0);
        assert_eq!(order, vec![0, 1, 2, 3, 4, 5, 6, 7]);
    }

    #[test]
    fn test_growth_linear_2n() {
        let (topo, _total) = topo_2n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::Linear)];
        let order = get_core_order(&topo, &specs, 0);
        assert_eq!(
            order,
            vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
        );
    }

    #[test]
    fn test_growth_linear_preserves_topo_order_with_nodes() {
        let (topo, _total) = topo_2n();
        let specs = vec![test_layer_spec_with_nodes(
            "L0",
            LayerGrowthAlgo::Linear,
            vec![1],
        )];
        let order = get_core_order(&topo, &specs, 0);
        // With nodes=[1] hard limit, only node 1 cores appear.
        assert_eq!(order, vec![8, 9, 10, 11, 12, 13, 14, 15]);
    }

    // --- Reverse ---

    #[test]
    fn test_growth_reverse_1n() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::Reverse)];
        let order = get_core_order(&topo, &specs, 0);
        assert_eq!(order, vec![7, 6, 5, 4, 3, 2, 1, 0]);
    }

    // --- Random ---

    #[test]
    fn test_growth_random_1n() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::Random)];
        let order = get_core_order(&topo, &specs, 0);
        // Random is a seeded shuffle — not sequential.
        assert_valid_core_order(&order, 8);
        assert_ne!(order, vec![0, 1, 2, 3, 4, 5, 6, 7], "should be shuffled");
    }

    #[test]
    fn test_growth_random_deterministic_with_same_idx() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::Random)];
        let order1 = get_core_order(&topo, &specs, 0);
        let order2 = get_core_order(&topo, &specs, 0);
        // Same layer_idx → same seed → same order.
        assert_eq!(order1, order2);
    }

    #[test]
    fn test_growth_random_different_idx_different_order() {
        let (topo, _total) = topo_1n();
        let specs = vec![
            test_layer_spec("L0", LayerGrowthAlgo::Random),
            test_layer_spec("L1", LayerGrowthAlgo::Random),
        ];
        let o0 = get_core_order(&topo, &specs, 0);
        let o1 = get_core_order(&topo, &specs, 1);
        // Different seeds should (very likely) produce different orders.
        assert_ne!(o0, o1);
    }

    #[test]
    fn test_growth_random_2n_node_contiguous() {
        let (topo, _total) = topo_2n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::Random)];
        let order = get_core_order(&topo, &specs, 0);
        assert_valid_core_order(&order, 16);
        // Cores must be grouped by node: one contiguous block of node-0
        // cores (0-7) and one of node-1 cores (8-15), in either order.
        let first_node: Vec<usize> = order.iter().take(8).copied().collect();
        let second_node: Vec<usize> = order.iter().skip(8).copied().collect();
        assert!(
            first_node.iter().all(|&c| c < 8) || first_node.iter().all(|&c| c >= 8),
            "first 8 cores should all belong to the same node: {:?}",
            order
        );
        assert!(
            second_node.iter().all(|&c| c < 8) || second_node.iter().all(|&c| c >= 8),
            "last 8 cores should all belong to the same node: {:?}",
            order
        );
    }

    #[test]
    fn test_growth_random_4n_node_contiguous() {
        let (topo, _total) = topo_4n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::Random)];
        let order = get_core_order(&topo, &specs, 0);
        assert_valid_core_order(&order, 16);
        // 4 blocks of 4 cores, each block should be same node.
        // Node 0: cores 0-3, Node 1: cores 4-7, Node 2: cores 8-11, Node 3: cores 12-15
        for block_start in (0..16).step_by(4) {
            let block = &order[block_start..block_start + 4];
            let node = block[0] / 4;
            assert!(
                block.iter().all(|&c| c / 4 == node),
                "block at {}-{} should all be node {}, got {:?}",
                block_start,
                block_start + 3,
                node,
                block
            );
        }
    }

    // --- Topo ---

    #[test]
    fn test_growth_topo_no_pref_falls_back_to_roundrobin() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::Topo)];
        let order = get_core_order(&topo, &specs, 0);
        // No nodes/llcs specified → falls back to RoundRobin.
        let rr_specs = vec![test_layer_spec("L0", LayerGrowthAlgo::RoundRobin)];
        let rr_order = get_core_order(&topo, &rr_specs, 0);
        assert_eq!(order, rr_order);
    }

    #[test]
    fn test_growth_topo_with_llc_pref() {
        let (topo, _total) = topo_2n();
        // Prefer LLC 2 (on node 1). Per-node ordering: node 0 cores come
        // first (no LLC pref there), then node 1 with LLC 2 cores first.
        let specs = vec![test_layer_spec_with_llcs(
            "L0",
            LayerGrowthAlgo::Topo,
            vec![2],
        )];
        let order = get_core_order(&topo, &specs, 0);
        assert_valid_core_order(&order, 16);
        // Node 1 starts at position 8. Within node 1, LLC 2 (cores 8-11)
        // should come before LLC 3 (cores 12-15).
        let n1_order: Vec<usize> = order.iter().copied().filter(|&c| c >= 8).collect();
        for &core in &n1_order[..4] {
            assert!(
                core >= 8 && core < 12,
                "LLC2 core {} should be first within node 1",
                core
            );
        }
    }

    #[test]
    fn test_growth_topo_with_node_pref() {
        let (topo, _total) = topo_2n();
        // Prefer node 1 — node_order puts node 1 first.
        let specs = vec![test_layer_spec_with_nodes(
            "L0",
            LayerGrowthAlgo::Topo,
            vec![1],
        )];
        let order = get_core_order(&topo, &specs, 0);
        // Node 1 cores (8-15) should appear first.
        for &core in &order[..8] {
            assert!(
                core >= 8 && core < 16,
                "core {} should be node 1 (8-15)",
                core
            );
        }
    }

    // --- RoundRobin ---

    #[test]
    fn test_growth_roundrobin_1n() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::RoundRobin)];
        let order = get_core_order(&topo, &specs, 0);
        // Seeded interleave across LLCs: alternates LLC0/LLC1 cores.
        assert_eq!(order, vec![2, 6, 3, 4, 1, 7, 0, 5]);
    }

    #[test]
    fn test_growth_roundrobin_2n_per_node() {
        let (topo, _total) = topo_2n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::RoundRobin)];
        let order = get_core_order(&topo, &specs, 0);
        // Per-node LLC interleave: N0 LLCs interleaved, then N1.
        // Cross-node spreading handled by even-split budget.
        assert_eq!(
            order,
            vec![2, 6, 3, 4, 1, 7, 0, 5, 12, 8, 15, 11, 14, 10, 13, 9]
        );
    }

    // --- BigLittle / LittleBig ---

    #[test]
    fn test_growth_big_little_1n() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::BigLittle)];
        let order = get_core_order(&topo, &specs, 0);
        // All cores are same type in test topo → topo order.
        assert_eq!(order, vec![0, 1, 2, 3, 4, 5, 6, 7]);
    }

    #[test]
    fn test_growth_little_big_1n() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::LittleBig)];
        let order = get_core_order(&topo, &specs, 0);
        // All cores same type in test topo → stable sort preserves topo order
        // regardless of sort direction.
        assert_eq!(order, vec![0, 1, 2, 3, 4, 5, 6, 7]);
    }

    // --- Spread algo degeneration ---
    //
    // Spread algos produce the same core_order as their per-node equivalents.
    // Cross-node distribution is handled by even-split budget, not core_order.

    #[test]
    fn test_spread_algos_match_per_node_equivalents() {
        let (topo, _total) = topo_2n();
        let pairs: Vec<(LayerGrowthAlgo, LayerGrowthAlgo)> = vec![
            (LayerGrowthAlgo::NodeSpread, LayerGrowthAlgo::Linear),
            (LayerGrowthAlgo::NodeSpreadReverse, LayerGrowthAlgo::Reverse),
            (LayerGrowthAlgo::NodeSpreadRandom, LayerGrowthAlgo::Random),
            (LayerGrowthAlgo::RoundRobin, LayerGrowthAlgo::RoundRobin),
        ];
        for (spread, base) in &pairs {
            let s_specs = vec![test_layer_spec("L0", spread.clone())];
            let b_specs = vec![test_layer_spec("L0", base.clone())];
            let s_order = get_core_order(&topo, &s_specs, 0);
            let b_order = get_core_order(&topo, &b_specs, 0);
            assert_eq!(
                s_order, b_order,
                "{:?} should produce same order as {:?}",
                spread, base
            );
        }
    }

    #[test]
    fn test_spread_algos_have_layer_offset() {
        // Two layers with the same spread algo should get different starting
        // positions (rotate_node_layer_offset), just like non-spread algos.
        let (topo, _total) = topo_2n();
        let specs = vec![
            test_layer_spec("L0", LayerGrowthAlgo::NodeSpread),
            test_layer_spec("L1", LayerGrowthAlgo::NodeSpread),
        ];
        let o0 = get_core_order(&topo, &specs, 0);
        let o1 = get_core_order(&topo, &specs, 1);
        // Same cores, different starting position.
        assert_ne!(o0, o1, "different layers should get different offsets");
        assert_valid_core_order(&o0, 16);
        assert_valid_core_order(&o1, 16);
    }

    // --- RandomTopo ---

    #[test]
    fn test_growth_random_topo_1n() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::RandomTopo)];
        let order = get_core_order(&topo, &specs, 0);
        // Randomized within topology levels — shuffled permutation.
        assert_valid_core_order(&order, 8);
        assert_ne!(order, vec![0, 1, 2, 3, 4, 5, 6, 7], "should be shuffled");
    }

    #[test]
    fn test_growth_random_topo_2n() {
        let (topo, _total) = topo_2n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::RandomTopo)];
        let order = get_core_order(&topo, &specs, 0);
        assert_valid_core_order(&order, 16);
        assert_ne!(
            order,
            vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            "should be shuffled"
        );
    }

    // --- StickyDynamic ---

    #[test]
    fn test_growth_sticky_dynamic_same_as_sticky() {
        let (topo, _total) = topo_1n();
        let specs_s = vec![test_layer_spec("L0", LayerGrowthAlgo::Sticky)];
        let specs_sd = vec![test_layer_spec("L0", LayerGrowthAlgo::StickyDynamic)];
        let order_s = get_core_order(&topo, &specs_s, 0);
        let order_sd = get_core_order(&topo, &specs_sd, 0);
        assert_eq!(
            order_s, order_sd,
            "StickyDynamic initial order should match Sticky"
        );
    }

    #[test]
    fn test_growth_sticky_dynamic_2n() {
        let (topo, _total) = topo_2n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::StickyDynamic)];
        let order = get_core_order(&topo, &specs, 0);
        // Initial order matches Sticky (verified in test_growth_sticky_dynamic_same_as_sticky).
        assert_eq!(
            order,
            vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
        );
    }

    // --- CpuSetSpread (host-dependent, verify basic invariants) ---

    #[test]
    fn test_growth_cpuset_spread_1n() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::CpuSetSpread)];
        let pool = CpuPool::new(topo.clone(), false).unwrap();
        let orders = LayerGrowthAlgo::layer_core_orders(&pool, &specs, &topo).unwrap();
        let order = &orders[&0];
        // CpuSetSpread reads /sys/fs/cgroup, results are host-dependent.
        // Just verify no duplicates and all cores are valid.
        let mut seen = std::collections::HashSet::new();
        for &core in order.iter().flatten() {
            assert!(core < 8, "core {} out of range", core);
            assert!(seen.insert(core), "duplicate core {}", core);
        }
    }

    #[test]
    fn test_growth_cpuset_spread_reverse_1n() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::CpuSetSpreadReverse)];
        let pool = CpuPool::new(topo.clone(), false).unwrap();
        let orders = LayerGrowthAlgo::layer_core_orders(&pool, &specs, &topo).unwrap();
        let order = &orders[&0];
        let mut seen = std::collections::HashSet::new();
        for &core in order.iter().flatten() {
            assert!(core < 8, "core {} out of range", core);
            assert!(seen.insert(core), "duplicate core {}", core);
        }
    }

    #[test]
    fn test_growth_cpuset_spread_random_1n() {
        let (topo, _total) = topo_1n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::CpuSetSpreadRandom)];
        let pool = CpuPool::new(topo.clone(), false).unwrap();
        let orders = LayerGrowthAlgo::layer_core_orders(&pool, &specs, &topo).unwrap();
        let order = &orders[&0];
        let mut seen = std::collections::HashSet::new();
        for &core in order.iter().flatten() {
            assert!(core < 8, "core {} out of range", core);
            assert!(seen.insert(core), "duplicate core {}", core);
        }
    }

    // --- All algorithms produce valid orders on 2N ---

    #[test]
    fn test_alloc_free_realloc_all_deterministic() {
        let deterministic_algos = vec![
            LayerGrowthAlgo::Sticky,
            LayerGrowthAlgo::Linear,
            LayerGrowthAlgo::Reverse,
            LayerGrowthAlgo::RoundRobin,
            LayerGrowthAlgo::BigLittle,
            LayerGrowthAlgo::LittleBig,
            LayerGrowthAlgo::NodeSpread,
            LayerGrowthAlgo::NodeSpreadReverse,
            LayerGrowthAlgo::StickyDynamic,
        ];
        for algo in deterministic_algos {
            let (topo, total) = topo_1n();
            let allowed = all_cpus_mask(total);
            let specs = vec![test_layer_spec("L0", algo.clone())];
            let order = get_core_order(&topo, &specs, 0);
            let mut pool = CpuPool::new(topo, false).unwrap();
            let initial = pool.available_cpus().weight();

            // Alloc 8, alloc 4, free first 8, re-alloc 4 → from freed set.
            let alloc1 = pool.alloc_cpus(&allowed, &order, 8).unwrap();
            let _alloc2 = pool.alloc_cpus(&allowed, &order, 4).unwrap();
            pool.free(&alloc1).unwrap();
            assert_eq!(
                pool.available_cpus().weight(),
                initial - _alloc2.weight(),
                "{:?}: wrong available count after free",
                algo
            );
            let alloc3 = pool.alloc_cpus(&allowed, &order, 4).unwrap();
            for cpu in alloc3.iter() {
                assert!(
                    alloc1.test_cpu(cpu),
                    "{:?}: cpu {} should come from freed alloc1",
                    algo,
                    cpu
                );
            }
        }
    }

    #[test]
    fn test_all_algos_valid_on_2n() {
        let (topo, _total) = topo_2n();
        let algos = vec![
            LayerGrowthAlgo::Sticky,
            LayerGrowthAlgo::Linear,
            LayerGrowthAlgo::Reverse,
            LayerGrowthAlgo::Random,
            LayerGrowthAlgo::Topo,
            LayerGrowthAlgo::RoundRobin,
            LayerGrowthAlgo::BigLittle,
            LayerGrowthAlgo::LittleBig,
            LayerGrowthAlgo::NodeSpread,
            LayerGrowthAlgo::NodeSpreadReverse,
            LayerGrowthAlgo::NodeSpreadRandom,
            LayerGrowthAlgo::RandomTopo,
            LayerGrowthAlgo::StickyDynamic,
        ];
        for algo in algos {
            let specs = vec![test_layer_spec("L0", algo.clone())];
            let pool = CpuPool::new(topo.clone(), false).unwrap();
            let orders = LayerGrowthAlgo::layer_core_orders(&pool, &specs, &topo)
                .unwrap_or_else(|e| panic!("{:?} failed: {}", algo, e));
            let order: Vec<usize> = orders[&0].iter().flatten().copied().collect();
            assert_valid_core_order(&order, 16);
        }
    }

    // --- 4N growth algorithm tests ---

    #[test]
    fn test_all_algos_valid_on_4n() {
        let (topo, _total) = topo_4n();
        let algos = vec![
            LayerGrowthAlgo::Sticky,
            LayerGrowthAlgo::Linear,
            LayerGrowthAlgo::Reverse,
            LayerGrowthAlgo::Random,
            LayerGrowthAlgo::Topo,
            LayerGrowthAlgo::RoundRobin,
            LayerGrowthAlgo::BigLittle,
            LayerGrowthAlgo::LittleBig,
            LayerGrowthAlgo::NodeSpread,
            LayerGrowthAlgo::NodeSpreadReverse,
            LayerGrowthAlgo::NodeSpreadRandom,
            LayerGrowthAlgo::RandomTopo,
            LayerGrowthAlgo::StickyDynamic,
        ];
        for algo in algos {
            let specs = vec![test_layer_spec("L0", algo.clone())];
            let pool = CpuPool::new(topo.clone(), false).unwrap();
            let orders = LayerGrowthAlgo::layer_core_orders(&pool, &specs, &topo)
                .unwrap_or_else(|e| panic!("{:?} failed: {}", algo, e));
            let order: Vec<usize> = orders[&0].iter().flatten().copied().collect();
            assert_valid_core_order(&order, 16);
        }
    }

    #[test]
    fn test_growth_sticky_4n() {
        let (topo, _total) = topo_4n();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::Sticky)];
        let order = get_core_order(&topo, &specs, 0);
        // Single layer, idx=0: identity order across all 16 cores.
        assert_eq!(
            order,
            vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
        );
    }

    #[test]
    fn test_growth_sticky_4n_multi_layer() {
        let (topo, _total) = topo_4n();
        let specs = vec![
            test_layer_spec("L0", LayerGrowthAlgo::Sticky),
            test_layer_spec("L1", LayerGrowthAlgo::Sticky),
            test_layer_spec("L2", LayerGrowthAlgo::Sticky),
            test_layer_spec("L3", LayerGrowthAlgo::Sticky),
        ];
        let o0 = get_core_order(&topo, &specs, 0);
        let o1 = get_core_order(&topo, &specs, 1);
        let o2 = get_core_order(&topo, &specs, 2);
        let o3 = get_core_order(&topo, &specs, 3);
        // Each layer should start at a different core.
        let starts: Vec<usize> = vec![o0[0], o1[0], o2[0], o3[0]];
        let mut unique = starts.clone();
        unique.sort();
        unique.dedup();
        assert_eq!(
            unique.len(),
            4,
            "4 layers should have 4 different starting cores, got {:?}",
            starts
        );
        // All should be valid permutations.
        for order in [&o0, &o1, &o2, &o3] {
            assert_valid_core_order(order, 16);
        }
    }

    // =========================================================================
    // Non-StickyDynamic per-node allocation tests
    //
    // These tests simulate the per-node shrink/grow logic from
    // refresh_cpumasks() for non-StickyDynamic layers. The simulation
    // mirrors the two-phase algorithm:
    //   1. Shrink per-node: free excess CPUs using reversed core_order
    //   2. Grow per-node: allocate CPUs using forward core_order in norder
    // =========================================================================

    /// Lightweight stand-in for a non-SD layer's state.
    struct NonSdLayerState {
        #[allow(dead_code)]
        name: String,
        cpus: Cpumask,
        nr_cpus: usize,
        nr_node_cpus: Vec<usize>,
        core_order: Vec<Vec<usize>>,
        allowed_cpus: Cpumask,
    }

    impl NonSdLayerState {
        fn new(
            name: &str,
            nr_nodes: usize,
            total_cpus: usize,
            core_order: Vec<Vec<usize>>,
        ) -> Self {
            let mut allowed = Cpumask::new();
            for cpu in 0..total_cpus {
                allowed.set_cpu(cpu).unwrap();
            }
            Self {
                name: name.to_string(),
                cpus: Cpumask::new(),
                nr_cpus: 0,
                nr_node_cpus: vec![0; nr_nodes],
                core_order,
                allowed_cpus: allowed,
            }
        }
    }

    /// Helper to construct a LayerAlloc with given per-node pinned and unpinned.
    fn make_alloc(pinned: &[usize], unpinned: &[usize]) -> alloc::LayerAlloc {
        alloc::LayerAlloc {
            pinned: pinned.to_vec(),
            unpinned_budget: unpinned.iter().sum(),
            unpinned: unpinned.to_vec(),
        }
    }

    /// Simulate per-node shrink/grow from refresh_cpumasks for non-SD layers.
    ///
    /// `ascending` is (layer_idx, target_cpus) sorted ascending by target.
    /// `layer_allocs[idx]` is the per-node allocation for layer idx.
    /// `norders[idx]` is the node iteration order for growing.
    /// `au` is the alloc unit (cpus per core).
    fn simulate_non_sd_recompute(
        pool: &mut CpuPool,
        layers: &mut [NonSdLayerState],
        ascending: &[(usize, usize)],
        layer_allocs: &[alloc::LayerAlloc],
        norders: &[Vec<usize>],
        au: usize,
        topo: &Topology,
    ) {
        let nr_nodes = topo.nodes.len();

        // Shrink per-node: free excess CPUs from each node (reverse iteration).
        for &(idx, _target) in ascending.iter().rev() {
            let layer = &mut layers[idx];
            let alloc = &layer_allocs[idx];

            for n in 0..nr_nodes {
                let desired = alloc.node_target(n) * au;
                let mut to_free = layer.nr_node_cpus[n].saturating_sub(desired);
                let node_span = &topo.nodes[&n].span;

                while to_free > 0 {
                    let node_cands = layer.cpus.and(node_span);
                    let cpus_to_free =
                        match pool.next_to_free(&node_cands, layer.core_order[n].iter().rev()) {
                            Ok(Some(ret)) => ret,
                            _ => break,
                        };
                    let nr = cpus_to_free.weight();
                    layer.cpus &= &cpus_to_free.not();
                    layer.nr_cpus -= nr;
                    for cpu in cpus_to_free.iter() {
                        let node_id = topo.all_cpus[&cpu].node_id;
                        layer.nr_node_cpus[node_id] -= 1;
                    }
                    pool.free(&cpus_to_free).unwrap();
                    to_free = to_free.saturating_sub(nr);
                }
            }
        }

        // Grow per-node: allocate CPUs in norder.
        for &(idx, _target) in ascending.iter() {
            let layer = &mut layers[idx];
            let alloc = &layer_allocs[idx];
            let norder = &norders[idx];

            for &node_id in norder.iter() {
                let node_target = alloc.node_target(node_id) * au;
                let cur_node = layer.nr_node_cpus[node_id];
                if node_target <= cur_node {
                    continue;
                }
                let mut nr_to_alloc = node_target - cur_node;
                let node_span = &topo.nodes[&node_id].span;
                let node_allowed = layer.allowed_cpus.and(node_span);

                while nr_to_alloc > 0 {
                    let nr_alloced = match pool.alloc_cpus(
                        &node_allowed,
                        &layer.core_order[node_id],
                        nr_to_alloc,
                    ) {
                        Some(new_cpus) => {
                            let nr = new_cpus.weight();
                            layer.cpus |= &new_cpus;
                            layer.nr_cpus += nr;
                            for cpu in new_cpus.iter() {
                                let nid = topo.all_cpus[&cpu].node_id;
                                layer.nr_node_cpus[nid] += 1;
                            }
                            nr
                        }
                        None => 0,
                    };
                    if nr_alloced == 0 {
                        break;
                    }
                    nr_to_alloc -= nr_alloced.min(nr_to_alloc);
                }
            }
        }
    }

    /// Assert CPU conservation: pool available + all layers' CPUs = total.
    fn assert_cpu_conservation(pool: &CpuPool, layers: &[NonSdLayerState], total_cpus: usize) {
        let pool_avail = pool.available_cpus().weight();
        let assigned: usize = layers.iter().map(|l| l.nr_cpus).sum();
        assert_eq!(
            pool_avail + assigned,
            total_cpus,
            "CPU conservation violated: {} available + {} assigned != {}",
            pool_avail,
            assigned,
            total_cpus
        );
    }

    /// Assert no CPU appears in multiple layers.
    fn assert_no_cpu_overlap(layers: &[NonSdLayerState], step: &str) {
        for i in 0..layers.len() {
            for j in (i + 1)..layers.len() {
                let overlap = layers[i].cpus.and(&layers[j].cpus);
                assert_eq!(
                    overlap.weight(),
                    0,
                    "{}: L{} and L{} overlap by {} CPUs",
                    step,
                    i,
                    j,
                    overlap.weight()
                );
            }
        }
    }

    /// Assert per-node CPU counts match expectations.
    fn assert_node_cpus(layer: &NonSdLayerState, expected: &[usize], label: &str) {
        for (n, &exp) in expected.iter().enumerate() {
            assert_eq!(
                layer.nr_node_cpus[n], exp,
                "{}: node {} expected {} cpus, got {}",
                label, n, exp, layer.nr_node_cpus[n]
            );
        }
    }

    // --- Non-SD per-node: 1N tests ---

    #[test]
    fn test_nonsd_1n_grow_partial() {
        // 1N: 16 CPUs, au=2. Grow to 4 CPUs (2 cores).
        let (topo, total) = topo_1n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![NonSdLayerState::new("L0", 1, total, core_order)];
        let allocs = vec![make_alloc(&[0], &[2])]; // 2 au unpinned on N0 = 4 CPUs
        let norders = vec![vec![0]];
        let ascending = vec![(0, 4)];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &ascending,
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 4);
        assert_node_cpus(&layers[0], &[4], "grow partial");
        assert_cpu_conservation(&pool, &layers, total);
    }

    #[test]
    fn test_nonsd_1n_grow_full() {
        // 1N: grow to all 16 CPUs.
        let (topo, total) = topo_1n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![NonSdLayerState::new("L0", 1, total, core_order)];
        let allocs = vec![make_alloc(&[0], &[8])]; // 8 au = 16 CPUs
        let norders = vec![vec![0]];
        let ascending = vec![(0, 16)];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &ascending,
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 16);
        assert_cpu_conservation(&pool, &layers, total);
    }

    #[test]
    fn test_nonsd_1n_grow_then_shrink() {
        let (topo, total) = topo_1n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![NonSdLayerState::new("L0", 1, total, core_order)];
        let norders = vec![vec![0]];

        // Grow to 12 CPUs (6 cores = 6 au).
        let allocs = vec![make_alloc(&[0], &[6])];
        let ascending = vec![(0, 12)];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &ascending,
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 12);
        assert_cpu_conservation(&pool, &layers, total);

        // Shrink to 4 CPUs (2 au).
        let allocs = vec![make_alloc(&[0], &[2])];
        let ascending = vec![(0, 4)];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &ascending,
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 4);
        assert_cpu_conservation(&pool, &layers, total);
    }

    #[test]
    fn test_nonsd_1n_roundtrip() {
        let (topo, total) = topo_1n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![NonSdLayerState::new("L0", 1, total, core_order)];
        let norders = vec![vec![0]];

        // Grow to full.
        let allocs = vec![make_alloc(&[0], &[8])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 16)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 16);

        // Shrink to zero.
        let allocs = vec![make_alloc(&[0], &[0])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 0)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 0);
        assert_eq!(pool.available_cpus().weight(), total);
    }

    // --- Non-SD per-node: 2N tests ---

    #[test]
    fn test_nonsd_2n_all_on_node0() {
        // 2N: 32 CPUs, au=2. All allocation on node 0 (16 CPUs max).
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![NonSdLayerState::new("L0", 2, total, core_order)];
        let allocs = vec![make_alloc(&[0, 0], &[8, 0])]; // 8 au on N0 = 16 CPUs
        let norders = vec![vec![0, 1]];
        let ascending = vec![(0, 16)];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &ascending,
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 16);
        assert_node_cpus(&layers[0], &[16, 0], "all on N0");
        assert_cpu_conservation(&pool, &layers, total);
    }

    #[test]
    fn test_nonsd_2n_split_across_nodes() {
        // 2N: 8 CPUs per node (4 au per node).
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![NonSdLayerState::new("L0", 2, total, core_order)];
        let allocs = vec![make_alloc(&[0, 0], &[4, 4])]; // 4 au each = 8 CPUs each
        let norders = vec![vec![0, 1]];
        let ascending = vec![(0, 16)];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &ascending,
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 16);
        assert_node_cpus(&layers[0], &[8, 8], "split");
        assert_cpu_conservation(&pool, &layers, total);
    }

    #[test]
    fn test_nonsd_2n_asymmetric() {
        // 2N: 12 CPUs on N0, 4 on N1.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![NonSdLayerState::new("L0", 2, total, core_order)];
        let allocs = vec![make_alloc(&[0, 0], &[6, 2])]; // 12 on N0, 4 on N1
        let norders = vec![vec![0, 1]];
        let ascending = vec![(0, 16)];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &ascending,
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 16);
        assert_node_cpus(&layers[0], &[12, 4], "asymmetric");
        assert_cpu_conservation(&pool, &layers, total);
    }

    #[test]
    fn test_nonsd_2n_grow_one_shrink_other() {
        // Two-step: first all on N0, then move some to N1.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![NonSdLayerState::new("L0", 2, total, core_order)];
        let norders = vec![vec![0, 1]];

        // Step 1: All 16 CPUs on N0.
        let allocs = vec![make_alloc(&[0, 0], &[8, 0])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 16)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_node_cpus(&layers[0], &[16, 0], "step 1");

        // Step 2: Rebalance to 8 per node.
        let allocs = vec![make_alloc(&[0, 0], &[4, 4])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 16)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_node_cpus(&layers[0], &[8, 8], "step 2");
        assert_cpu_conservation(&pool, &layers, total);
    }

    #[test]
    fn test_nonsd_2n_shrink_releases_correct_node() {
        // Start with 8+8, shrink to 8+0. Node 1 should be freed.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![NonSdLayerState::new("L0", 2, total, core_order)];
        let norders = vec![vec![0, 1]];

        // Grow: 8 per node.
        let allocs = vec![make_alloc(&[0, 0], &[4, 4])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 16)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_node_cpus(&layers[0], &[8, 8], "grow");

        // Shrink: keep N0, release N1.
        let allocs = vec![make_alloc(&[0, 0], &[4, 0])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 8)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 8);
        assert_node_cpus(&layers[0], &[8, 0], "shrink");
        assert_cpu_conservation(&pool, &layers, total);
    }

    // --- Non-SD per-node: 2N multi-layer tests ---

    #[test]
    fn test_nonsd_2n_two_layers_same_node() {
        // Two layers both want N0 CPUs.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![
            NonSdLayerState::new("L0", 2, total, core_order.clone()),
            NonSdLayerState::new("L1", 2, total, core_order),
        ];
        let allocs = vec![
            make_alloc(&[0, 0], &[4, 0]), // L0: 8 CPUs on N0
            make_alloc(&[0, 0], &[4, 0]), // L1: 8 CPUs on N0
        ];
        let norders = vec![vec![0, 1], vec![0, 1]];
        // Ascending by target: both at 8.
        let ascending = vec![(0, 8), (1, 8)];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &ascending,
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 8);
        assert_eq!(layers[1].nr_cpus, 8);
        assert_node_cpus(&layers[0], &[8, 0], "L0");
        assert_node_cpus(&layers[1], &[8, 0], "L1");
        assert_no_cpu_overlap(&layers, "same node");
        assert_cpu_conservation(&pool, &layers, total);
    }

    #[test]
    fn test_nonsd_2n_two_layers_different_nodes() {
        // L0 on N0, L1 on N1.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![
            NonSdLayerState::new("L0", 2, total, core_order.clone()),
            NonSdLayerState::new("L1", 2, total, core_order),
        ];
        let allocs = vec![
            make_alloc(&[0, 0], &[4, 0]), // L0: 8 CPUs on N0
            make_alloc(&[0, 0], &[0, 4]), // L1: 8 CPUs on N1
        ];
        let norders = vec![vec![0, 1], vec![1, 0]];
        let ascending = vec![(0, 8), (1, 8)];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &ascending,
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 8);
        assert_eq!(layers[1].nr_cpus, 8);
        assert_node_cpus(&layers[0], &[8, 0], "L0");
        assert_node_cpus(&layers[1], &[0, 8], "L1");
        assert_no_cpu_overlap(&layers, "diff nodes");
        assert_cpu_conservation(&pool, &layers, total);
    }

    #[test]
    fn test_nonsd_2n_one_shrinks_other_grows() {
        // L0 starts with 16 on N0, then shrinks; L1 grows into freed space.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![
            NonSdLayerState::new("L0", 2, total, core_order.clone()),
            NonSdLayerState::new("L1", 2, total, core_order),
        ];
        let norders = vec![vec![0, 1], vec![0, 1]];

        // Step 1: L0 gets 16 on N0.
        let allocs = vec![make_alloc(&[0, 0], &[8, 0]), make_alloc(&[0, 0], &[0, 0])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(1, 0), (0, 16)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 16);
        assert_eq!(layers[1].nr_cpus, 0);

        // Step 2: L0 shrinks to 8 on N0, L1 grows to 8 on N0.
        let allocs = vec![make_alloc(&[0, 0], &[4, 0]), make_alloc(&[0, 0], &[4, 0])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 8), (1, 8)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 8);
        assert_eq!(layers[1].nr_cpus, 8);
        assert_node_cpus(&layers[0], &[8, 0], "L0 after");
        assert_node_cpus(&layers[1], &[8, 0], "L1 after");
        assert_no_cpu_overlap(&layers, "swap");
        assert_cpu_conservation(&pool, &layers, total);
    }

    // --- Non-SD per-node: stress tests ---

    #[test]
    fn test_nonsd_2n_10_cycle_oscillation() {
        // Single layer oscillates between N0-heavy and N1-heavy, 10 cycles.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![NonSdLayerState::new("L0", 2, total, core_order)];
        let norders = vec![vec![0, 1]];

        for cycle in 0..10 {
            // Phase A: 12 on N0, 4 on N1.
            let allocs = vec![make_alloc(&[0, 0], &[6, 2])];
            simulate_non_sd_recompute(
                &mut pool,
                &mut layers,
                &[(0, 16)],
                &allocs,
                &norders,
                au,
                &topo,
            );
            assert_node_cpus(&layers[0], &[12, 4], &format!("cycle {} A", cycle));
            assert_cpu_conservation(&pool, &layers, total);

            // Phase B: 4 on N0, 12 on N1.
            let allocs = vec![make_alloc(&[0, 0], &[2, 6])];
            simulate_non_sd_recompute(
                &mut pool,
                &mut layers,
                &[(0, 16)],
                &allocs,
                &norders,
                au,
                &topo,
            );
            assert_node_cpus(&layers[0], &[4, 12], &format!("cycle {} B", cycle));
            assert_cpu_conservation(&pool, &layers, total);
        }
    }

    #[test]
    fn test_nonsd_2n_multi_layer_competition() {
        // 3 layers compete on 2 nodes over 5 rebalancing steps.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![
            NonSdLayerState::new("L0", 2, total, core_order.clone()),
            NonSdLayerState::new("L1", 2, total, core_order.clone()),
            NonSdLayerState::new("L2", 2, total, core_order),
        ];
        let norders = vec![vec![0, 1], vec![0, 1], vec![1, 0]];

        struct Step {
            allocs: Vec<alloc::LayerAlloc>,
            ascending: Vec<(usize, usize)>,
        }

        let steps = vec![
            // Step 0: L0=8(N0), L1=8(N0), L2=16(N1).
            Step {
                allocs: vec![
                    make_alloc(&[0, 0], &[4, 0]),
                    make_alloc(&[0, 0], &[4, 0]),
                    make_alloc(&[0, 0], &[0, 8]),
                ],
                ascending: vec![(0, 8), (1, 8), (2, 16)],
            },
            // Step 1: Rebalance — L0=4(N0), L1=12(N0+N1), L2=8(N1).
            Step {
                allocs: vec![
                    make_alloc(&[0, 0], &[2, 0]),
                    make_alloc(&[0, 0], &[4, 2]),
                    make_alloc(&[0, 0], &[0, 4]),
                ],
                ascending: vec![(0, 4), (2, 8), (1, 12)],
            },
            // Step 2: L2 grows, others shrink.
            Step {
                allocs: vec![
                    make_alloc(&[0, 0], &[1, 0]),
                    make_alloc(&[0, 0], &[1, 0]),
                    make_alloc(&[0, 0], &[4, 8]),
                ],
                ascending: vec![(0, 2), (1, 2), (2, 24)],
            },
            // Step 3: Even split.
            Step {
                allocs: vec![
                    make_alloc(&[0, 0], &[3, 2]),
                    make_alloc(&[0, 0], &[2, 3]),
                    make_alloc(&[0, 0], &[3, 3]),
                ],
                ascending: vec![(0, 10), (1, 10), (2, 12)],
            },
            // Step 4: All shrink to zero.
            Step {
                allocs: vec![
                    make_alloc(&[0, 0], &[0, 0]),
                    make_alloc(&[0, 0], &[0, 0]),
                    make_alloc(&[0, 0], &[0, 0]),
                ],
                ascending: vec![(0, 0), (1, 0), (2, 0)],
            },
        ];

        for (i, step) in steps.iter().enumerate() {
            simulate_non_sd_recompute(
                &mut pool,
                &mut layers,
                &step.ascending,
                &step.allocs,
                &norders,
                au,
                &topo,
            );
            assert_cpu_conservation(&pool, &layers, total);
            assert_no_cpu_overlap(&layers, &format!("step {}", i));
            // Verify per-node counts match alloc targets.
            for (idx, alloc) in step.allocs.iter().enumerate() {
                for n in 0..2 {
                    assert_eq!(
                        layers[idx].nr_node_cpus[n],
                        alloc.node_target(n) * au,
                        "step {} L{} node {} mismatch",
                        i,
                        idx,
                        n
                    );
                }
            }
        }
    }

    #[test]
    fn test_nonsd_2n_staircase() {
        // Staircase: grow across nodes then shrink back. Verify per-node at each step.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let core_order = core_order_per_node(&topo);
        let mut layers = vec![NonSdLayerState::new("L0", 2, total, core_order)];
        let norders = vec![vec![0, 1]];

        // (N0_unpinned_au, N1_unpinned_au) at each step.
        let steps: Vec<(usize, usize)> = vec![
            (2, 0), // 4 CPUs on N0
            (4, 0), // 8 CPUs on N0
            (4, 2), // 8 N0 + 4 N1
            (4, 4), // 8 N0 + 8 N1
            (8, 4), // 16 N0 + 8 N1
            (8, 8), // 16 N0 + 16 N1 = full
            (8, 4), // shrink N1
            (4, 4), // shrink N0
            (4, 0), // shrink N1 to 0
            (0, 0), // all freed
        ];

        for (i, &(n0_au, n1_au)) in steps.iter().enumerate() {
            let total_cpus = (n0_au + n1_au) * au;
            let allocs = vec![make_alloc(&[0, 0], &[n0_au, n1_au])];
            let ascending = vec![(0, total_cpus)];
            simulate_non_sd_recompute(
                &mut pool,
                &mut layers,
                &ascending,
                &allocs,
                &norders,
                au,
                &topo,
            );
            assert_node_cpus(
                &layers[0],
                &[n0_au * au, n1_au * au],
                &format!("step {}", i),
            );
            assert_cpu_conservation(&pool, &layers, total);
        }
    }

    // =========================================================================
    // StickyDynamic lifecycle tests
    //
    // The StickyDynamic algorithm in recompute_layer_core_order() is tightly
    // coupled to the Scheduler struct. These tests simulate the three-phase
    // algorithm (free → redistribute → spillover) using lightweight state
    // to verify correctness independent of the full runtime.
    // =========================================================================

    /// Mirrors the target computation from Scheduler::compute_target_llcs.
    fn compute_target_llcs(target: usize, node_llcs: &BTreeMap<usize, Arc<Llc>>) -> (usize, usize) {
        let mut remaining = target;
        let mut full = 0;
        for llc in node_llcs.values() {
            let cpus_in_llc = llc.span.weight();
            if remaining >= cpus_in_llc {
                full += 1;
                remaining -= cpus_in_llc;
            } else {
                let mut extra = 0;
                for (_, core) in llc.cores.iter() {
                    if remaining == 0 {
                        break;
                    }
                    extra += 1;
                    remaining = remaining.saturating_sub(core.cpus.len());
                }
                return (full, extra);
            }
        }
        (full, 0)
    }

    // --- StickyDynamic per-node test infrastructure ---

    /// Lightweight stand-in for Layer's per-node StickyDynamic state.
    struct SdLayerState {
        #[allow(dead_code)]
        name: String,
        /// Per-node assigned LLCs: assigned_llcs[n] = list of LLC IDs on node n.
        assigned_llcs: Vec<Vec<usize>>,
    }

    impl SdLayerState {
        fn new(name: &str, nr_nodes: usize) -> Self {
            Self {
                name: name.to_string(),
                assigned_llcs: vec![vec![]; nr_nodes],
            }
        }

        fn total_llcs(&self) -> usize {
            self.assigned_llcs.iter().map(|v| v.len()).sum()
        }

        fn all_llcs(&self) -> Vec<usize> {
            self.assigned_llcs.iter().flatten().copied().collect()
        }

        fn llcs_on_node(&self, n: usize) -> usize {
            self.assigned_llcs[n].len()
        }
    }

    fn make_sd_alloc(pinned: &[usize], unpinned: &[usize]) -> alloc::LayerAlloc {
        alloc::LayerAlloc {
            pinned: pinned.to_vec(),
            unpinned_budget: unpinned.iter().sum(),
            unpinned: unpinned.to_vec(),
        }
    }

    /// Simulate the three-phase per-node StickyDynamic algorithm from
    /// Scheduler::recompute_layer_core_order(). Sort order is derived from
    /// layer_allocs (largest total acquires first), matching production code.
    fn simulate_sd_recompute(
        pool: &mut CpuPool,
        layers: &mut [SdLayerState],
        layer_allocs: &[alloc::LayerAlloc],
        au: usize,
        topo: &Topology,
    ) {
        let nr_nodes = topo.nodes.len();

        // Derive ascending sort order from allocs (matches main.rs).
        let mut ascending: Vec<(usize, usize)> = layer_allocs
            .iter()
            .enumerate()
            .map(|(i, a)| (i, a.total() * au))
            .collect();
        ascending.sort_by(|a, b| a.1.cmp(&b.1));

        // Phase 1 — Free per-node: return excess LLCs.
        for &(idx, _) in ascending.iter().rev() {
            let layer = &mut layers[idx];
            let alloc = &layer_allocs[idx];

            for n in 0..nr_nodes {
                let assigned_on_n = layer.assigned_llcs[n].len();
                let target_full_n = compute_target_llcs(
                    alloc.node_target(n) * au,
                    &topo.nodes.get(&n).unwrap().llcs,
                )
                .0;
                let mut to_free = assigned_on_n.saturating_sub(target_full_n);

                while to_free > 0 {
                    if let Some(llc) = layer.assigned_llcs[n].pop() {
                        pool.return_llc(llc);
                        to_free -= 1;
                    } else {
                        break;
                    }
                }
            }
        }

        // Phase 2 — Acquire per-node: claim LLCs from cpu_pool.
        for &(idx, _) in ascending.iter().rev() {
            let layer = &mut layers[idx];
            let alloc = &layer_allocs[idx];

            for n in 0..nr_nodes {
                let cur_on_n = layer.assigned_llcs[n].len();
                let target_full_n = compute_target_llcs(
                    alloc.node_target(n) * au,
                    &topo.nodes.get(&n).unwrap().llcs,
                )
                .0;
                let mut to_alloc = target_full_n.saturating_sub(cur_on_n);

                while to_alloc > 0 {
                    if let Some(llc) = pool.take_llc_from_node(n) {
                        layer.assigned_llcs[n].push(llc);
                        to_alloc -= 1;
                    } else {
                        break;
                    }
                }
            }
        }

        // Phase 3 — Spillover per-node: consume extra cores from free LLCs.
        for &(idx, _) in ascending.iter() {
            let alloc = &layer_allocs[idx];

            for n in 0..nr_nodes {
                let mut extra = compute_target_llcs(
                    alloc.node_target(n) * au,
                    &topo.nodes.get(&n).unwrap().llcs,
                )
                .1;

                if let Some(node_llcs) = pool.free_llcs.get_mut(&n) {
                    for entry in node_llcs.iter_mut() {
                        if extra == 0 {
                            break;
                        }
                        let llc_id = entry.0;
                        let llc = topo.all_llcs.get(&llc_id).unwrap();
                        let avail = llc.cores.len() - entry.1;
                        let used = extra.min(avail);
                        entry.1 += used;
                        extra -= used;
                    }
                }
            }
        }

        // Reset consumed entries.
        for node_llcs in pool.free_llcs.values_mut() {
            for entry in node_llcs.iter_mut() {
                entry.1 = 0;
            }
        }
    }

    fn assert_sd_llc_conservation(pool: &CpuPool, layers: &[SdLayerState], total_llcs: usize) {
        let layer_llcs: usize = layers.iter().map(|l| l.total_llcs()).sum();
        let free_llcs = pool.total_free_llcs();
        assert_eq!(
            layer_llcs + free_llcs,
            total_llcs,
            "LLC conservation: {} layer + {} free != {} total",
            layer_llcs,
            free_llcs,
            total_llcs
        );
    }

    fn assert_sd_no_duplicate_llcs(layers: &[SdLayerState], step: &str) {
        let mut seen = std::collections::HashSet::new();
        for layer in layers {
            for llc in layer.all_llcs() {
                assert!(seen.insert(llc), "duplicate LLC {} at step '{}'", llc, step);
            }
        }
    }

    // --- compute_target_llcs ---

    #[test]
    fn test_compute_target_llcs_1n() {
        let (topo, _total) = topo_1n();
        let node = topo.nodes.get(&0).unwrap();
        // 1N: 8 cores, 2 LLCs → 4 cores/LLC, 2 HTs/core → 8 cpus/LLC
        assert_eq!(compute_target_llcs(0, &node.llcs), (0, 0));
        assert_eq!(compute_target_llcs(1, &node.llcs), (0, 1)); // 1 cpu = 0 full + 1 extra core
        assert_eq!(compute_target_llcs(2, &node.llcs), (0, 1)); // 2 cpus = 1 core
        assert_eq!(compute_target_llcs(3, &node.llcs), (0, 2)); // 3 cpus = 2 cores (ceil)
        assert_eq!(compute_target_llcs(8, &node.llcs), (1, 0)); // exactly 1 LLC
        assert_eq!(compute_target_llcs(9, &node.llcs), (1, 1)); // 1 LLC + 1 extra core
        assert_eq!(compute_target_llcs(16, &node.llcs), (2, 0)); // exactly 2 LLCs
    }

    #[test]
    fn test_compute_target_llcs_2n() {
        let (topo, _total) = topo_2n();
        // 2N: 16 cores, 4 LLCs → 4 cores/LLC, 2 HTs/core → 8 cpus/LLC
        assert_eq!(compute_target_llcs(0, &topo.all_llcs), (0, 0));
        assert_eq!(compute_target_llcs(8, &topo.all_llcs), (1, 0));
        assert_eq!(compute_target_llcs(16, &topo.all_llcs), (2, 0));
        assert_eq!(compute_target_llcs(24, &topo.all_llcs), (3, 0));
        assert_eq!(compute_target_llcs(32, &topo.all_llcs), (4, 0));
        assert_eq!(compute_target_llcs(10, &topo.all_llcs), (1, 1)); // 8+2 = 1 LLC + 1 core
        assert_eq!(compute_target_llcs(15, &topo.all_llcs), (1, 4)); // 8+7 = 1 LLC + 4 cores (ceil(7/2))
    }

    #[test]
    fn test_compute_target_llcs_4n() {
        let (topo, _total) = topo_4n();
        // 4N: 16 cores, 8 LLCs → 2 cores/LLC, 2 HTs/core → 4 cpus/LLC
        assert_eq!(compute_target_llcs(0, &topo.all_llcs), (0, 0));
        assert_eq!(compute_target_llcs(4, &topo.all_llcs), (1, 0)); // exactly 1 LLC
        assert_eq!(compute_target_llcs(5, &topo.all_llcs), (1, 1)); // 1 LLC + 1 extra core
        assert_eq!(compute_target_llcs(8, &topo.all_llcs), (2, 0)); // exactly 2 LLCs
        assert_eq!(compute_target_llcs(32, &topo.all_llcs), (8, 0)); // all 8 LLCs
    }

    #[test]
    fn test_compute_target_llcs_het() {
        // LLC0: 2 cores × 2 HTs = 4 CPUs
        // LLC1: 4 cores × 2 HTs = 8 CPUs
        // Total: 12 CPUs across 2 unequal LLCs
        let (topo, _total) = make_het_test_topo(&[&[2, 4]], 2);
        let node = topo.nodes.get(&0).unwrap();
        assert_eq!(compute_target_llcs(0, &node.llcs), (0, 0));
        assert_eq!(compute_target_llcs(4, &node.llcs), (1, 0)); // exactly LLC0
        assert_eq!(compute_target_llcs(5, &node.llcs), (1, 1)); // LLC0 + 1 core from LLC1
        assert_eq!(compute_target_llcs(8, &node.llcs), (1, 2)); // LLC0 + 2 cores from LLC1
        assert_eq!(compute_target_llcs(12, &node.llcs), (2, 0)); // both LLCs full
    }

    #[test]
    fn test_compute_target_llcs_het_multi_node() {
        // Node 0: LLC0=2cores, LLC1=2cores → 4 CPUs per LLC, 8 total
        // Node 1: LLC2=4cores, LLC3=4cores → 8 CPUs per LLC, 16 total
        // Total: 24 CPUs across 2 nodes with different LLC sizes
        let (topo, _total) = make_het_test_topo(&[&[2, 2], &[4, 4]], 2);
        let node0 = topo.nodes.get(&0).unwrap();
        let node1 = topo.nodes.get(&1).unwrap();

        // node 0: each LLC has 4 CPUs, 2 cores
        assert_eq!(compute_target_llcs(0, &node0.llcs), (0, 0));
        assert_eq!(compute_target_llcs(4, &node0.llcs), (1, 0)); // exactly LLC0
        assert_eq!(compute_target_llcs(5, &node0.llcs), (1, 1)); // LLC0 + 1 core from LLC1
        assert_eq!(compute_target_llcs(8, &node0.llcs), (2, 0)); // both LLCs full

        // node 1: each LLC has 8 CPUs, 4 cores — different size than node 0
        assert_eq!(compute_target_llcs(0, &node1.llcs), (0, 0));
        assert_eq!(compute_target_llcs(8, &node1.llcs), (1, 0)); // exactly LLC2
        assert_eq!(compute_target_llcs(10, &node1.llcs), (1, 1)); // LLC2 + 1 core
        assert_eq!(compute_target_llcs(16, &node1.llcs), (2, 0)); // both LLCs full
    }

    // --- StickyDynamic per-node lifecycle: 1N ---

    #[test]
    fn test_sd_1n_grow_one_llc() {
        let (topo, _total) = topo_1n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 1)];

        // 4 alloc units on N0 = 8 cpus = 1 LLC.
        let allocs = vec![make_sd_alloc(&[0], &[4])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 1);
        assert_eq!(pool.total_free_llcs(), 1);
        assert_sd_llc_conservation(&pool, &layers, 2);
    }

    #[test]
    fn test_sd_1n_grow_two_llcs() {
        let (topo, _total) = topo_1n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 1)];

        let allocs = vec![make_sd_alloc(&[0], &[8])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 2);
        assert_eq!(pool.total_free_llcs(), 0);
        assert_sd_llc_conservation(&pool, &layers, 2);
    }

    #[test]
    fn test_sd_1n_grow_then_shrink() {
        let (topo, _total) = topo_1n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 1)];

        // Grow to 2 LLCs.
        let allocs = vec![make_sd_alloc(&[0], &[8])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 2);

        // Shrink to 1 LLC.
        let allocs = vec![make_sd_alloc(&[0], &[4])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 1);
        assert_eq!(pool.total_free_llcs(), 1);
    }

    #[test]
    fn test_sd_1n_roundtrip_zero_full_zero() {
        let (topo, _total) = topo_1n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let initial_free = pool.total_free_llcs();
        let mut layers = vec![SdLayerState::new("L0", 1)];

        // Grow to full.
        let allocs = vec![make_sd_alloc(&[0], &[8])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 2);

        // Shrink to zero.
        let allocs = vec![make_sd_alloc(&[0], &[0])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 0);
        assert_eq!(pool.total_free_llcs(), initial_free);
    }

    #[test]
    fn test_sd_1n_spillover() {
        // 1N: target 10 cpus = 1 full LLC + 1 extra core.
        let (topo, _total) = topo_1n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 1)];

        let allocs = vec![make_sd_alloc(&[0], &[5])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 1);
        // Spillover consumed from the free LLC but reset.
        assert_sd_llc_conservation(&pool, &layers, 2);
    }

    // --- StickyDynamic per-node lifecycle: 2N ---

    #[test]
    fn test_sd_2n_all_on_node0() {
        // All allocation on N0.
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2)];

        let allocs = vec![make_sd_alloc(&[0, 0], &[8, 0])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 2);
        assert_eq!(layers[0].llcs_on_node(0), 2);
        assert_eq!(layers[0].llcs_on_node(1), 0);
        assert_eq!(pool.total_free_llcs(), 2); // N1 free
        assert_sd_llc_conservation(&pool, &layers, 4);
    }

    #[test]
    fn test_sd_2n_split_across_nodes() {
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2)];

        let allocs = vec![make_sd_alloc(&[0, 0], &[4, 4])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 2);
        assert_eq!(layers[0].llcs_on_node(0), 1);
        assert_eq!(layers[0].llcs_on_node(1), 1);
        assert_sd_llc_conservation(&pool, &layers, 4);
    }

    #[test]
    fn test_sd_2n_asymmetric() {
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2)];

        let allocs = vec![make_sd_alloc(&[0, 0], &[8, 4])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 3);
        assert_eq!(layers[0].llcs_on_node(0), 2);
        assert_eq!(layers[0].llcs_on_node(1), 1);
        assert_sd_llc_conservation(&pool, &layers, 4);
    }

    #[test]
    fn test_sd_2n_grow_shrink_one_node() {
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2)];

        // Step 1: 2 LLCs on N0, 0 on N1.
        let allocs = vec![make_sd_alloc(&[0, 0], &[8, 0])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].llcs_on_node(0), 2);
        assert_eq!(layers[0].llcs_on_node(1), 0);

        // Step 2: 1 LLC on N0, 1 LLC on N1.
        let allocs = vec![make_sd_alloc(&[0, 0], &[4, 4])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].llcs_on_node(0), 1);
        assert_eq!(layers[0].llcs_on_node(1), 1);
        assert_sd_llc_conservation(&pool, &layers, 4);
    }

    #[test]
    fn test_sd_2n_two_layers_same_node() {
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2), SdLayerState::new("L1", 2)];

        // Both want 1 LLC on N0.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[4, 0]),
            make_sd_alloc(&[0, 0], &[4, 0]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 1);
        assert_eq!(layers[1].total_llcs(), 1);
        assert_eq!(layers[0].llcs_on_node(0), 1);
        assert_eq!(layers[1].llcs_on_node(0), 1);
        assert_sd_no_duplicate_llcs(&layers, "same node");
        assert_sd_llc_conservation(&pool, &layers, 4);
    }

    #[test]
    fn test_sd_2n_two_layers_different_nodes() {
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2), SdLayerState::new("L1", 2)];

        // L0 wants all of N0, L1 wants all of N1.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[8, 0]),
            make_sd_alloc(&[0, 0], &[0, 8]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 2);
        assert_eq!(layers[1].total_llcs(), 2);
        assert_eq!(layers[0].llcs_on_node(0), 2);
        assert_eq!(layers[1].llcs_on_node(1), 2);
        assert_eq!(pool.total_free_llcs(), 0);
        assert_sd_no_duplicate_llcs(&layers, "diff nodes");
    }

    #[test]
    fn test_sd_2n_one_shrinks_other_grows() {
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2), SdLayerState::new("L1", 2)];

        // Step 1: L0 takes all 4 LLCs.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[8, 8]),
            make_sd_alloc(&[0, 0], &[0, 0]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 4);
        assert_eq!(layers[1].total_llcs(), 0);

        // Step 2: L0 shrinks, L1 grows — split evenly.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[4, 4]),
            make_sd_alloc(&[0, 0], &[4, 4]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 2);
        assert_eq!(layers[1].total_llcs(), 2);
        assert_sd_no_duplicate_llcs(&layers, "shrink/grow");
        assert_sd_llc_conservation(&pool, &layers, 4);
    }

    #[test]
    fn test_sd_2n_swap_allocations() {
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2), SdLayerState::new("L1", 2)];

        // Step 1: L0 on N0, L1 on N1.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[8, 0]),
            make_sd_alloc(&[0, 0], &[0, 8]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].llcs_on_node(0), 2);
        assert_eq!(layers[1].llcs_on_node(1), 2);

        // Step 2: Swap — L0 to N1, L1 to N0.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[0, 8]),
            make_sd_alloc(&[0, 0], &[8, 0]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].llcs_on_node(1), 2);
        assert_eq!(layers[0].llcs_on_node(0), 0);
        assert_eq!(layers[1].llcs_on_node(0), 2);
        assert_eq!(layers[1].llcs_on_node(1), 0);
        assert_sd_no_duplicate_llcs(&layers, "swap");
        assert_sd_llc_conservation(&pool, &layers, 4);
    }

    // --- StickyDynamic per-node lifecycle: 4N ---

    #[test]
    fn test_sd_4n_spread_all_nodes() {
        let (topo, _total) = topo_4n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 4)];

        let allocs = vec![make_sd_alloc(&[0, 0, 0, 0], &[2, 2, 2, 2])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 4);
        for n in 0..4 {
            assert_eq!(layers[0].llcs_on_node(n), 1, "node {} should have 1 LLC", n);
        }
        assert_sd_llc_conservation(&pool, &layers, 8);
    }

    #[test]
    fn test_sd_4n_all_on_single_node() {
        let (topo, _total) = topo_4n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 4)];

        let allocs = vec![make_sd_alloc(&[0, 0, 0, 0], &[4, 0, 0, 0])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 2);
        assert_eq!(layers[0].llcs_on_node(0), 2);
        assert_sd_llc_conservation(&pool, &layers, 8);
    }

    #[test]
    fn test_sd_4n_overflow_target_exceeds_node() {
        // Target on N0 exceeds node capacity (2 LLCs = 8 cpus).
        let (topo, _total) = topo_4n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 4)];

        // 8 au on N0 = 16 cpus, but N0 only has 2 LLCs (8 cpus).
        let allocs = vec![make_sd_alloc(&[0, 0, 0, 0], &[8, 0, 0, 0])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].llcs_on_node(0), 2);
        assert_eq!(layers[0].total_llcs(), 2); // Can't exceed N0's supply
        assert_sd_llc_conservation(&pool, &layers, 8);
    }

    #[test]
    fn test_sd_4n_four_layers_one_per_node() {
        let (topo, _total) = topo_4n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![
            SdLayerState::new("L0", 4),
            SdLayerState::new("L1", 4),
            SdLayerState::new("L2", 4),
            SdLayerState::new("L3", 4),
        ];

        let allocs = vec![
            make_sd_alloc(&[0, 0, 0, 0], &[4, 0, 0, 0]),
            make_sd_alloc(&[0, 0, 0, 0], &[0, 4, 0, 0]),
            make_sd_alloc(&[0, 0, 0, 0], &[0, 0, 4, 0]),
            make_sd_alloc(&[0, 0, 0, 0], &[0, 0, 0, 4]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        for (i, layer) in layers.iter().enumerate() {
            assert_eq!(layer.total_llcs(), 2, "L{} should have 2 LLCs", i);
            assert_eq!(
                layer.llcs_on_node(i),
                2,
                "L{} should have 2 LLCs on node {}",
                i,
                i
            );
        }
        assert_eq!(pool.total_free_llcs(), 0);
        assert_sd_no_duplicate_llcs(&layers, "4 layers");
    }

    // --- StickyDynamic stress tests ---

    #[test]
    fn test_sd_stress_10_cycle_oscillation() {
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2)];

        for cycle in 0..10 {
            // Phase A: 2 LLCs on N0, 1 on N1.
            let allocs = vec![make_sd_alloc(&[0, 0], &[8, 4])];
            simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
            assert_eq!(layers[0].llcs_on_node(0), 2, "cycle {} A: N0", cycle);
            assert_eq!(layers[0].llcs_on_node(1), 1, "cycle {} A: N1", cycle);
            assert_sd_llc_conservation(&pool, &layers, 4);

            // Phase B: 1 LLC on N0, 0 LLCs on N1.
            let allocs = vec![make_sd_alloc(&[0, 0], &[4, 0])];
            simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
            assert_eq!(layers[0].llcs_on_node(0), 1, "cycle {} B: N0", cycle);
            assert_eq!(layers[0].llcs_on_node(1), 0, "cycle {} B: N1", cycle);
            assert_sd_llc_conservation(&pool, &layers, 4);
        }
    }

    #[test]
    fn test_sd_stress_multi_layer_competition() {
        // 3 layers compete on 2N over 5 rebalancing steps.
        // 2N: 4 LLCs (2/node), au=2.
        // The SD algorithm acquires in reverse target order (largest first).
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![
            SdLayerState::new("L0", 2),
            SdLayerState::new("L1", 2),
            SdLayerState::new("L2", 2),
        ];

        // Step 0: L0 splits (1 N0 + 1 N1, total=16 largest), L1 gets 1 N0, L2 gets 1 N1.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[4, 4]),
            make_sd_alloc(&[0, 0], &[4, 0]),
            make_sd_alloc(&[0, 0], &[0, 4]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].llcs_on_node(0), 1, "s0 L0 N0");
        assert_eq!(layers[0].llcs_on_node(1), 1, "s0 L0 N1");
        assert_eq!(layers[1].llcs_on_node(0), 1, "s0 L1 N0");
        assert_eq!(layers[2].llcs_on_node(1), 1, "s0 L2 N1");
        assert_sd_llc_conservation(&pool, &layers, 4);
        assert_sd_no_duplicate_llcs(&layers, "s0");

        // Step 1: All shrink to 0.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[0, 0]),
            make_sd_alloc(&[0, 0], &[0, 0]),
            make_sd_alloc(&[0, 0], &[0, 0]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 0);
        assert_eq!(layers[1].total_llcs(), 0);
        assert_eq!(layers[2].total_llcs(), 0);
        assert_sd_llc_conservation(&pool, &layers, 4);

        // Step 2: L1 gets 2 on N0 (largest, acquires first), L0 gets 1 N1, L2 gets 1 N1.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[0, 4]),
            make_sd_alloc(&[0, 0], &[8, 0]),
            make_sd_alloc(&[0, 0], &[0, 4]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[1].llcs_on_node(0), 2, "s2 L1 N0");
        assert_eq!(layers[0].llcs_on_node(1), 1, "s2 L0 N1");
        assert_eq!(layers[2].llcs_on_node(1), 1, "s2 L2 N1");
        assert_sd_llc_conservation(&pool, &layers, 4);
        assert_sd_no_duplicate_llcs(&layers, "s2");

        // Step 3: L1 shrinks to 1 N0, L0 gets 1 N0, L2 stays 1 N1.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[4, 0]),
            make_sd_alloc(&[0, 0], &[4, 0]),
            make_sd_alloc(&[0, 0], &[0, 4]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].llcs_on_node(0), 1, "s3 L0 N0");
        assert_eq!(layers[1].llcs_on_node(0), 1, "s3 L1 N0");
        assert_eq!(layers[2].llcs_on_node(1), 1, "s3 L2 N1");
        assert_sd_llc_conservation(&pool, &layers, 4);
        assert_sd_no_duplicate_llcs(&layers, "s3");

        // Step 4: All back to 0.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[0, 0]),
            make_sd_alloc(&[0, 0], &[0, 0]),
            make_sd_alloc(&[0, 0], &[0, 0]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 0);
        assert_eq!(layers[1].total_llcs(), 0);
        assert_eq!(layers[2].total_llcs(), 0);
        assert_sd_llc_conservation(&pool, &layers, 4);
    }

    #[test]
    fn test_sd_stress_staircase() {
        // 4N: grow across nodes then shrink back, verify each step.
        let (topo, _total) = topo_4n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 4)];

        // Staircase up: 1 node at a time.
        let staircase = vec![
            vec![2, 0, 0, 0], // 1 LLC on N0
            vec![2, 2, 0, 0], // + 1 on N1
            vec![2, 2, 2, 0], // + 1 on N2
            vec![2, 2, 2, 2], // + 1 on N3
        ];
        for (i, unpinned) in staircase.iter().enumerate() {
            let allocs = vec![make_sd_alloc(&[0, 0, 0, 0], unpinned)];
            simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
            let total_cpus: usize = unpinned.iter().sum::<usize>() * au;
            let expected_llcs = compute_target_llcs(total_cpus, &topo.all_llcs).0;
            assert_eq!(
                layers[0].total_llcs(),
                expected_llcs,
                "staircase up step {}",
                i
            );
            for n in 0..4 {
                let expected_n =
                    compute_target_llcs(unpinned[n] * au, &topo.nodes.get(&n).unwrap().llcs).0;
                assert_eq!(
                    layers[0].llcs_on_node(n),
                    expected_n,
                    "step {} node {}",
                    i,
                    n
                );
            }
            assert_sd_llc_conservation(&pool, &layers, 8);
        }

        // Staircase down.
        for (i, unpinned) in staircase.iter().rev().enumerate() {
            let allocs = vec![make_sd_alloc(&[0, 0, 0, 0], unpinned)];
            simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
            for n in 0..4 {
                let expected_n =
                    compute_target_llcs(unpinned[n] * au, &topo.nodes.get(&n).unwrap().llcs).0;
                assert_eq!(
                    layers[0].llcs_on_node(n),
                    expected_n,
                    "down step {} node {}",
                    i,
                    n
                );
            }
            assert_sd_llc_conservation(&pool, &layers, 8);
        }
    }

    #[test]
    fn test_sd_stress_idempotent() {
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2)];

        let allocs = vec![make_sd_alloc(&[0, 0], &[4, 4])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        let stable = layers[0].assigned_llcs.clone();

        for i in 0..5 {
            simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
            assert_eq!(
                layers[0].assigned_llcs, stable,
                "iteration {}: should be identical",
                i
            );
            assert_sd_llc_conservation(&pool, &layers, 4);
        }
    }

    // --- StickyDynamic edge cases ---

    #[test]
    fn test_sd_zero_target_everywhere() {
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2)];

        let allocs = vec![make_sd_alloc(&[0, 0], &[0, 0])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 0);
        assert_eq!(pool.total_free_llcs(), 4);
    }

    #[test]
    fn test_sd_all_llcs_taken_by_l0() {
        // L0 takes all, L1 gets zero.
        let (topo, _total) = topo_2n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 2), SdLayerState::new("L1", 2)];

        let allocs = vec![
            make_sd_alloc(&[0, 0], &[8, 8]),
            make_sd_alloc(&[0, 0], &[0, 0]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 4);
        assert_eq!(layers[1].total_llcs(), 0);
        assert_sd_llc_conservation(&pool, &layers, 4);

        // L1 wants some but L0 still holds all — L1 gets nothing.
        let allocs = vec![
            make_sd_alloc(&[0, 0], &[8, 8]),
            make_sd_alloc(&[0, 0], &[4, 0]),
        ];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        assert_eq!(layers[0].total_llcs(), 4);
        assert_eq!(layers[1].total_llcs(), 0);
        assert_sd_llc_conservation(&pool, &layers, 4);
    }

    #[test]
    fn test_sd_target_exceeds_supply_per_node() {
        // Target per node exceeds supply. Should get capped.
        let (topo, _total) = topo_4n();
        let au = 2;
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let mut layers = vec![SdLayerState::new("L0", 4)];

        // 6 au per node = 12 cpus/node, but each node only has 2 LLCs (8 cpus).
        let allocs = vec![make_sd_alloc(&[0, 0, 0, 0], &[6, 6, 6, 6])];
        simulate_sd_recompute(&mut pool, &mut layers, &allocs, au, &topo);
        // Should get 2 LLCs per node (8 total = all).
        assert_eq!(layers[0].total_llcs(), 8);
        for n in 0..4 {
            assert_eq!(layers[0].llcs_on_node(n), 2, "node {} capped at 2", n);
        }
        assert_eq!(pool.total_free_llcs(), 0);
        assert_sd_llc_conservation(&pool, &layers, 8);
    }

    // --- largest_remainder ---

    #[test]
    fn test_lr_exact_sum() {
        let result = largest_remainder(10, &[3.0, 3.0, 4.0]);
        assert_eq!(result.iter().sum::<usize>(), 10);
    }

    #[test]
    fn test_lr_proportional() {
        let result = largest_remainder(100, &[1.0, 2.0, 3.0]);
        // ~17, ~33, ~50 — sum must be exactly 100.
        assert_eq!(result.iter().sum::<usize>(), 100);
        assert!(result[0] >= 16 && result[0] <= 17);
        assert!(result[1] >= 33 && result[1] <= 34);
        assert!(result[2] >= 49 && result[2] <= 50);
    }

    #[test]
    fn test_lr_equal_quotas() {
        let result = largest_remainder(10, &[1.0, 1.0, 1.0]);
        assert_eq!(result.iter().sum::<usize>(), 10);
        // 10/3 = 3.333... → two get 3, one gets 4.
        assert!(result.iter().all(|&v| v == 3 || v == 4));
        assert_eq!(result.iter().filter(|&&v| v == 4).count(), 1);
    }

    #[test]
    fn test_lr_zero_quotas() {
        let result = largest_remainder(10, &[0.0, 0.0, 0.0]);
        assert_eq!(result, vec![0, 0, 0]);
    }

    #[test]
    fn test_lr_single_entry() {
        let result = largest_remainder(42, &[7.0]);
        assert_eq!(result, vec![42]);
    }

    #[test]
    fn test_lr_large_remainder() {
        // 7/3 = 2.333... → floors = [2,2,2] = 6, remainder = 1.
        let result = largest_remainder(7, &[1.0, 1.0, 1.0]);
        assert_eq!(result.iter().sum::<usize>(), 7);
        assert_eq!(result.iter().filter(|&&v| v == 3).count(), 1);
        assert_eq!(result.iter().filter(|&&v| v == 2).count(), 2);
    }

    #[test]
    fn test_lr_empty() {
        let result = largest_remainder(10, &[]);
        assert!(result.is_empty());
    }

    #[test]
    fn test_lr_one_zero_one_nonzero() {
        let result = largest_remainder(10, &[0.0, 5.0]);
        assert_eq!(result, vec![0, 10]);
    }

    #[test]
    fn test_lr_total_zero() {
        let result = largest_remainder(0, &[3.0, 7.0]);
        assert_eq!(result, vec![0, 0]);
    }

    // --- round_targets_to_alloc_units ---

    #[test]
    fn test_rta_unit_1() {
        let targets = vec![(10, 2), (15, 5)];
        let result = round_targets_to_alloc_units(&targets, 1, 100);
        assert_eq!(result, targets);
    }

    #[test]
    fn test_rta_round_up() {
        // alloc_unit=2: 3→4, 5→6
        let targets = vec![(3, 1), (5, 0)];
        let result = round_targets_to_alloc_units(&targets, 2, 100);
        assert_eq!(result, vec![(4, 2), (6, 0)]);
    }

    #[test]
    fn test_rta_already_aligned() {
        let targets = vec![(4, 2), (6, 0)];
        let result = round_targets_to_alloc_units(&targets, 2, 100);
        assert_eq!(result, targets);
    }

    #[test]
    fn test_rta_caps_at_total() {
        let targets = vec![(99, 0)];
        let result = round_targets_to_alloc_units(&targets, 2, 16);
        assert_eq!(result, vec![(16, 0)]);
    }

    // --- alloc_unit ---

    #[test]
    fn test_alloc_unit_partial() {
        let (topo, _) = topo_1n();
        let pool = CpuPool::new(topo, true).unwrap();
        assert_eq!(pool.alloc_unit(), 1);
    }

    #[test]
    fn test_alloc_unit_full_core() {
        let (topo, _) = topo_1n();
        let pool = CpuPool::new(topo, false).unwrap();
        // 2 HTs per core.
        assert_eq!(pool.alloc_unit(), 2);
    }

    // =====================================================================
    // Spread + non-spread per-node integration tests
    //
    // Verify that spread layers (even-split budget) and non-spread layers
    // (water_fill budget) coexist correctly through the per-node grow/shrink
    // simulation. Spread layers should get equal CPUs per node; non-spread
    // layers fill the remaining capacity.
    // =====================================================================

    #[test]
    fn test_nonsd_2n_spread_even_split() {
        // One spread layer on 2N: target=16, should get 8 per node.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::NodeSpread)];
        let order = get_core_order_per_node(&topo, &specs, 0);
        let mut layers = vec![NonSdLayerState::new("L0", 2, total, order)];
        let norders = vec![vec![0, 1]];

        // Even split: 4 au per node = 8 CPUs per node.
        let allocs = vec![make_alloc(&[0, 0], &[4, 4])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 16)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 16);
        assert_node_cpus(&layers[0], &[8, 8], "spread even split");
        assert_cpu_conservation(&pool, &layers, total);
    }

    #[test]
    fn test_nonsd_2n_spread_and_linear_coexist() {
        // Spread layer + Linear layer on 2N. Spread gets 8/8, Linear gets
        // remaining capacity proportional to its budget on preferred node.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();

        let spread_specs = vec![test_layer_spec("S", LayerGrowthAlgo::NodeSpread)];
        let spread_order = get_core_order_per_node(&topo, &spread_specs, 0);
        let linear_specs = vec![test_layer_spec("L", LayerGrowthAlgo::Linear)];
        let linear_order = get_core_order_per_node(&topo, &linear_specs, 0);

        let mut layers = vec![
            NonSdLayerState::new("S", 2, total, spread_order),
            NonSdLayerState::new("L", 2, total, linear_order),
        ];
        let norders = vec![vec![0, 1], vec![0, 1]];

        // Spread: 8 CPUs (4 au) per node. Linear: 8 CPUs all on N0.
        let allocs = vec![make_alloc(&[0, 0], &[4, 4]), make_alloc(&[0, 0], &[4, 0])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 16), (1, 8)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 16);
        assert_node_cpus(&layers[0], &[8, 8], "spread");
        assert_eq!(layers[1].nr_cpus, 8);
        assert_node_cpus(&layers[1], &[8, 0], "linear");
        assert_no_cpu_overlap(&layers, "spread+linear");
        assert_cpu_conservation(&pool, &layers, total);
    }

    #[test]
    fn test_nonsd_2n_spread_grow_shrink_roundtrip() {
        // Spread layer: grow to 16, shrink to 8, grow back to 16.
        // Per-node counts should be even at each step.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::NodeSpread)];
        let order = get_core_order_per_node(&topo, &specs, 0);
        let mut layers = vec![NonSdLayerState::new("L0", 2, total, order)];
        let norders = vec![vec![0, 1]];

        // Grow to 16 (8 per node).
        let allocs = vec![make_alloc(&[0, 0], &[4, 4])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 16)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_node_cpus(&layers[0], &[8, 8], "grow to 16");
        let cpus_at_16 = layers[0].cpus.clone();

        // Shrink to 8 (4 per node).
        let allocs = vec![make_alloc(&[0, 0], &[2, 2])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 8)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_node_cpus(&layers[0], &[4, 4], "shrink to 8");
        assert_cpu_conservation(&pool, &layers, total);

        // Grow back to 16.
        let allocs = vec![make_alloc(&[0, 0], &[4, 4])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 16)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_node_cpus(&layers[0], &[8, 8], "grow back to 16");
        assert_eq!(layers[0].cpus, cpus_at_16, "roundtrip same CPUs");
    }

    #[test]
    fn test_nonsd_2n_spread_roundrobin_order() {
        // RoundRobin (spread) on 2N: verify LLC interleave within each node.
        // With even split, each node gets half. Within each node the LLC
        // interleave determines which cores are picked first.
        let (topo, total) = topo_2n();
        let mut pool = CpuPool::new(topo.clone(), false).unwrap();
        let au = pool.alloc_unit();
        let specs = vec![test_layer_spec("L0", LayerGrowthAlgo::RoundRobin)];
        let order = get_core_order_per_node(&topo, &specs, 0);
        let mut layers = vec![NonSdLayerState::new("L0", 2, total, order)];
        let norders = vec![vec![0, 1]];

        // Grow to 8 CPUs: 4 per node (2 cores per node). With LLC interleave
        // these should come from different LLCs within each node.
        let allocs = vec![make_alloc(&[0, 0], &[2, 2])];
        simulate_non_sd_recompute(
            &mut pool,
            &mut layers,
            &[(0, 8)],
            &allocs,
            &norders,
            au,
            &topo,
        );
        assert_eq!(layers[0].nr_cpus, 8);
        assert_node_cpus(&layers[0], &[4, 4], "roundrobin even split");
        assert_cpu_conservation(&pool, &layers, total);
    }
}
