From 448debf109a1403dc82d4cd845a02bd3c0cf0dda Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 29 Jul 2025 09:26:23 +0200 Subject: [PATCH] Solver queue (#3092) Implements a queue over constraint system items such that items are re-queued when their variables are updated. --------- Co-authored-by: Georg Wiese --- constraint-solver/Cargo.toml | 1 + .../src/indexed_constraint_system.rs | 204 +++++++++++++++++- constraint-solver/src/solver.rs | 54 ++--- 3 files changed, 215 insertions(+), 44 deletions(-) diff --git a/constraint-solver/Cargo.toml b/constraint-solver/Cargo.toml index f2ab0e9fc..9e2fe4f02 100644 --- a/constraint-solver/Cargo.toml +++ b/constraint-solver/Cargo.toml @@ -15,6 +15,7 @@ num-traits = "0.2.15" derive_more = "0.99.17" auto_enums = "0.8.5" log = "0.4.27" +bitvec = "1.0.1" [dev-dependencies] pretty_assertions = "1.4.0" diff --git a/constraint-solver/src/indexed_constraint_system.rs b/constraint-solver/src/indexed_constraint_system.rs index 7347fa128..10820ee5d 100644 --- a/constraint-solver/src/indexed_constraint_system.rs +++ b/constraint-solver/src/indexed_constraint_system.rs @@ -1,5 +1,11 @@ -use std::{collections::HashMap, fmt::Display, hash::Hash}; +use std::{ + cmp, + collections::{HashMap, VecDeque}, + fmt::Display, + hash::Hash, +}; +use bitvec::vec::BitVec; use itertools::Itertools; use crate::{ @@ -30,12 +36,52 @@ pub struct IndexedConstraintSystem { variable_occurrences: HashMap>, } +/// Structure on top of [`IndexedConstraintSystem`] that +/// tracks changes to variables and how they may affect constraints. +/// +/// In particular, the assumption is that items in the constraint system +/// need to be "handled". Initially, all items need to be "handled" +/// and are put in a queue. Handling an item can cause an update to a variable, +/// which causes all constraints referencing that variable to be put back into the +/// queue. +#[derive(Clone, Default)] +pub struct IndexedConstraintSystemWithQueue { + constraint_system: IndexedConstraintSystem, + queue: ConstraintSystemQueue, +} + +/// A reference to an item in the constraint system. #[derive(Debug, Clone, Copy, PartialEq, Eq, Ord, PartialOrd, Hash)] enum ConstraintSystemItem { AlgebraicConstraint(usize), BusInteraction(usize), } +impl ConstraintSystemItem { + /// Returns an index that is unique across both algebraic constraints and bus interactions. + fn flat_id(&self) -> usize { + match self { + ConstraintSystemItem::AlgebraicConstraint(i) => 2 * i, + ConstraintSystemItem::BusInteraction(i) => 2 * i + 1, + } + } + + /// Turns this indexed-based item into a reference to the actual constraint. + fn to_constraint_ref<'a, T, V>( + self, + constraint_system: &'a ConstraintSystem, + ) -> ConstraintRef<'a, T, V> { + match self { + ConstraintSystemItem::AlgebraicConstraint(i) => { + ConstraintRef::AlgebraicConstraint(&constraint_system.algebraic_constraints[i]) + } + ConstraintSystemItem::BusInteraction(i) => { + ConstraintRef::BusInteraction(&constraint_system.bus_interactions[i]) + } + } + } +} + impl From> for IndexedConstraintSystem { @@ -207,14 +253,7 @@ impl IndexedConstraintSystem ConstraintRef::AlgebraicConstraint( - &self.constraint_system.algebraic_constraints[i], - ), - ConstraintSystemItem::BusInteraction(i) => { - ConstraintRef::BusInteraction(&self.constraint_system.bus_interactions[i]) - } - }) + .map(|&item| item.to_constraint_ref(&self.constraint_system)) } } @@ -350,6 +389,153 @@ impl Display } } +impl>> From + for IndexedConstraintSystemWithQueue +{ + fn from(constraint_system: C) -> Self { + let constraint_system = constraint_system.into(); + let queue = ConstraintSystemQueue::new(constraint_system.system()); + Self { + constraint_system, + queue, + } + } +} + +impl IndexedConstraintSystemWithQueue +where + T: RuntimeConstant + Substitutable, + V: Clone + Ord + Hash, +{ + /// Returns a reference to the underlying indexed constraint system. + pub fn system(&self) -> &IndexedConstraintSystem { + &self.constraint_system + } + + /// Removes the next item from the queue and returns it. + pub fn pop_front<'a>(&'a mut self) -> Option> { + self.queue + .pop_front() + .map(|item| item.to_constraint_ref(&self.constraint_system.constraint_system)) + } + + /// Notifies the system that a variable has been updated and causes all constraints + /// referencing that variable to be put back into the queue. + /// + /// Note that this function does not have to be called if the system is modified directly. + pub fn variable_updated(&mut self, variable: &V) { + if let Some(items) = self.constraint_system.variable_occurrences.get(variable) { + for item in items { + self.queue.push(*item); + } + } + } + + /// Substitutes a variable with a known value in the whole system. + /// This function also updates the queue accordingly. + pub fn substitute_by_unknown(&mut self, variable: &V, substitution: &GroupedExpression) { + self.constraint_system + .substitute_by_unknown(variable, substitution); + self.variable_updated(variable); + } + + pub fn add_algebraic_constraints( + &mut self, + constraints: impl IntoIterator>, + ) { + let initial_len = self + .constraint_system + .constraint_system + .algebraic_constraints + .len(); + self.constraint_system + .add_algebraic_constraints(constraints.into_iter().enumerate().map(|(i, c)| { + self.queue + .push(ConstraintSystemItem::AlgebraicConstraint(initial_len + i)); + c + })); + } + + pub fn retain_algebraic_constraints( + &mut self, + mut f: impl FnMut(&GroupedExpression) -> bool, + ) { + self.constraint_system.retain_algebraic_constraints(&mut f); + if !self.queue.queue.is_empty() { + // Removing items will destroy the indices, which is only safe if + // the queue is empty. Otherwise, we just put all items back into the queue. + self.queue = ConstraintSystemQueue::new(self.constraint_system.system()); + } + } + + pub fn retain_bus_interactions( + &mut self, + mut f: impl FnMut(&BusInteraction>) -> bool, + ) { + self.constraint_system.retain_bus_interactions(&mut f); + if !self.queue.queue.is_empty() { + // Removing items will destroy the indices, which is only safe if + // the queue is empty. Otherwise, we just put all items back into the queue. + self.queue = ConstraintSystemQueue::new(self.constraint_system.system()); + } + } +} + +impl Display + for IndexedConstraintSystemWithQueue +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.constraint_system) + } +} + +/// The actual queue used in `IndexedConstraintSystemWithQueue`. +/// +/// It keeps track that there are no duplicates in the queue by maintaining +/// a flat bitvector of items in the queue. +#[derive(Default, Clone)] +struct ConstraintSystemQueue { + queue: VecDeque, + in_queue: BitVec, +} + +impl ConstraintSystemQueue { + fn new(constraint_system: &ConstraintSystem) -> Self { + let num_algebraic = constraint_system.algebraic_constraints.len(); + let num_bus = constraint_system.bus_interactions.len(); + let queue = (0..num_algebraic) + .map(ConstraintSystemItem::AlgebraicConstraint) + .chain((0..num_bus).map(ConstraintSystemItem::BusInteraction)) + .collect::>() + .into(); + // The maximum value of `item.flat_id()` is `2 * max(num_algebraic, num_bus) + 1` + let mut in_queue = BitVec::repeat(false, 2 * cmp::max(num_algebraic, num_bus) + 2); + for item in &queue { + let item: &ConstraintSystemItem = item; + in_queue.set(item.flat_id(), true); + } + Self { queue, in_queue } + } + + fn push(&mut self, item: ConstraintSystemItem) { + if self.in_queue.len() <= item.flat_id() { + self.in_queue.resize(item.flat_id() + 1, false); + } + if !self.in_queue[item.flat_id()] { + self.queue.push_back(item); + self.in_queue.set(item.flat_id(), true); + } + } + + fn pop_front(&mut self) -> Option { + let item = self.queue.pop_front(); + if let Some(item) = &item { + self.in_queue.set(item.flat_id(), false); + } + item + } +} + #[cfg(test)] mod tests { use powdr_number::GoldilocksField; diff --git a/constraint-solver/src/solver.rs b/constraint-solver/src/solver.rs index 2451954a2..a90bc41e3 100644 --- a/constraint-solver/src/solver.rs +++ b/constraint-solver/src/solver.rs @@ -1,11 +1,11 @@ use powdr_number::{ExpressionConvertible, FieldElement}; use crate::constraint_system::{ - BusInteractionHandler, ConstraintSystem, DefaultBusInteractionHandler, + BusInteractionHandler, ConstraintRef, ConstraintSystem, DefaultBusInteractionHandler, }; use crate::effect::Effect; use crate::grouped_expression::GroupedExpression; -use crate::indexed_constraint_system::IndexedConstraintSystem; +use crate::indexed_constraint_system::IndexedConstraintSystemWithQueue; use crate::range_constraint::RangeConstraint; use crate::runtime_constant::{ReferencedSymbols, RuntimeConstant, Substitutable}; use crate::utils::known_variables; @@ -100,7 +100,7 @@ pub type VariableAssignment = (V, GroupedExpression); struct SolverImpl { /// The constraint system to solve. During the solving process, any expressions will /// be simplified as much as possible. - constraint_system: IndexedConstraintSystem, + constraint_system: IndexedConstraintSystemWithQueue, /// The handler for bus interactions. bus_interaction_handler: BusInterHandler, /// The currently known range constraints of the variables. @@ -121,7 +121,7 @@ impl, V: Ord + Clone + Hash + Eq + Dis ); SolverImpl { - constraint_system: IndexedConstraintSystem::from(constraint_system), + constraint_system: IndexedConstraintSystemWithQueue::from(constraint_system), range_constraints: Default::default(), bus_interaction_handler: Default::default(), assignments_to_return: Default::default(), @@ -177,7 +177,7 @@ where } fn retain_variables(&mut self, variables_to_keep: &HashSet) { - assert!(self.assignments_to_return.is_empty(),); + assert!(self.assignments_to_return.is_empty()); self.range_constraints .range_constraints .retain(|v, _| variables_to_keep.contains(v)); @@ -208,8 +208,6 @@ where let mut progress = false; // Try solving constraints in isolation. progress |= self.solve_in_isolation()?; - // Try inferring new information using bus interactions. - progress |= self.solve_bus_interactions()?; // Try to find equivalent variables using quadratic constraints. progress |= self.try_solve_quadratic_equivalences(); @@ -229,13 +227,17 @@ where /// Tries to make progress by solving each constraint in isolation. fn solve_in_isolation(&mut self) -> Result { let mut progress = false; - for i in 0..self.constraint_system.algebraic_constraints().len() { - // TODO: Improve efficiency by only running skipping constraints that - // have not received any updates since they were last processed. - let effects = self.constraint_system.algebraic_constraints()[i] - .solve(&self.range_constraints) - .map_err(Error::QseSolvingError)? - .effects; + while let Some(item) = self.constraint_system.pop_front() { + let effects = match item { + ConstraintRef::AlgebraicConstraint(c) => { + c.solve(&self.range_constraints) + .map_err(Error::QseSolvingError)? + .effects + } + ConstraintRef::BusInteraction(b) => b + .solve(&self.bus_interaction_handler, &self.range_constraints) + .map_err(|_| Error::BusInteractionError)?, + }; for effect in effects { progress |= self.apply_effect(effect); } @@ -243,29 +245,10 @@ where Ok(progress) } - /// Tries to infer new information using bus interactions. - fn solve_bus_interactions(&mut self) -> Result { - let mut progress = false; - let effects = self - .constraint_system - .bus_interactions() - .iter() - .map(|bus_interaction| { - bus_interaction.solve(&self.bus_interaction_handler, &self.range_constraints) - }) - // Collect to satisfy borrow checker - .collect::, _>>() - .map_err(|_e| Error::BusInteractionError)?; - for effect in effects.into_iter().flatten() { - progress |= self.apply_effect(effect); - } - Ok(progress) - } - /// Tries to find equivalent variables using quadratic constraints. fn try_solve_quadratic_equivalences(&mut self) -> bool { let equivalences = quadratic_equivalences::find_quadratic_equalities( - self.constraint_system.algebraic_constraints(), + self.constraint_system.system().algebraic_constraints(), &*self, ); for (x, y) in &equivalences { @@ -279,7 +262,7 @@ where /// apply it. This might be expensive. fn exhaustive_search(&mut self) -> Result { let assignments = exhaustive_search::get_unique_assignments( - &self.constraint_system, + self.constraint_system.system(), &*self, &self.bus_interaction_handler, )?; @@ -321,6 +304,7 @@ where } else { // The range constraint was updated. log::trace!("({variable}: {range_constraint})"); + self.constraint_system.variable_updated(variable); } true } else {