Avoid running boolean extraction twice. (#3200)

In a future PR (#3208) we will try to run boolean extraction on any
newly changed constraint. This can have the risk that we perform boolean
extraction on a constraint where we already have run boolean extraction.
This PR stores a map of substitutions we performed to avoid introducing
a new boolean variable that then just turns out to be equivalent to an
already existing one.

Closes 3202

---------

Co-authored-by: schaeff <thibaut@powdrlabs.com>
This commit is contained in:
chriseth
2025-08-21 15:56:15 +02:00
committed by GitHub
parent 81ae444b4c
commit b5aa6759b4
4 changed files with 294 additions and 91 deletions

View File

@@ -26,6 +26,24 @@ pub fn apply_substitutions<T: RuntimeConstant + Substitutable<V>, V: Hash + Eq +
indexed_constraint_system.into()
}
/// Applies multiple substitutions to all expressions in a sequence of expressions.
pub fn apply_substitutions_to_expressions<
T: RuntimeConstant + Substitutable<V>,
V: Hash + Eq + Clone + Ord,
>(
expressions: impl IntoIterator<Item = GroupedExpression<T, V>>,
substitutions: impl IntoIterator<Item = (V, GroupedExpression<T, V>)>,
) -> Vec<GroupedExpression<T, V>> {
apply_substitutions(
ConstraintSystem {
algebraic_constraints: expressions.into_iter().collect(),
bus_interactions: Vec::new(),
},
substitutions,
)
.algebraic_constraints
}
/// Structure on top of a [`ConstraintSystem`] that stores indices
/// to more efficiently update the constraints.
#[derive(Clone)]

View File

@@ -8,7 +8,7 @@ use crate::grouped_expression::{GroupedExpression, RangeConstraintProvider};
use crate::indexed_constraint_system::IndexedConstraintSystemWithQueue;
use crate::range_constraint::RangeConstraint;
use crate::runtime_constant::{ReferencedSymbols, RuntimeConstant, Substitutable};
use crate::solver::boolean_extractor::try_extract_boolean;
use crate::solver::boolean_extractor::BooleanExtractor;
use crate::solver::linearizer::Linearizer;
use crate::solver::var_transformation::Variable;
use crate::solver::{exhaustive_search, quadratic_equivalences, Error, Solver, VariableAssignment};
@@ -45,6 +45,8 @@ pub struct BaseSolver<T: RuntimeConstant, V, BusInterHandler, VarDisp> {
equivalent_expressions_cache: HashMap<GroupedExpression<T, V>, Vec<GroupedExpression<T, V>>>,
/// A dispenser for fresh variables.
var_dispenser: VarDisp,
/// The boolean extraction component.
boolean_extractor: BooleanExtractor<T, V>,
/// The linearizing component.
linearizer: Linearizer<T, V>,
}
@@ -93,6 +95,7 @@ impl<T: RuntimeConstant, V, B, VD: Default> BaseSolver<T, V, B, VD> {
assignments_to_return: Default::default(),
equivalent_expressions_cache: Default::default(),
var_dispenser: Default::default(),
boolean_extractor: Default::default(),
linearizer: Default::default(),
bus_interaction_handler,
}
@@ -133,12 +136,13 @@ where
self.loop_until_no_progress()?;
let assignments = std::mem::take(&mut self.assignments_to_return);
// Apply the deduced assignments to the substitutions we performed
// while linearizing.
// while linearizing and boolean extracting.
// We assume that the user of the solver applies the assignments to
// their expressions and thus "incoming" expressions used in the functions
// `range_constraint_for_expression` and `are_expressions_known_to_be_different`
// will have the assignments applied.
self.linearizer.apply_assignments(&assignments);
self.boolean_extractor.apply_assignments(&assignments);
Ok(assignments)
}
@@ -151,7 +155,11 @@ where
let constraints = constraints
.into_iter()
.filter(|c| !c.is_zero())
.flat_map(|constr| self.extract_boolean(constr))
.flat_map(|constr| {
self.try_extract_boolean(&constr)
.into_iter()
.chain(std::iter::once(constr))
})
// needed because of unique access to the var dispenser / self.
.collect_vec()
.into_iter()
@@ -257,19 +265,20 @@ where
+ ExpressionConvertible<T::FieldType, V>
+ Substitutable<V>,
{
/// Performs boolean extraction on `constr`, i.e. tries to turn quadratic constraints into affine constraints
/// Tries to performs boolean extraction on `constr`, i.e. tries to turn quadratic constraints into affine constraints
/// by introducing new boolean variables.
/// This function will always return the original constraint as well as any extracted constraints.
fn extract_boolean(
fn try_extract_boolean(
&mut self,
constr: GroupedExpression<T, V>,
) -> impl Iterator<Item = GroupedExpression<T, V>> {
let extracted = try_extract_boolean(&constr, &mut || {
let v = self.var_dispenser.next_boolean();
self.add_range_constraint(&v, RangeConstraint::from_mask(1));
v
});
std::iter::once(constr).chain(extracted)
constr: &GroupedExpression<T, V>,
) -> Option<GroupedExpression<T, V>> {
let (constr, var) = self
.boolean_extractor
.try_extract_boolean(constr, || self.var_dispenser.next_boolean())?;
if let Some(var) = var {
// If we created a boolean variable, we constrain it to be boolean.
self.add_range_constraint(&var, RangeConstraint::from_mask(1));
}
Some(constr)
}
/// Performs linearization of `constr`, i.e. replaces all non-affine sub-components of the constraint

View File

@@ -1,76 +1,170 @@
use std::hash::Hash;
use std::{cmp::min, collections::HashMap, hash::Hash};
use crate::{grouped_expression::GroupedExpression, runtime_constant::RuntimeConstant};
use itertools::Itertools;
use powdr_number::{FieldElement, LargeInt};
/// Tries to simplify a quadratic constraint by transforming it into an affine
/// constraint that makes use of a new boolean variable.
/// NOTE: The boolean constraint is not part of the output.
///
/// For example `(a + b) * (a + b + 10) = 0` can be transformed into
/// `a + b + z * 10 = 0`, where `z` is a new boolean variable.
///
/// @param constraint The quadratic constraint to transform.
/// @param var_dispenser A function that returns a new variable that is assumed to be boolean-constrained.
/// It will only be called if the transformation is performed.
pub fn try_extract_boolean<T: RuntimeConstant, V: Ord + Clone + Hash + Eq>(
constraint: &GroupedExpression<T, V>,
mut var_dispenser: impl FnMut() -> V,
) -> Option<GroupedExpression<T, V>> {
let (left, right) = constraint.try_as_single_product()?;
// We want to check if `left` and `right` differ by a constant offset.
// Since multiplying the whole constraint by a non-zero constant does
// not change the constraint, we also transform `left` by a constant
// (non-zero) factor.
// So we are looking for a offset `c` and a non-zero constant factor `f`
// such that `f * left = right + c`.
// Then we can write the original constraint `left * right = 0` as
// `(right + c) * right = 0` (we can just ignore `f`).
// This is in turn equivalent to `right + z * c = 0`, where `z` is
// a new boolean variable.
use crate::{
grouped_expression::GroupedExpression,
indexed_constraint_system::apply_substitutions_to_expressions,
runtime_constant::{RuntimeConstant, Substitutable},
solver::VariableAssignment,
};
// For example, if the constraint was `(2 * a + 2 * b) * (a + b + 10) = 0`, we would
// set `factor = 1 / 2`, such that `left * factor - right` is a constant.
pub struct BooleanExtractor<T, V> {
/// If (expr, Some(z)) is in the map, it means that
/// we have transformed a constraint `left * right = 0` into
/// `right + z * offset = 0`, where `z` is a new boolean variable
/// and `expr = -right / offset = z`.
///
/// If (expr, None) is in the map, it means that
/// we have transformed a constraint `right * right = 0` into
/// `right = 0`, which is a special case where we do not need
/// a new boolean variable.
substitutions: HashMap<GroupedExpression<T, V>, Option<V>>,
}
// First, try to find a good factor so that `left` and `right`
// likely cancel out except for a constant. As a good guess,
// we try to match the coefficient of the first variable.
let factor = match (left.components().1.next(), right.components().1.next()) {
(Some((left_var, left_coeff)), Some((right_var, right_coeff))) if left_var == right_var => {
right_coeff.field_div(left_coeff)
impl<T, V> Default for BooleanExtractor<T, V> {
fn default() -> Self {
Self {
substitutions: HashMap::new(),
}
_ => T::one(),
};
}
}
// `constr = 0` is equivalent to `left * right = 0`
let offset = &(left.clone() * &factor) - right;
// We only do the transformation if `offset` is known, because
// otherwise the constraint stays quadratic.
let offset = offset.try_to_known()?;
// We know that `offset + right = left` and thus
// `constr = 0` is equivalent to `right * (right + offset) = 0`
// which is equivalent to `right + z * offset = 0` for a new
// boolean variable `z`.
impl<T: RuntimeConstant + Hash, V: Ord + Clone + Hash + Eq> BooleanExtractor<T, V> {
/// Tries to simplify a quadratic constraint by transforming it into an affine
/// constraint that makes use of a new boolean variable.
/// NOTE: The boolean constraint is not part of the output.
///
/// Returns the new constraint and the new variable if required.
///
/// If the same simplification has been performed before, it will
/// return None (in particular, it will not request a new variable).
///
/// For example `(a + b) * (a + b + 10) = 0` can be transformed into
/// `a + b + z * 10 = 0`, where `z` is a new boolean variable.
///
/// @param constraint The quadratic constraint to transform.
/// @param var_dispenser A function that returns a new variable that is assumed to be boolean-constrained.
/// It will only be called if the transformation is performed.
pub fn try_extract_boolean(
&mut self,
constraint: &GroupedExpression<T, V>,
mut var_dispenser: impl FnMut() -> V,
) -> Option<(GroupedExpression<T, V>, Option<V>)> {
let (left, right) = constraint.try_as_single_product()?;
// We want to check if `left` and `right` differ by a constant offset.
// Since multiplying the whole constraint by a non-zero constant does
// not change the constraint, we also transform `left` by a constant
// (non-zero) factor.
// So we are looking for an offset `c` and a non-zero constant factor `f`
// such that `f * left = right + c`.
// Then we can write the original constraint `left * right = 0` as
// `(right + c) * right = 0` (we can just ignore `f`).
// This is in turn equivalent to `right + z * c = 0`, where `z` is
// a new boolean variable.
if offset.is_zero() {
// In this special case, we do not need a new variable.
Some(right.clone())
} else if (right.clone() * -T::one().field_div(offset))
.try_to_simple_unknown()
.is_some()
{
// In this case we don't gain anything because the new variable `z` will just
// be equivalent to the single variable in `right`.
return None;
} else {
let z = var_dispenser();
// For example, if the constraint was `(2 * a + 2 * b) * (a + b + 10) = 0`, we would
// set `factor = 1 / 2`, such that `left * factor - right` is a constant.
// We return `right + z * offset == 0`, which is equivalent to the original constraint.
Some(right + &(GroupedExpression::from_unknown_variable(z) * offset))
// First, try to find a good factor so that `left` and `right`
// likely cancel out except for a constant. As a good guess,
// we try to match the coefficient of the first variable.
let factor = match (left.components().1.next(), right.components().1.next()) {
(Some((left_var, left_coeff)), Some((right_var, right_coeff)))
if left_var == right_var =>
{
right_coeff.field_div(left_coeff)
}
_ => T::one(),
};
// `constr = 0` is equivalent to `left * right = 0`
let offset = &(left.clone() * &factor) - right;
// We only do the transformation if `offset` is known, because
// otherwise the constraint stays quadratic.
let offset = offset.try_to_known()?;
// We know that `offset + right = left` and thus
// `constr = 0` is equivalent to `right * (right + offset) = 0`
// which is equivalent to `right + z * offset = 0` for a new
// boolean variable `z`.
if offset.is_zero() {
// In this special case, we do not need a new variable.
if self.substitutions.contains_key(right) {
None
} else {
self.substitutions.insert(right.clone(), None);
Some((right.clone(), None))
}
} else {
// We can substitute the initial constraint using a new boolean variable `z`
// either by
// `0 = right + z * offset`
// or by
// `0 = right + (1 - z) * offset = right + offset - z * offset`,
// which is equivalent to
// `0 = -right - offset + z * offset`.
// We use the one that has a smaller constant offset in the resulting expression.
let expr = [
right.clone(),
-right - GroupedExpression::from_runtime_constant(offset.clone()),
]
.into_iter()
.min_by_key(|e| {
// Return the abs of the constant offset, or None on larger fields.
e.components().2.try_to_number().and_then(try_to_abs_u64)
})
.unwrap();
let key = -&expr * T::one().field_div(offset);
if self.substitutions.contains_key(&key) {
// We have already performed this transformation before.
return None;
}
if key.try_to_simple_unknown().is_some() {
// In this case we don't gain anything because the new variable `z` will just
// be equivalent to the single variable in `right`.
None
} else {
let z = var_dispenser();
self.substitutions.insert(key, Some(z.clone()));
// We return `expr + z * offset == 0`, which is equivalent to the original constraint.
Some((
expr + (GroupedExpression::from_unknown_variable(z.clone()) * offset),
Some(z),
))
}
}
}
}
fn try_to_abs_u64<T: FieldElement>(x: T) -> Option<u64> {
let modulus = T::modulus().try_into_u64()?;
let x = x.to_integer().try_into_u64()?;
Some(min(x, modulus - x))
}
impl<T: RuntimeConstant + Substitutable<V> + Hash, V: Clone + Eq + Ord + Hash>
BooleanExtractor<T, V>
{
/// Applies the assignments to the stored substitutions.
pub fn apply_assignments(&mut self, assignments: &[VariableAssignment<T, V>]) {
if assignments.is_empty() {
return;
}
let (exprs, vars): (Vec<_>, Vec<_>) = self.substitutions.drain().unzip();
let exprs = apply_substitutions_to_expressions(exprs, assignments.iter().cloned());
self.substitutions = exprs.into_iter().zip_eq(vars).collect();
}
}
#[cfg(test)]
mod tests {
use crate::test_utils::{constant, var};
use super::*;
@@ -79,31 +173,121 @@ mod tests {
fn test_extract_boolean() {
let mut var_dispenser = || "z";
let expr = (var("a") + var("b")) * (var("a") + var("b") + constant(10));
let result = try_extract_boolean(&expr, &mut var_dispenser);
let mut extractor: BooleanExtractor<_, _> = Default::default();
let result = extractor.try_extract_boolean(&expr, &mut var_dispenser);
assert!(result.is_some());
let result = result.unwrap();
assert_eq!(result.to_string(), "a + b - 10 * z + 10");
let (result, z) = result.unwrap();
assert_eq!(result.to_string(), "-(a + b + 10 * z)");
assert_eq!(z, Some("z"));
}
#[test]
fn test_extract_boolean_square() {
let mut var_dispenser = || "z";
let expr = (var("a") + var("b")) * (var("a") + var("b"));
let result = try_extract_boolean(&expr, &mut var_dispenser);
let mut extractor: BooleanExtractor<_, _> = Default::default();
let result = extractor.try_extract_boolean(&expr, &mut var_dispenser);
assert!(result.is_some());
let result = result.unwrap();
let (result, z) = result.unwrap();
assert_eq!(result.to_string(), "a + b");
assert_eq!(z, None);
}
#[test]
fn test_extract_boolean_useless() {
let mut var_dispenser = || "z";
let expr = (var("a") - constant(1)) * (var("a"));
let result = try_extract_boolean(&expr, &mut var_dispenser);
let mut extractor: BooleanExtractor<_, _> = Default::default();
let result = extractor.try_extract_boolean(&expr, &mut var_dispenser);
assert!(result.is_none());
let expr = (constant(2) * var("a") - constant(2)) * (constant(2) * var("a"));
let result = try_extract_boolean(&expr, &mut var_dispenser);
let result = extractor.try_extract_boolean(&expr, &mut var_dispenser);
assert!(result.is_none());
}
#[test]
fn do_not_extract_twice() {
let mut var_dispenser = || "z";
let expr = (var("a") + var("b")) * (var("a") + var("b") + constant(10));
let mut extractor: BooleanExtractor<_, _> = Default::default();
let result = extractor.try_extract_boolean(&expr, &mut var_dispenser);
assert!(result.is_some());
let (result, z) = result.unwrap();
assert_eq!(result.to_string(), "-(a + b + 10 * z)");
assert_eq!(z, Some("z"));
assert!(extractor
.try_extract_boolean(&expr, &mut var_dispenser)
.is_none());
// left and right swapped
assert!(extractor
.try_extract_boolean(
&(var("a") + var("b") + constant(10) * (var("a") + var("b"))),
&mut var_dispenser
)
.is_none());
let expr2 = (constant(2) * (var("a") + var("b"))) * (var("a") + var("b") + constant(10));
assert!(extractor
.try_extract_boolean(&expr2, &mut var_dispenser)
.is_none());
let expr3 = (var("a") + var("b")) * (constant(2) * (var("a") + var("b") + constant(10)));
assert!(extractor
.try_extract_boolean(&expr3, &mut var_dispenser)
.is_none());
// This is different because the effective constant is different.
let expr4 = (var("a") + var("b")) * (constant(2) * (var("a") + var("b") + constant(20)));
assert_eq!(
extractor
.try_extract_boolean(&expr4, &mut var_dispenser)
.unwrap()
.0
.to_string(),
"-(2 * a + 2 * b + 40 * z)"
);
}
#[test]
fn do_not_extract_squares_twice() {
let mut var_dispenser = || "z";
let expr = (var("a") + var("b")) * (var("a") + var("b"));
let mut extractor: BooleanExtractor<_, _> = Default::default();
let result = extractor.try_extract_boolean(&expr, &mut var_dispenser);
assert!(result.is_some());
let (result, z) = result.unwrap();
assert_eq!(result.to_string(), "a + b");
assert_eq!(z, None);
let result = extractor.try_extract_boolean(&expr, &mut var_dispenser);
assert!(result.is_none());
}
#[test]
fn apply_assignments() {
let mut counter = 0;
let vars = (0..10).map(|i| format!("z_{i}")).collect_vec();
let mut var_dispenser = || {
counter += 1;
vars[counter - 1].as_str()
};
let expr =
(var("a") + var("b") + var("k")) * (var("a") + var("b") + var("k") - constant(2));
let mut extractor: BooleanExtractor<_, _> = Default::default();
let result = extractor.try_extract_boolean(&expr, &mut var_dispenser);
assert!(result.is_some());
let (result, z) = result.unwrap();
assert_eq!(result.to_string(), "-(a + b + k - 2 * z_0)");
assert_eq!(z, Some("z_0"));
extractor.apply_assignments(&[("k", -constant(9))]);
let expr2 =
(var("a") + var("b") - constant(9)) * (var("a") + var("b") - constant(9) - constant(2));
let result = extractor.try_extract_boolean(&expr2, &mut var_dispenser);
assert!(result.is_none());
}
}

View File

@@ -3,8 +3,7 @@ use std::hash::Hash;
use itertools::Itertools;
use crate::constraint_system::ConstraintSystem;
use crate::indexed_constraint_system::apply_substitutions;
use crate::indexed_constraint_system::apply_substitutions_to_expressions;
use crate::runtime_constant::Substitutable;
use crate::solver::VariableAssignment;
use crate::{grouped_expression::GroupedExpression, runtime_constant::RuntimeConstant};
@@ -158,14 +157,7 @@ impl<T: RuntimeConstant + Substitutable<V> + Hash, V: Clone + Eq + Ord + Hash> L
return;
}
let (exprs, vars): (Vec<_>, Vec<_>) = self.substitutions.drain().unzip();
let exprs = apply_substitutions(
ConstraintSystem {
algebraic_constraints: exprs,
bus_interactions: vec![],
},
assignments.iter().cloned(),
)
.algebraic_constraints;
let exprs = apply_substitutions_to_expressions(exprs, assignments.iter().cloned());
self.substitutions = exprs
.into_iter()
.zip_eq(vars)