pumpkin-propagators 0.5.0

use pumpkin_checking::AtomicConstraint;
use pumpkin_checking::CheckerVariable;
use pumpkin_checking::InferenceChecker;
use pumpkin_checking::IntExt;
use pumpkin_checking::VariableState;
use pumpkin_core::asserts::pumpkin_assert_simple;
use pumpkin_core::declare_inference_label;
use pumpkin_core::predicate;
use pumpkin_core::predicates::Predicate;
use pumpkin_core::predicates::PropositionalConjunction;
use pumpkin_core::proof::ConstraintTag;
use pumpkin_core::proof::InferenceCode;
use pumpkin_core::propagation::DomainEvents;
use pumpkin_core::propagation::Domains;
use pumpkin_core::propagation::EnqueueDecision;
use pumpkin_core::propagation::EventsToRegister;
use pumpkin_core::propagation::ExplanationContext;
use pumpkin_core::propagation::LazyExplanation;
use pumpkin_core::propagation::LocalId;
use pumpkin_core::propagation::NotificationContext;
use pumpkin_core::propagation::OpaqueDomainEvent;
use pumpkin_core::propagation::Priority;
use pumpkin_core::propagation::PropagationContext;
use pumpkin_core::propagation::Propagator;
use pumpkin_core::propagation::PropagatorConstructor;
use pumpkin_core::propagation::PropagatorConstructorContext;
use pumpkin_core::propagation::PropagatorSpec;
use pumpkin_core::propagation::ReadDomains;
use pumpkin_core::propagation::RuntimeCheckers;
use pumpkin_core::propagation::TrailedInteger;
use pumpkin_core::state::PropagationStatusCP;
use pumpkin_core::state::PropagatorConflict;
use pumpkin_core::variables::IntegerVariable;

declare_inference_label!(LinearBounds);

/// The [`PropagatorConstructor`] for the [`LinearLessOrEqualPropagator`].
#[derive(Clone, Debug)]
pub struct LinearLessOrEqualPropagatorArgs<Var> {
    pub x: Box<[Var]>,
    pub c: i32,
    pub constraint_tag: ConstraintTag,
}

impl<Var> PropagatorConstructor for LinearLessOrEqualPropagatorArgs<Var>
where
    Var: IntegerVariable + 'static,
{
    type PropagatorImpl = LinearLessOrEqualPropagator<Var>;

    fn create(
        self,
        mut context: PropagatorConstructorContext,
    ) -> PropagatorSpec<Self::PropagatorImpl> {
        let LinearLessOrEqualPropagatorArgs {
            x,
            c,
            constraint_tag,
        } = self;

        let mut lower_bound_left_hand_side = 0_i64;
        let mut current_bounds = vec![];

        let mut registration = EventsToRegister::builder();
        for (i, x_i) in x.iter().enumerate() {
            registration =
                registration.add(x_i, DomainEvents::LOWER_BOUND, LocalId::from(i as u32));
            lower_bound_left_hand_side += context.lower_bound(x_i) as i64;
            current_bounds.push(context.new_trailed_integer(context.lower_bound(x_i) as i64));
        }

        let lower_bound_left_hand_side = context.new_trailed_integer(lower_bound_left_hand_side);

        let mut checkers = RuntimeCheckers::builder();
        let inference_code = checkers.add_inference_checker(
            constraint_tag,
            LinearBounds,
            LinearLessOrEqualInferenceChecker::new(x.clone(), c),
        );

        let propagator = LinearLessOrEqualPropagator {
            x,
            c,
            lower_bound_left_hand_side,
            current_bounds: current_bounds.into(),
            inference_code,
            reason_buffer: Vec::default(),
        };

        PropagatorSpec {
            registration: registration.build(),
            checkers: checkers.build(),
            propagator,
        }
    }
}

/// Propagator for the constraint `\sum x_i <= c`.
#[derive(Clone, Debug)]
pub struct LinearLessOrEqualPropagator<Var> {
    x: Box<[Var]>,
    c: i32,

    /// The lower bound of the sum of the left-hand side. This is incremental state.
    lower_bound_left_hand_side: TrailedInteger,
    /// The value at index `i` is the bound for `x[i]`.
    current_bounds: Box<[TrailedInteger]>,
    /// A buffer for storing the reason for a propagation.
    reason_buffer: Vec<Predicate>,

    inference_code: InferenceCode,
}

impl<Var> LinearLessOrEqualPropagator<Var>
where
    Var: IntegerVariable,
{
    fn create_conflict(&self, context: Domains) -> PropagatorConflict {
        PropagatorConflict {
            conjunction: self
                .x
                .iter()
                .map(|var| predicate![var >= context.lower_bound(var)])
                .collect(),
            inference_code: self.inference_code.clone(),
        }
    }
}

impl<Var: 'static> Propagator for LinearLessOrEqualPropagator<Var>
where
    Var: IntegerVariable,
{
    fn detect_inconsistency(&self, domains: Domains) -> Option<PropagatorConflict> {
        if (self.c as i64) < domains.read_trailed_integer(self.lower_bound_left_hand_side) {
            Some(self.create_conflict(domains))
        } else {
            None
        }
    }

    fn notify(
        &mut self,
        mut context: NotificationContext,
        local_id: LocalId,
        _event: OpaqueDomainEvent,
    ) -> EnqueueDecision {
        let index = local_id.unpack() as usize;
        let x_i = &self.x[index];

        let old_bound = context.read_trailed_integer(self.current_bounds[index]);
        let new_bound = context.lower_bound(x_i) as i64;

        pumpkin_assert_simple!(
            old_bound < new_bound,
            "propagator should only be triggered when lower bounds are tightened, old_bound={old_bound}, new_bound={new_bound}"
        );

        context.write_trailed_integer(
            self.lower_bound_left_hand_side,
            context.read_trailed_integer(self.lower_bound_left_hand_side) + (new_bound - old_bound),
        );
        context.write_trailed_integer(self.current_bounds[index], new_bound);

        EnqueueDecision::Enqueue
    }

    fn priority(&self) -> Priority {
        Priority::High
    }

    fn name(&self) -> &str {
        "LinearLeq"
    }

    fn lazy_explanation(&mut self, code: u64, context: ExplanationContext) -> LazyExplanation<'_> {
        let i = code as usize;

        self.reason_buffer.clear();

        self.reason_buffer
            .extend(self.x.iter().enumerate().filter_map(|(j, x_j)| {
                if j != i {
                    Some(predicate![
                        x_j >= context
                            .lower_bound_at_trail_position(x_j, context.get_trail_position())
                    ])
                } else {
                    None
                }
            }));

        LazyExplanation {
            predicates: self.reason_buffer.as_slice(),
            inference_code: self.inference_code.clone(),
        }
    }

    fn propagate(&mut self, mut context: PropagationContext) -> PropagationStatusCP {
        if let Some(conflict) = self.detect_inconsistency(context.domains()) {
            return Err(conflict.into());
        }

        let lower_bound_left_hand_side = match TryInto::<i32>::try_into(
            context.read_trailed_integer(self.lower_bound_left_hand_side),
        ) {
            Ok(bound) => bound,
            Err(_)
                if context
                    .read_trailed_integer(self.lower_bound_left_hand_side)
                    .is_positive() =>
            {
                // We cannot fit the `lower_bound_left_hand_side` into an i32 due to an
                // overflow (hence the check that the lower-bound on the left-hand side is
                // positive)
                //
                // This means that the lower-bounds of the current variables will always be
                // higher than the right-hand side (with a maximum value of i32). We thus
                // return a conflict
                return Err(self.create_conflict(context.domains()).into());
            }
            Err(_) => {
                // We cannot fit the `lower_bound_left_hand_side` into an i32 due to an
                // underflow
                //
                // This means that the constraint is always satisfied
                return Ok(());
            }
        };

        for (i, x_i) in self.x.iter().enumerate() {
            let bound = self.c - (lower_bound_left_hand_side - context.lower_bound(x_i));

            if context.upper_bound(x_i) > bound {
                context.post(predicate![x_i <= bound], i)?;
            }
        }

        Ok(())
    }

    fn propagate_from_scratch(&self, mut context: PropagationContext) -> PropagationStatusCP {
        let lower_bound_left_hand_side = self
            .x
            .iter()
            .map(|var| context.lower_bound(var) as i64)
            .sum::<i64>();

        let lower_bound_left_hand_side = match TryInto::<i32>::try_into(lower_bound_left_hand_side)
        {
            Ok(bound) => bound,
            Err(_)
                if context
                    .read_trailed_integer(self.lower_bound_left_hand_side)
                    .is_positive() =>
            {
                // We cannot fit the `lower_bound_left_hand_side` into an i32 due to an
                // overflow (hence the check that the lower-bound on the left-hand side is
                // positive)
                //
                // This means that the lower-bounds of the current variables will always be
                // higher than the right-hand side (with a maximum value of i32). We thus
                // return a conflict
                return Err(self.create_conflict(context.domains()).into());
            }
            Err(_) => {
                // We cannot fit the `lower_bound_left_hand_side` into an i32 due to an
                // underflow
                //
                // This means that the constraint is always satisfied
                return Ok(());
            }
        };

        for (i, x_i) in self.x.iter().enumerate() {
            let bound = self.c - (lower_bound_left_hand_side - context.lower_bound(x_i));

            if context.upper_bound(x_i) > bound {
                let reason: PropositionalConjunction = self
                    .x
                    .iter()
                    .enumerate()
                    .filter_map(|(j, x_j)| {
                        if j != i {
                            Some(predicate![x_j >= context.lower_bound(x_j)])
                        } else {
                            None
                        }
                    })
                    .collect();

                context.post(predicate![x_i <= bound], (reason, &self.inference_code))?;
            }
        }

        Ok(())
    }
}

#[derive(Debug, Clone)]
pub struct LinearLessOrEqualInferenceChecker<Var> {
    terms: Box<[Var]>,
    bound: i32,
}

impl<Var> LinearLessOrEqualInferenceChecker<Var> {
    pub fn new(terms: Box<[Var]>, bound: i32) -> Self {
        LinearLessOrEqualInferenceChecker { terms, bound }
    }
}

impl<Var, Atomic> InferenceChecker<Atomic> for LinearLessOrEqualInferenceChecker<Var>
where
    Var: CheckerVariable<Atomic>,
    Atomic: AtomicConstraint,
{
    fn check(
        &self,
        variable_state: VariableState<Atomic>,
        _: &[Atomic],
        _: Option<&Atomic>,
    ) -> bool {
        // Next, we evaluate the linear inequality. The lower bound of the
        // left-hand side must exceed the bound in the constraint. Note that the accumulator is an
        // IntExt, and if the lower bound of one of the terms is -infty, then the left-hand side
        // will be -infty regardless of the other terms.
        let left_hand_side: IntExt<i64> = self
            .terms
            .iter()
            .map(|variable| variable.induced_lower_bound(&variable_state).into())
            .sum();

        left_hand_side > i64::from(self.bound)
    }
}

#[cfg(test)]
mod tests {
    use pumpkin_core::conjunction;
    use pumpkin_core::predicate;
    use pumpkin_core::predicates::Predicate;
    use pumpkin_core::predicates::PropositionalConjunction;
    use pumpkin_core::propagation::CurrentNogood;
    use pumpkin_core::state::State;

    use super::*;
    use crate::StateExt;

    #[test]
    fn test_bounds_are_propagated() {
        let mut state = State::default();
        let x = state.new_interval_variable(1, 5, None);
        let y = state.new_interval_variable(0, 10, None);

        let constraint_tag = state.new_constraint_tag();

        let _ = state.add_propagator(LinearLessOrEqualPropagatorArgs {
            x: [x, y].into(),
            c: 7,
            constraint_tag,
        });
        state.propagate_to_fixed_point().expect("no empty domains");

        state.assert_bounds(x, 1, 5);
        state.assert_bounds(y, 0, 6);
    }

    #[test]
    fn test_explanations() {
        let mut state = State::default();
        let x = state.new_interval_variable(1, 5, None);
        let y = state.new_interval_variable(0, 10, None);
        let constraint_tag = state.new_constraint_tag();

        let _ = state.add_propagator(LinearLessOrEqualPropagatorArgs {
            x: [x, y].into(),
            c: 7,
            constraint_tag,
        });
        state.propagate_to_fixed_point().expect("no empty domains");

        let mut reason_buffer: Vec<Predicate> = vec![];
        let _ = state.get_propagation_reason(
            predicate![y <= 6],
            &mut reason_buffer,
            CurrentNogood::empty(),
        );
        let reason: PropositionalConjunction = reason_buffer.into();

        assert_eq!(conjunction!([x >= 1]), reason);
    }

    #[test]
    fn overflow_leads_to_conflict() {
        let mut state = State::default();

        let x = state.new_interval_variable(i32::MAX, i32::MAX, None);
        let y = state.new_interval_variable(1, 1, None);
        let constraint_tag = state.new_constraint_tag();

        let _ = state.add_propagator(LinearLessOrEqualPropagatorArgs {
            x: [x, y].into(),
            c: i32::MAX,
            constraint_tag,
        });
        let _ = state
            .propagate_to_fixed_point()
            .expect_err("Expected overflow to be detected");
    }

    #[test]
    fn underflow_leads_to_no_propagation() {
        let mut state = State::default();

        let x = state.new_interval_variable(i32::MIN, i32::MIN, None);
        let y = state.new_interval_variable(-1, -1, None);
        let constraint_tag = state.new_constraint_tag();

        let _ = state.add_propagator(LinearLessOrEqualPropagatorArgs {
            x: [x, y].into(),
            c: i32::MIN,
            constraint_tag,
        });
        state
            .propagate_to_fixed_point()
            .expect("Expected no error to be detected");
    }
}