Add detection for linear views in the FlatZinc

huub-solver · May 14, 2024 · b4ebd52 · b4ebd52
1 parent a26289c
commit b4ebd52
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 14 deletions.
diff --git a/crates/huub/src/model/flatzinc.rs b/crates/huub/src/model/flatzinc.rs
@@ -9,7 +9,7 @@ use super::{
 	reformulate::ReformulationError,
 	ModelView,
 };
-use crate::{Constraint, IntVal, Model, Solver, SolverView};
+use crate::{Constraint, IntVal, Model, NonZeroIntVal, Solver, SolverView};
 
 impl Model {
 	pub fn from_fzn<S: Ord + Deref<Target = str> + Clone + Debug>(
@@ -24,27 +24,60 @@ impl Model {
 			.iter()
 			.map(|c| {
 				match (c.id.deref(), c.defines.as_ref()) {
-					("bool_not", Some(l)) => {
-						if let [a, b] = c.args.as_slice() {
-							match (a, b) {
-								(Argument::Literal(Literal::Identifier(x)), b)
-								| (b, Argument::Literal(Literal::Identifier(x)))
-									if x == l =>
-								{
-									let b = arg_bool(fzn, &mut prb, &mut map, b)?;
-									let _ = map.insert(l.clone(), (!b).into());
-									return Ok(true);
-								}
-								_ => {}
+					("bool_not", Some(l)) => match c.args.as_slice() {
+						[Argument::Literal(Literal::Identifier(x)), b]
+						| [b, Argument::Literal(Literal::Identifier(x))]
+							if x == l =>
+						{
+							let b = arg_bool(fzn, &mut prb, &mut map, b)?;
+							let _ = map.insert(l.clone(), (!b).into());
+							return Ok(true);
+						}
+						_ => {}
+					},
+					("int_lin_eq", Some(l))
+						if c.args
+							.get(1)
+							.map(|v| arg_has_length(fzn, v, 2))
+							.unwrap_or(false) =>
+					{
+						let [coeff, vars, sum] = c.args.as_slice() else {
+							return Ok(false);
+						};
+						let coeff = arg_array(fzn, coeff)?;
+						let vars = arg_array(fzn, vars)?;
+						let (c, (cy, vy)) = match vars.as_slice() {
+							[Literal::Identifier(v), y] if v == l => {
+								(par_int(fzn, &coeff[0])?, (par_int(fzn, &coeff[1])?, y))
 							}
+							[y, Literal::Identifier(v)] if v == l => {
+								(par_int(fzn, &coeff[1])?, (par_int(fzn, &coeff[0])?, y))
+							}
+							_ => return Ok(false),
+						};
+						let sum = arg_par_int(fzn, sum)?;
+						// c * l + cy * y = sum === l = (sum - cy * y) / c
+						if cy % c != 0 || sum % c != 0 {
+							return Ok(false);
 						}
+						let offset = sum / c;
+						let view = if let Some(scale) = NonZeroIntVal::new(-cy / c) {
+							let y = lit_int(fzn, &mut prb, &mut map, vy)?;
+							y * scale + offset
+						} else {
+							IntView::Const(offset)
+						};
+						let _ = map.insert(l.clone(), view.into());
+						return Ok(true);
 					}
 					_ => {}
 				}
 				Ok(false)
 			})
 			.collect::<Result<Vec<_>, FlatZincError>>()?;
 
+		eprintln!("{:?}", map);
+
 		// Traditional relational constraints
 		for (i, c) in fzn.constraints.iter().enumerate() {
 			if processed[i] {
@@ -327,6 +360,24 @@ fn arg_array<'a, S: Ord + Deref<Target = str> + Clone + Debug>(
 	}
 }
 
+fn arg_has_length<'a, S: Ord + Deref<Target = str> + Clone + Debug>(
+	fzn: &'a FlatZinc<S>,
+	arg: &'a Argument<S>,
+	len: usize,
+) -> bool {
+	match arg {
+		Argument::Array(x) => x.len() == len,
+		Argument::Literal(Literal::Identifier(ident)) => {
+			if let Some(arr) = fzn.arrays.get(ident) {
+				arr.contents.len() == len
+			} else {
+				false
+			}
+		}
+		_ => false,
+	}
+}
+
 fn new_var<S: Ord + Deref<Target = str> + Clone + Debug>(
 	prb: &mut Model,
 	var: &Variable<S>,

diff --git a/crates/huub/src/model/int.rs b/crates/huub/src/model/int.rs
@@ -1,3 +1,5 @@
+use std::ops::{Add, Mul};
+
 use flatzinc_serde::RangeList;
 use pindakaas::{
 	solver::{PropagatorAccess, Solver as SolverTrait},
@@ -8,7 +10,7 @@ use super::reformulate::{ReifContext, VariableMap};
 use crate::{
 	helpers::linear_transform::LinearTransform,
 	solver::{view, SatSolver},
-	IntVal, Model, ReformulationError, Solver,
+	IntVal, Model, NonZeroIntVal, ReformulationError, Solver,
 };
 
 impl IntView {
@@ -51,6 +53,28 @@ pub enum IntView {
 	Linear(LinearTransform, IntVar),
 }
 
+impl Add<IntVal> for IntView {
+	type Output = Self;
+	fn add(self, rhs: IntVal) -> Self::Output {
+		match self {
+			Self::Var(x) => Self::Linear(LinearTransform::offset(rhs), x),
+			Self::Const(v) => Self::Const(v + rhs),
+			Self::Linear(t, x) => Self::Linear(t + rhs, x),
+		}
+	}
+}
+
+impl Mul<NonZeroIntVal> for IntView {
+	type Output = Self;
+	fn mul(self, rhs: NonZeroIntVal) -> Self::Output {
+		match self {
+			Self::Var(x) => Self::Linear(LinearTransform::scaled(rhs), x),
+			Self::Const(v) => Self::Const(v * rhs.get()),
+			Self::Linear(t, x) => Self::Linear(t * rhs, x),
+		}
+	}
+}
+
 impl Model {
 	pub(crate) fn get_int_lower_bound(&self, iv: &IntView) -> IntVal {
 		match *iv {