loadwise-core 0.1.0

use rand::Rng;

use super::{SelectionContext, Strategy};
use crate::LoadMetric;

/// Power of Two Choices: randomly pick two nodes, select the one with lower load.
///
/// Provides near-optimal load distribution with O(1) selection time,
/// a significant improvement over pure random while avoiding the overhead
/// of scanning all nodes like [`LeastLoad`](super::LeastLoad).
///
/// When fewer than two candidates are eligible, the sole eligible node is
/// returned directly.
#[derive(Debug)]
pub struct PowerOfTwoChoices;

impl PowerOfTwoChoices {
    pub fn new() -> Self {
        Self
    }
}

impl Default for PowerOfTwoChoices {
    fn default() -> Self {
        Self
    }
}

impl<N: LoadMetric> Strategy<N> for PowerOfTwoChoices {
    fn select(&self, candidates: &[N], ctx: &SelectionContext) -> Option<usize> {
        let eligible: Vec<usize> = (0..candidates.len())
            .filter(|i| !ctx.is_excluded(*i))
            .collect();
        match eligible.len() {
            0 => None,
            1 => Some(eligible[0]),
            len => {
                let mut rng = rand::rng();
                let ai = rng.random_range(0..len);
                let mut bi = rng.random_range(0..len - 1);
                if bi >= ai {
                    bi += 1;
                }
                let a = eligible[ai];
                let b = eligible[bi];
                if candidates[a].load_score() <= candidates[b].load_score() {
                    Some(a)
                } else {
                    Some(b)
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    struct L(f64);
    impl LoadMetric for L {
        fn load_score(&self) -> f64 {
            self.0
        }
    }

    #[test]
    fn picks_lower_of_two() {
        let p2c = PowerOfTwoChoices::new();
        // With only 2 candidates, always picks the lower-load one
        let nodes = [L(10.0), L(1.0)];
        let ctx = SelectionContext::default();
        assert_eq!(p2c.select(&nodes, &ctx), Some(1));
    }

    #[test]
    fn skips_excluded() {
        let p2c = PowerOfTwoChoices::new();
        let nodes = [L(1.0), L(2.0), L(3.0)];
        let ctx = SelectionContext::builder().exclude(vec![0]).build();
        // Best node (0) is excluded, should pick from {1, 2}
        let idx = p2c.select(&nodes, &ctx).unwrap();
        assert!(idx == 1 || idx == 2);
    }

    #[test]
    fn all_excluded_returns_none() {
        let p2c = PowerOfTwoChoices::new();
        let nodes = [L(1.0)];
        let ctx = SelectionContext::builder().exclude(vec![0]).build();
        assert_eq!(p2c.select(&nodes, &ctx), None);
    }
}