solverforge_scoring/constraint/

grouped.rs

// Zero-erasure grouped constraint for group-by operations.
//
// Provides incremental scoring for constraints that group entities and
// apply collectors to compute aggregate scores.
// All type information is preserved at compile time - no Arc, no dyn.

use std::collections::HashMap;
use std::hash::Hash;
use std::marker::PhantomData;

use solverforge_core::score::Score;
use solverforge_core::{ConstraintRef, ImpactType};

use crate::api::constraint_set::IncrementalConstraint;
use crate::stream::collector::{Accumulator, UniCollector};

// Zero-erasure constraint that groups entities by key and scores based on collector results.
//
// This enables incremental scoring for group-by operations:
// - Tracks which entities belong to which group
// - Maintains collector state per group
// - Computes score deltas when entities are added/removed
//
// All type parameters are concrete - no trait objects, no Arc allocations.
//
// # Type Parameters
//
// - `S` - Solution type
// - `A` - Entity type
// - `K` - Group key type
// - `E` - Extractor function for entities
// - `KF` - Key function
// - `C` - Collector type
// - `W` - Weight function
// - `Sc` - Score type
//
// # Example
//
// ```
// use solverforge_scoring::constraint::grouped::GroupedUniConstraint;
// use solverforge_scoring::stream::collector::count;
// use solverforge_scoring::api::constraint_set::IncrementalConstraint;
// use solverforge_core::{ConstraintRef, ImpactType};
// use solverforge_core::score::SimpleScore;
//
// #[derive(Clone, Hash, PartialEq, Eq)]
// struct Shift { employee_id: usize }
//
// #[derive(Clone)]
// struct Solution { shifts: Vec<Shift> }
//
// // Penalize based on squared workload per employee
// let constraint = GroupedUniConstraint::new(
//     ConstraintRef::new("", "Balanced workload"),
//     ImpactType::Penalty,
//     |s: &Solution| &s.shifts,
//     |shift: &Shift| shift.employee_id,
//     count::<Shift>(),
//     |count: &usize| SimpleScore::of((*count * *count) as i64),
//     false,
// );
//
// let solution = Solution {
//     shifts: vec![
//         Shift { employee_id: 1 },
//         Shift { employee_id: 1 },
//         Shift { employee_id: 1 },
//         Shift { employee_id: 2 },
//     ],
// };
//
// // Employee 1: 3 shifts -> 9 penalty
// // Employee 2: 1 shift -> 1 penalty
// // Total: -10
// assert_eq!(constraint.evaluate(&solution), SimpleScore::of(-10));
// ```
pub struct GroupedUniConstraint<S, A, K, E, KF, C, W, Sc>
where
    C: UniCollector<A>,
    Sc: Score,
{
    constraint_ref: ConstraintRef,
    impact_type: ImpactType,
    extractor: E,
    key_fn: KF,
    collector: C,
    weight_fn: W,
    is_hard: bool,
    expected_descriptor: Option<usize>,
    // Group key -> accumulator (scores computed on-the-fly, no cloning)
    groups: HashMap<K, C::Accumulator>,
    // Entity index -> group key (for tracking which group an entity belongs to)
    entity_groups: HashMap<usize, K>,
    // Entity index -> extracted value (for correct retraction after entity mutation)
    entity_values: HashMap<usize, C::Value>,
    _phantom: PhantomData<(fn() -> S, fn() -> A, fn() -> Sc)>,
}

impl<S, A, K, E, KF, C, W, Sc> GroupedUniConstraint<S, A, K, E, KF, C, W, Sc>
where
    S: Send + Sync + 'static,
    A: Clone + Send + Sync + 'static,
    K: Clone + Eq + Hash + Send + Sync + 'static,
    E: Fn(&S) -> &[A] + Send + Sync,
    KF: Fn(&A) -> K + Send + Sync,
    C: UniCollector<A> + Send + Sync + 'static,
    C::Accumulator: Send + Sync,
    C::Result: Send + Sync,
    W: Fn(&C::Result) -> Sc + Send + Sync,
    Sc: Score + 'static,
{
    // Creates a new zero-erasure grouped constraint.
    //
    // # Arguments
    //
    // * `constraint_ref` - Identifier for this constraint
    // * `impact_type` - Whether to penalize or reward
    // * `extractor` - Function to get entity slice from solution
    // * `key_fn` - Function to extract group key from entity
    // * `collector` - Collector to aggregate entities per group
    // * `weight_fn` - Function to compute score from collector result
    // * `is_hard` - Whether this is a hard constraint
    pub fn new(
        constraint_ref: ConstraintRef,
        impact_type: ImpactType,
        extractor: E,
        key_fn: KF,
        collector: C,
        weight_fn: W,
        is_hard: bool,
    ) -> Self {
        Self {
            constraint_ref,
            impact_type,
            extractor,
            key_fn,
            collector,
            weight_fn,
            is_hard,
            expected_descriptor: None,
            groups: HashMap::new(),
            entity_groups: HashMap::new(),
            entity_values: HashMap::new(),
            _phantom: PhantomData,
        }
    }

    pub fn with_descriptor(mut self, descriptor_index: usize) -> Self {
        self.expected_descriptor = Some(descriptor_index);
        self
    }

    // Computes the score contribution for a group's result.
    fn compute_score(&self, result: &C::Result) -> Sc {
        let base = (self.weight_fn)(result);
        match self.impact_type {
            ImpactType::Penalty => -base,
            ImpactType::Reward => base,
        }
    }
}

impl<S, A, K, E, KF, C, W, Sc> IncrementalConstraint<S, Sc>
    for GroupedUniConstraint<S, A, K, E, KF, C, W, Sc>
where
    S: Send + Sync + 'static,
    A: Clone + Send + Sync + 'static,
    K: Clone + Eq + Hash + Send + Sync + 'static,
    E: Fn(&S) -> &[A] + Send + Sync,
    KF: Fn(&A) -> K + Send + Sync,
    C: UniCollector<A> + Send + Sync + 'static,
    C::Accumulator: Send + Sync,
    C::Result: Send + Sync,
    C::Value: Send + Sync,
    W: Fn(&C::Result) -> Sc + Send + Sync,
    Sc: Score + 'static,
{
    fn evaluate(&self, solution: &S) -> Sc {
        let entities = (self.extractor)(solution);

        // Group entities by key
        let mut groups: HashMap<K, C::Accumulator> = HashMap::new();

        for entity in entities {
            let key = (self.key_fn)(entity);
            let value = self.collector.extract(entity);
            let acc = groups
                .entry(key)
                .or_insert_with(|| self.collector.create_accumulator());
            acc.accumulate(&value);
        }

        // Sum scores for all groups
        let mut total = Sc::zero();
        for acc in groups.values() {
            let result = acc.finish();
            total = total + self.compute_score(&result);
        }

        total
    }

    fn match_count(&self, solution: &S) -> usize {
        let entities = (self.extractor)(solution);

        // Count unique groups
        let mut groups: HashMap<K, ()> = HashMap::new();
        for entity in entities {
            let key = (self.key_fn)(entity);
            groups.insert(key, ());
        }

        groups.len()
    }

    fn initialize(&mut self, solution: &S) -> Sc {
        self.reset();

        let entities = (self.extractor)(solution);
        let mut total = Sc::zero();

        for (idx, entity) in entities.iter().enumerate() {
            total = total + self.insert_entity(entities, idx, entity);
        }

        total
    }

    fn on_insert(&mut self, solution: &S, entity_index: usize, descriptor_index: usize) -> Sc {
        if let Some(expected) = self.expected_descriptor {
            if descriptor_index != expected {
                return Sc::zero();
            }
        }
        let entities = (self.extractor)(solution);
        if entity_index >= entities.len() {
            return Sc::zero();
        }

        let entity = &entities[entity_index];
        self.insert_entity(entities, entity_index, entity)
    }

    fn on_retract(&mut self, solution: &S, entity_index: usize, descriptor_index: usize) -> Sc {
        if let Some(expected) = self.expected_descriptor {
            if descriptor_index != expected {
                return Sc::zero();
            }
        }
        let entities = (self.extractor)(solution);
        self.retract_entity(entities, entity_index)
    }

    fn reset(&mut self) {
        self.groups.clear();
        self.entity_groups.clear();
        self.entity_values.clear();
    }

    fn name(&self) -> &str {
        &self.constraint_ref.name
    }

    fn is_hard(&self) -> bool {
        self.is_hard
    }

    fn constraint_ref(&self) -> ConstraintRef {
        self.constraint_ref.clone()
    }
}

impl<S, A, K, E, KF, C, W, Sc> GroupedUniConstraint<S, A, K, E, KF, C, W, Sc>
where
    S: Send + Sync + 'static,
    A: Clone + Send + Sync + 'static,
    K: Clone + Eq + Hash + Send + Sync + 'static,
    E: Fn(&S) -> &[A] + Send + Sync,
    KF: Fn(&A) -> K + Send + Sync,
    C: UniCollector<A> + Send + Sync + 'static,
    C::Accumulator: Send + Sync,
    C::Result: Send + Sync,
    C::Value: Send + Sync,
    W: Fn(&C::Result) -> Sc + Send + Sync,
    Sc: Score + 'static,
{
    fn insert_entity(&mut self, _entities: &[A], entity_index: usize, entity: &A) -> Sc {
        let key = (self.key_fn)(entity);
        let value = self.collector.extract(entity);
        let impact = self.impact_type;

        // Get or create group accumulator
        let acc = self
            .groups
            .entry(key.clone())
            .or_insert_with(|| self.collector.create_accumulator());

        // Compute old score from current state (inlined to avoid borrow conflict)
        let old_base = (self.weight_fn)(&acc.finish());
        let old = match impact {
            ImpactType::Penalty => -old_base,
            ImpactType::Reward => old_base,
        };

        // Accumulate and compute new score
        acc.accumulate(&value);
        let new_base = (self.weight_fn)(&acc.finish());
        let new_score = match impact {
            ImpactType::Penalty => -new_base,
            ImpactType::Reward => new_base,
        };

        // Track entity -> group mapping and cache value for correct retraction
        self.entity_groups.insert(entity_index, key);
        self.entity_values.insert(entity_index, value);

        // Return delta (both scores computed fresh, no cloning)
        new_score - old
    }

    fn retract_entity(&mut self, _entities: &[A], entity_index: usize) -> Sc {
        // Find which group this entity belonged to
        let Some(key) = self.entity_groups.remove(&entity_index) else {
            return Sc::zero();
        };

        // Use cached value (entity may have been mutated since insert)
        let Some(value) = self.entity_values.remove(&entity_index) else {
            return Sc::zero();
        };
        let impact = self.impact_type;

        // Get the group accumulator
        let Some(acc) = self.groups.get_mut(&key) else {
            return Sc::zero();
        };

        // Compute old score from current state (inlined to avoid borrow conflict)
        let old_base = (self.weight_fn)(&acc.finish());
        let old = match impact {
            ImpactType::Penalty => -old_base,
            ImpactType::Reward => old_base,
        };

        // Retract and compute new score
        acc.retract(&value);
        let new_base = (self.weight_fn)(&acc.finish());
        let new_score = match impact {
            ImpactType::Penalty => -new_base,
            ImpactType::Reward => new_base,
        };

        // Return delta (both scores computed fresh, no cloning)
        new_score - old
    }
}

impl<S, A, K, E, KF, C, W, Sc> std::fmt::Debug for GroupedUniConstraint<S, A, K, E, KF, C, W, Sc>
where
    C: UniCollector<A>,
    Sc: Score,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GroupedUniConstraint")
            .field("name", &self.constraint_ref.name)
            .field("impact_type", &self.impact_type)
            .field("groups", &self.groups.len())
            .finish()
    }
}