Lines
53.23 %
Functions
30.77 %
use itertools::Itertools;
use std::collections::{HashMap, VecDeque};
use std::fmt::{Debug, Display};
use std::hash::Hash;
use std::sync::{Arc, RwLock};
use crate::ast::Domain;
use crate::context::Context;
use crate::{bug, into_matrix_expr};
use derivative::Derivative;
use indexmap::IndexSet;
use itertools::izip;
use parking_lot::{RwLockReadGuard, RwLockWriteGuard};
use serde::{Deserialize, Serialize};
use serde_with::serde_as;
use uniplate::{Biplate, Tree, Uniplate};
use super::serde::{HasId, ObjId, PtrAsInner};
use super::{
Atom, DeclarationPtr, Expression, Literal, Metadata, Moo, Name, Objective, ReturnType,
SatEncodingDecision, SymbolTable, SymbolTablePtr, Typeable,
comprehension::Comprehension,
declaration::DeclarationKind,
pretty::{
pretty_domain_letting_declaration, pretty_expressions_as_top_level, pretty_sat_decisions,
pretty_value_letting_declaration, pretty_variable_declaration,
},
};
/// An Essence model.
#[serde_as]
#[derive(Derivative, Clone, Debug, Serialize, Deserialize)]
#[derivative(PartialEq, Eq)]
pub struct Model {
constraints: Moo<Expression>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
instantiation_conditions: Vec<Expression>,
#[serde_as(as = "PtrAsInner")]
symbols: SymbolTablePtr,
sat_decisions: Vec<SatEncodingDecision>,
pub search_order: Option<Vec<Name>>,
pub dominance: Option<Expression>,
pub objective: Option<Objective>,
#[serde(skip, default = "default_context")]
#[derivative(PartialEq = "ignore")]
pub context: Arc<RwLock<Context<'static>>>,
}
fn default_context() -> Arc<RwLock<Context<'static>>> {
Arc::new(RwLock::new(Context::default()))
impl Model {
fn new_empty(symbols: SymbolTablePtr, context: Arc<RwLock<Context<'static>>>) -> Model {
Model {
constraints: Moo::new(Expression::Root(Metadata::new(), vec![])),
instantiation_conditions: Vec::new(),
symbols,
sat_decisions: Vec::new(),
search_order: None,
dominance: None,
objective: None,
context,
/// Creates a new top-level model from the given context.
pub fn new(context: Arc<RwLock<Context<'static>>>) -> Model {
Self::new_empty(SymbolTablePtr::new(), context)
/// Creates a new model whose symbol table has `parent` as parent scope.
pub fn new_in_parent_scope(parent: SymbolTablePtr) -> Model {
Self::new_empty(SymbolTablePtr::with_parent(parent), default_context())
/// The symbol table for this model as a pointer.
pub fn symbols_ptr_unchecked(&self) -> &SymbolTablePtr {
&self.symbols
/// The symbol table for this model as a mutable pointer.
pub fn symbols_ptr_unchecked_mut(&mut self) -> &mut SymbolTablePtr {
&mut self.symbols
/// The symbol table for this model as a reference.
pub fn symbols(&self) -> RwLockReadGuard<'_, SymbolTable> {
self.symbols.read()
/// The symbol table for this model as a mutable reference.
pub fn symbols_mut(&mut self) -> RwLockWriteGuard<'_, SymbolTable> {
self.symbols.write()
/// The root node of this model.
pub fn root(&self) -> &Expression {
&self.constraints
/// The root node of this model, as a mutable reference.
///
/// The caller is responsible for ensuring that the root node remains an [`Expression::Root`].
pub fn root_mut_unchecked(&mut self) -> &mut Expression {
Moo::make_mut(&mut self.constraints)
/// Replaces the root node with `new_root`, returning the old root node.
pub fn replace_root(&mut self, new_root: Expression) -> Expression {
let Expression::Root(_, _) = new_root else {
tracing::error!(new_root=?new_root,"new_root is not an Expression::Root");
panic!("new_root is not an Expression::Root");
std::mem::replace(self.root_mut_unchecked(), new_root)
/// The top-level constraints in this model.
pub fn constraints(&self) -> &Vec<Expression> {
let Expression::Root(_, constraints) = self.constraints.as_ref() else {
bug!("The top level expression in a model should be Expr::Root");
constraints
/// The SAT encoding decisions in this model.
pub fn sat_decisions(&self) -> &Vec<SatEncodingDecision> {
&self.sat_decisions
/// The top-level constraints in this model as a mutable vector.
pub fn constraints_mut(&mut self) -> &mut Vec<Expression> {
let Expression::Root(_, constraints) = Moo::make_mut(&mut self.constraints) else {
/// The SAT encoding decisions in this model as a mutable vector.
pub fn sat_decisions_mut(&mut self) -> &mut Vec<SatEncodingDecision> {
&mut self.sat_decisions
/// Replaces the top-level constraints with `new_constraints`, returning the old ones.
pub fn replace_constraints(&mut self, new_constraints: Vec<Expression>) -> Vec<Expression> {
std::mem::replace(self.constraints_mut(), new_constraints)
/// Replaces the SAT encoding decisions with `new_sat_decisions`, returning the old ones.
pub fn replace_sat_decisions(
&mut self,
new_sat_decisions: Vec<SatEncodingDecision>,
) -> Vec<SatEncodingDecision> {
std::mem::replace(self.sat_decisions_mut(), new_sat_decisions)
/// Adds a top-level constraint.
pub fn add_constraint(&mut self, constraint: Expression) {
self.constraints_mut().push(constraint);
/// Adds a SAT encoding decision.
pub fn add_sat_decision(&mut self, decision: SatEncodingDecision) {
self.sat_decisions_mut().push(decision);
/// Adds top-level constraints.
pub fn add_constraints(&mut self, constraints: Vec<Expression>) {
self.constraints_mut().extend(constraints);
/// Conditions introduced by top-level `where` statements.
pub fn instantiation_conditions(&self) -> &[Expression] {
&self.instantiation_conditions
/// Adds a condition that must hold after parameter instantiation.
pub fn add_instantiation_condition(&mut self, condition: Expression) {
self.instantiation_conditions.push(condition);
/// Removes and returns all pending instantiation conditions.
pub fn take_instantiation_conditions(&mut self) -> Vec<Expression> {
std::mem::take(&mut self.instantiation_conditions)
/// Adds SAT encoding decisions.
pub fn add_sat_decisions(&mut self, decisions: Vec<SatEncodingDecision>) {
self.sat_decisions_mut().extend(decisions);
/// Adds a new symbol to the symbol table.
pub fn add_symbol(&mut self, decl: DeclarationPtr) -> Option<()> {
self.symbols_mut().insert(decl)
/// Converts the constraints in this model to a single expression suitable for use inside
/// another expression tree.
pub fn into_single_expression(self) -> Expression {
let constraints = self.constraints().clone();
match constraints.len() {
0 => Expression::Atomic(Metadata::new(), Atom::Literal(Literal::Bool(true))),
1 => constraints[0].clone(),
_ => Expression::And(Metadata::new(), Moo::new(into_matrix_expr![constraints])),
/// Collects all ObjId values from the model using uniplate traversal.
pub fn collect_stable_id_mapping(&self) -> HashMap<ObjId, ObjId> {
fn visit_symbol_table(symbol_table: SymbolTablePtr, id_list: &mut IndexSet<ObjId>) {
if !id_list.insert(symbol_table.id()) {
return;
let table_ref = symbol_table.read();
table_ref.iter_local().for_each(|(_, decl)| {
id_list.insert(decl.id());
});
let mut id_list: IndexSet<ObjId> = IndexSet::new();
visit_symbol_table(self.symbols_ptr_unchecked().clone(), &mut id_list);
let mut exprs: VecDeque<Expression> = self.universe_bi();
exprs.extend(self.sat_decisions.universe_bi());
if let Some(dominance) = &self.dominance {
exprs.push_back(dominance.clone());
if let Some(objective) = &self.objective {
exprs.push_back(objective.expression.clone());
for symbol_table in Biplate::<SymbolTablePtr>::universe_bi(&exprs) {
visit_symbol_table(symbol_table, &mut id_list);
for declaration in Biplate::<DeclarationPtr>::universe_bi(&exprs) {
id_list.insert(declaration.id());
let mut id_map = HashMap::new();
for (stable_id, original_id) in id_list.into_iter().enumerate() {
let type_name = original_id.type_name;
id_map.insert(
original_id,
ObjId {
object_id: stable_id as u32,
type_name,
);
id_map
impl Default for Model {
fn default() -> Self {
Self::new(default_context())
impl Typeable for Model {
fn return_type(&self) -> ReturnType {
ReturnType::Bool
impl Hash for Model {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.constraints.hash(state);
self.instantiation_conditions.hash(state);
self.symbols.hash(state);
self.sat_decisions.hash(state);
self.search_order.hash(state);
self.dominance.hash(state);
self.objective.hash(state);
// At time of writing (03/02/2025), the Uniplate derive macro doesn't like the lifetimes inside
// context, and we do not yet have a way of ignoring this field.
impl Uniplate for Model {
fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
let self2 = self.clone();
(Tree::Zero, Box::new(move |_| self2.clone()))
impl Biplate<Expression> for Model {
fn biplate(&self) -> (Tree<Expression>, Box<dyn Fn(Tree<Expression>) -> Self>) {
let (symtab_tree, symtab_ctx) =
<SymbolTable as Biplate<Expression>>::biplate(&self.symbols());
let dom_tree = match &self.dominance {
Some(expr) => Tree::One(expr.clone()),
None => Tree::Zero,
let obj_tree = match &self.objective {
Some(objective) => Tree::One(objective.expression.clone()),
let instantiation_tree = Tree::Many(
self.instantiation_conditions
.iter()
.cloned()
.map(Tree::One)
.collect(),
let tree = Tree::Many(VecDeque::from([
Tree::One(self.root().clone()),
instantiation_tree,
symtab_tree,
dom_tree,
obj_tree,
]));
let obj_direction = self.objective.as_ref().map(|objective| objective.direction);
let ctx = Box::new(move |x| {
let Tree::Many(xs) = x else {
panic!("Expected a tree with five children");
if xs.len() != 5 {
let Tree::One(root) = xs[0].clone() else {
panic!("Expected root expression tree");
let Tree::Many(instantiation_conditions) = xs[1].clone() else {
panic!("Expected instantiation-condition expression tree");
let symtab = symtab_ctx(xs[2].clone());
let dominance = match xs[3].clone() {
Tree::One(expr) => Some(expr),
Tree::Zero => None,
_ => panic!("Expected dominance tree"),
let objective_expr = match xs[4].clone() {
_ => panic!("Expected objective tree"),
let mut self3 = self2.clone();
let Expression::Root(_, _) = root else {
bug!("root expression not root");
*self3.root_mut_unchecked() = root;
self3.instantiation_conditions = instantiation_conditions
.into_iter()
.map(|tree| match tree {
Tree::One(expr) => expr,
_ => panic!("Expected instantiation condition expression"),
})
.collect();
*self3.symbols_mut() = symtab;
self3.dominance = dominance;
self3.objective = match (obj_direction, objective_expr) {
(Some(direction), Some(expression)) => Some(Objective {
direction,
expression,
}),
_ => None,
self3
(tree, ctx)
impl Biplate<Atom> for Model {
fn biplate(&self) -> (Tree<Atom>, Box<dyn Fn(Tree<Atom>) -> Self>) {
let (expression_tree, rebuild_self) = <Model as Biplate<Expression>>::biplate(self);
let (expression_list, rebuild_expression_tree) = expression_tree.list();
let (atom_trees, reconstruct_exprs): (VecDeque<_>, VecDeque<_>) = expression_list
.map(|e| <Expression as Biplate<Atom>>::biplate(e))
.unzip();
let tree = Tree::Many(atom_trees);
let ctx = Box::new(move |atom_tree: Tree<Atom>| {
let Tree::Many(atoms) = atom_tree else {
panic!();
assert_eq!(
atoms.len(),
reconstruct_exprs.len(),
"the number of children should not change when using Biplate"
let expression_list: VecDeque<Expression> = izip!(atoms, &reconstruct_exprs)
.map(|(atom, recons)| recons(atom))
let expression_tree = rebuild_expression_tree(expression_list);
rebuild_self(expression_tree)
impl Biplate<Comprehension> for Model {
fn biplate(
&self,
) -> (
Tree<Comprehension>,
Box<dyn Fn(Tree<Comprehension>) -> Self>,
) {
let (f1_tree, f1_ctx) = <_ as Biplate<Comprehension>>::biplate(&self.constraints);
let (f2_tree, f2_ctx) = <SymbolTable as Biplate<Comprehension>>::biplate(&self.symbols());
let tree = Tree::Many(VecDeque::from([f1_tree, f2_tree]));
let root = f1_ctx(xs[0].clone());
let symtab = f2_ctx(xs[1].clone());
let Expression::Root(_, _) = &*root else {
self3.constraints = root;
impl Display for Model {
#[allow(clippy::unwrap_used)]
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let symbols = self.symbols();
// Enum members are introduced by their type declaration, before any value uses them.
for (name, decl) in symbols.iter_local() {
if let DeclarationKind::EnumeratedType(ty) = &*decl.kind() {
writeln!(
f,
"letting {name} be new type enum {{{}}}",
ty.variants.iter().join(", ")
)?;
} else if let DeclarationKind::UnnamedType(size) = &decl.kind() as &DeclarationKind {
writeln!(f, "letting {name} be new type of size {size}")?;
match &decl.kind() as &DeclarationKind {
DeclarationKind::ValueLetting(
Expression::Atomic(_, Atom::Literal(Literal::EnumVariant(value))),
_,
) if value.name() == name => {}
DeclarationKind::Find(_) | DeclarationKind::FindAuxiliary(_) => {
"{}",
pretty_variable_declaration(&symbols, name).unwrap()
DeclarationKind::ValueLetting(_, _) | DeclarationKind::TemporaryValueLetting(_) => {
pretty_value_letting_declaration(&symbols, name).unwrap()
DeclarationKind::DomainLetting(_) => {
pretty_domain_letting_declaration(&symbols, name).unwrap()
DeclarationKind::Given(d) => {
writeln!(f, "given {name}: {d}")?;
DeclarationKind::Quantified(inner) => {
writeln!(f, "quantified {name}: {}", inner.domain())?;
DeclarationKind::QuantifiedExpr(expr) => {
writeln!(f, "quantified expr {name} <- {}", expr)?;
DeclarationKind::EnumeratedType(_) | DeclarationKind::UnnamedType(_) => {}
if !self.constraints().is_empty() {
writeln!(f, "\nsuch that\n")?;
writeln!(f, "{}", pretty_expressions_as_top_level(self.constraints()))?;
if !self.instantiation_conditions.is_empty() {
writeln!(f, "\nwhere\n")?;
pretty_expressions_as_top_level(&self.instantiation_conditions)
if !self.sat_decisions().is_empty() {
writeln!(f, "\nsat encoding decisions:\n")?;
writeln!(f, "{}", pretty_sat_decisions(self.sat_decisions()))?;
Ok(())
/// A model that is de/serializable using `serde`.
/// To turn this into a rewritable model, it needs to be initialised using
/// [`initialise`](SerdeModel::initialise).
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SerdeModel {
search_order: Option<Vec<Name>>,
dominance: Option<Expression>,
objective: Option<Objective>,
impl SerdeModel {
/// Initialises the model for rewriting.
pub fn initialise(mut self, context: Arc<RwLock<Context<'static>>>) -> Option<Model> {
let mut tables: HashMap<ObjId, SymbolTablePtr> = HashMap::new();
// Root model symbol table is always definitive.
tables.insert(self.symbols.id(), self.symbols.clone());
let mut exprs: VecDeque<Expression> = self.constraints.universe_bi();
exprs.extend(self.instantiation_conditions.clone());
// Some expressions (e.g. abstract comprehensions) contain additional symbol tables.
for table in Biplate::<SymbolTablePtr>::universe_bi(&exprs) {
tables.entry(table.id()).or_insert(table);
for table in tables.clone().into_values() {
let mut table_mut = table.write();
let parent_mut = table_mut.parent_mut_unchecked();
if let Some(parent) = parent_mut {
let parent_id = parent.id();
*parent = tables.get(&parent_id).unwrap().clone();
let mut all_declarations: HashMap<ObjId, DeclarationPtr> = HashMap::new();
for table in tables.values() {
for (_, decl) in table.read().iter_local() {
let id = decl.id();
all_declarations.insert(id, decl.clone());
let rebind = |decl: DeclarationPtr| {
all_declarations
.get(&id)
.unwrap_or_else(|| {
panic!(
"A declaration used in the expression tree should exist in the symbol table. The missing declaration has id {id}."
)
.clone()
self.constraints = self.constraints.transform_bi(&rebind);
self.sat_decisions = self.sat_decisions.transform_bi(&rebind);
Some(Model {
constraints: self.constraints,
instantiation_conditions: self.instantiation_conditions,
symbols: self.symbols,
sat_decisions: self.sat_decisions,
search_order: self.search_order,
dominance: self.dominance,
objective: self.objective,
impl From<Model> for SerdeModel {
fn from(val: Model) -> Self {
SerdeModel {
constraints: val.constraints,
instantiation_conditions: val.instantiation_conditions,
symbols: val.symbols,
sat_decisions: val.sat_decisions,
search_order: val.search_order,
dominance: val.dominance,
objective: val.objective,
impl Display for SerdeModel {
let model = Model {
constraints: self.constraints.clone(),
instantiation_conditions: self.instantiation_conditions.clone(),
symbols: self.symbols.clone(),
sat_decisions: self.sat_decisions.clone(),
search_order: self.search_order.clone(),
dominance: self.dominance.clone(),
objective: self.objective.clone(),
context: default_context(),
std::fmt::Display::fmt(&model, f)
/// Collects all ObjId values from the model and maps them to stable sequential IDs.
model.collect_stable_id_mapping()
/// A struct for the information about expressions
#[derive(Serialize)]
pub struct ExprInfo {
pretty: String,
domain: Option<Moo<Domain>>,
children: Vec<ExprInfo>,
impl ExprInfo {
pub fn create(expr: &Expression) -> ExprInfo {
let pretty = expr.to_string();
let domain = expr.domain_of();
let children = expr.children().iter().map(Self::create).collect();
ExprInfo {
pretty,
domain,
children,