conjure_cp_core/ast/
model.rs

1use std::collections::{HashMap, VecDeque};
2use std::fmt::{Debug, Display};
3use std::hash::Hash;
4use std::sync::{Arc, RwLock};
5
6use crate::ast::Domain;
7use crate::context::Context;
8use crate::{bug, into_matrix_expr};
9use derivative::Derivative;
10use indexmap::IndexSet;
11use itertools::izip;
12use parking_lot::{RwLockReadGuard, RwLockWriteGuard};
13use serde::{Deserialize, Serialize};
14use serde_with::serde_as;
15use uniplate::{Biplate, Tree, Uniplate};
16
17use super::serde::{HasId, ObjId, PtrAsInner};
18use super::{
19    Atom, CnfClause, DeclarationPtr, Expression, Literal, Metadata, Moo, Name, Objective,
20    ReturnType, SymbolTable, SymbolTablePtr, Typeable,
21    comprehension::Comprehension,
22    declaration::DeclarationKind,
23    pretty::{
24        pretty_clauses, pretty_domain_letting_declaration, pretty_expressions_as_top_level,
25        pretty_value_letting_declaration, pretty_variable_declaration,
26    },
27};
28
29/// An Essence model.
30#[serde_as]
31#[derive(Derivative, Clone, Debug, Serialize, Deserialize)]
32#[derivative(PartialEq, Eq)]
33pub struct Model {
34    constraints: Moo<Expression>,
35    #[serde(default, skip_serializing_if = "Vec::is_empty")]
36    instantiation_conditions: Vec<Expression>,
37    #[serde_as(as = "PtrAsInner")]
38    symbols: SymbolTablePtr,
39    cnf_clauses: Vec<CnfClause>,
40
41    pub search_order: Option<Vec<Name>>,
42    pub dominance: Option<Expression>,
43    pub objective: Option<Objective>,
44
45    #[serde(skip, default = "default_context")]
46    #[derivative(PartialEq = "ignore")]
47    pub context: Arc<RwLock<Context<'static>>>,
48}
49
50fn default_context() -> Arc<RwLock<Context<'static>>> {
51    Arc::new(RwLock::new(Context::default()))
52}
53
54impl Model {
55    fn new_empty(symbols: SymbolTablePtr, context: Arc<RwLock<Context<'static>>>) -> Model {
56        Model {
57            constraints: Moo::new(Expression::Root(Metadata::new(), vec![])),
58            instantiation_conditions: Vec::new(),
59            symbols,
60            cnf_clauses: Vec::new(),
61            search_order: None,
62            dominance: None,
63            objective: None,
64            context,
65        }
66    }
67
68    /// Creates a new top-level model from the given context.
69    pub fn new(context: Arc<RwLock<Context<'static>>>) -> Model {
70        Self::new_empty(SymbolTablePtr::new(), context)
71    }
72
73    /// Creates a new model whose symbol table has `parent` as parent scope.
74    pub fn new_in_parent_scope(parent: SymbolTablePtr) -> Model {
75        Self::new_empty(SymbolTablePtr::with_parent(parent), default_context())
76    }
77
78    /// The symbol table for this model as a pointer.
79    pub fn symbols_ptr_unchecked(&self) -> &SymbolTablePtr {
80        &self.symbols
81    }
82
83    /// The symbol table for this model as a mutable pointer.
84    pub fn symbols_ptr_unchecked_mut(&mut self) -> &mut SymbolTablePtr {
85        &mut self.symbols
86    }
87
88    /// The symbol table for this model as a reference.
89    pub fn symbols(&self) -> RwLockReadGuard<'_, SymbolTable> {
90        self.symbols.read()
91    }
92
93    /// The symbol table for this model as a mutable reference.
94    pub fn symbols_mut(&mut self) -> RwLockWriteGuard<'_, SymbolTable> {
95        self.symbols.write()
96    }
97
98    /// The root node of this model.
99    pub fn root(&self) -> &Expression {
100        &self.constraints
101    }
102
103    /// The root node of this model, as a mutable reference.
104    ///
105    /// The caller is responsible for ensuring that the root node remains an [`Expression::Root`].
106    pub fn root_mut_unchecked(&mut self) -> &mut Expression {
107        Moo::make_mut(&mut self.constraints)
108    }
109
110    /// Replaces the root node with `new_root`, returning the old root node.
111    pub fn replace_root(&mut self, new_root: Expression) -> Expression {
112        let Expression::Root(_, _) = new_root else {
113            tracing::error!(new_root=?new_root,"new_root is not an Expression::Root");
114            panic!("new_root is not an Expression::Root");
115        };
116
117        std::mem::replace(self.root_mut_unchecked(), new_root)
118    }
119
120    /// The top-level constraints in this model.
121    pub fn constraints(&self) -> &Vec<Expression> {
122        let Expression::Root(_, constraints) = self.constraints.as_ref() else {
123            bug!("The top level expression in a model should be Expr::Root");
124        };
125        constraints
126    }
127
128    /// The cnf clauses in this model.
129    pub fn clauses(&self) -> &Vec<CnfClause> {
130        &self.cnf_clauses
131    }
132
133    /// The top-level constraints in this model as a mutable vector.
134    pub fn constraints_mut(&mut self) -> &mut Vec<Expression> {
135        let Expression::Root(_, constraints) = Moo::make_mut(&mut self.constraints) else {
136            bug!("The top level expression in a model should be Expr::Root");
137        };
138
139        constraints
140    }
141
142    /// The cnf clauses in this model as a mutable vector.
143    pub fn clauses_mut(&mut self) -> &mut Vec<CnfClause> {
144        &mut self.cnf_clauses
145    }
146
147    /// Replaces the top-level constraints with `new_constraints`, returning the old ones.
148    pub fn replace_constraints(&mut self, new_constraints: Vec<Expression>) -> Vec<Expression> {
149        std::mem::replace(self.constraints_mut(), new_constraints)
150    }
151
152    /// Replaces the cnf clauses with `new_clauses`, returning the old ones.
153    pub fn replace_clauses(&mut self, new_clauses: Vec<CnfClause>) -> Vec<CnfClause> {
154        std::mem::replace(self.clauses_mut(), new_clauses)
155    }
156
157    /// Adds a top-level constraint.
158    pub fn add_constraint(&mut self, constraint: Expression) {
159        self.constraints_mut().push(constraint);
160    }
161
162    /// Adds a cnf clause.
163    pub fn add_clause(&mut self, clause: CnfClause) {
164        self.clauses_mut().push(clause);
165    }
166
167    /// Adds top-level constraints.
168    pub fn add_constraints(&mut self, constraints: Vec<Expression>) {
169        self.constraints_mut().extend(constraints);
170    }
171
172    /// Conditions introduced by top-level `where` statements.
173    pub fn instantiation_conditions(&self) -> &[Expression] {
174        &self.instantiation_conditions
175    }
176
177    /// Adds a condition that must hold after parameter instantiation.
178    pub fn add_instantiation_condition(&mut self, condition: Expression) {
179        self.instantiation_conditions.push(condition);
180    }
181
182    /// Removes and returns all pending instantiation conditions.
183    pub fn take_instantiation_conditions(&mut self) -> Vec<Expression> {
184        std::mem::take(&mut self.instantiation_conditions)
185    }
186
187    /// Adds cnf clauses.
188    pub fn add_clauses(&mut self, clauses: Vec<CnfClause>) {
189        self.clauses_mut().extend(clauses);
190    }
191
192    /// Adds a new symbol to the symbol table.
193    pub fn add_symbol(&mut self, decl: DeclarationPtr) -> Option<()> {
194        self.symbols_mut().insert(decl)
195    }
196
197    /// Converts the constraints in this model to a single expression suitable for use inside
198    /// another expression tree.
199    pub fn into_single_expression(self) -> Expression {
200        let constraints = self.constraints().clone();
201        match constraints.len() {
202            0 => Expression::Atomic(Metadata::new(), Atom::Literal(Literal::Bool(true))),
203            1 => constraints[0].clone(),
204            _ => Expression::And(Metadata::new(), Moo::new(into_matrix_expr![constraints])),
205        }
206    }
207
208    /// Collects all ObjId values from the model using uniplate traversal.
209    pub fn collect_stable_id_mapping(&self) -> HashMap<ObjId, ObjId> {
210        fn visit_symbol_table(symbol_table: SymbolTablePtr, id_list: &mut IndexSet<ObjId>) {
211            if !id_list.insert(symbol_table.id()) {
212                return;
213            }
214
215            let table_ref = symbol_table.read();
216            table_ref.iter_local().for_each(|(_, decl)| {
217                id_list.insert(decl.id());
218            });
219        }
220
221        let mut id_list: IndexSet<ObjId> = IndexSet::new();
222
223        visit_symbol_table(self.symbols_ptr_unchecked().clone(), &mut id_list);
224
225        let mut exprs: VecDeque<Expression> = self.universe_bi();
226        if let Some(dominance) = &self.dominance {
227            exprs.push_back(dominance.clone());
228        }
229        if let Some(objective) = &self.objective {
230            exprs.push_back(objective.expression.clone());
231        }
232
233        for symbol_table in Biplate::<SymbolTablePtr>::universe_bi(&exprs) {
234            visit_symbol_table(symbol_table, &mut id_list);
235        }
236        for declaration in Biplate::<DeclarationPtr>::universe_bi(&exprs) {
237            id_list.insert(declaration.id());
238        }
239
240        let mut id_map = HashMap::new();
241        for (stable_id, original_id) in id_list.into_iter().enumerate() {
242            let type_name = original_id.type_name;
243            id_map.insert(
244                original_id,
245                ObjId {
246                    object_id: stable_id as u32,
247                    type_name,
248                },
249            );
250        }
251
252        id_map
253    }
254}
255
256impl Default for Model {
257    fn default() -> Self {
258        Self::new(default_context())
259    }
260}
261
262impl Typeable for Model {
263    fn return_type(&self) -> ReturnType {
264        ReturnType::Bool
265    }
266}
267
268impl Hash for Model {
269    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
270        self.constraints.hash(state);
271        self.instantiation_conditions.hash(state);
272        self.symbols.hash(state);
273        self.cnf_clauses.hash(state);
274        self.search_order.hash(state);
275        self.dominance.hash(state);
276        self.objective.hash(state);
277    }
278}
279
280// At time of writing (03/02/2025), the Uniplate derive macro doesn't like the lifetimes inside
281// context, and we do not yet have a way of ignoring this field.
282impl Uniplate for Model {
283    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
284        let self2 = self.clone();
285        (Tree::Zero, Box::new(move |_| self2.clone()))
286    }
287}
288
289impl Biplate<Expression> for Model {
290    fn biplate(&self) -> (Tree<Expression>, Box<dyn Fn(Tree<Expression>) -> Self>) {
291        let (symtab_tree, symtab_ctx) =
292            <SymbolTable as Biplate<Expression>>::biplate(&self.symbols());
293
294        let dom_tree = match &self.dominance {
295            Some(expr) => Tree::One(expr.clone()),
296            None => Tree::Zero,
297        };
298
299        let obj_tree = match &self.objective {
300            Some(objective) => Tree::One(objective.expression.clone()),
301            None => Tree::Zero,
302        };
303
304        let instantiation_tree = Tree::Many(
305            self.instantiation_conditions
306                .iter()
307                .cloned()
308                .map(Tree::One)
309                .collect(),
310        );
311
312        let tree = Tree::Many(VecDeque::from([
313            Tree::One(self.root().clone()),
314            instantiation_tree,
315            symtab_tree,
316            dom_tree,
317            obj_tree,
318        ]));
319
320        let obj_direction = self.objective.as_ref().map(|objective| objective.direction);
321
322        let self2 = self.clone();
323        let ctx = Box::new(move |x| {
324            let Tree::Many(xs) = x else {
325                panic!("Expected a tree with five children");
326            };
327            if xs.len() != 5 {
328                panic!("Expected a tree with five children");
329            }
330
331            let Tree::One(root) = xs[0].clone() else {
332                panic!("Expected root expression tree");
333            };
334
335            let Tree::Many(instantiation_conditions) = xs[1].clone() else {
336                panic!("Expected instantiation-condition expression tree");
337            };
338            let symtab = symtab_ctx(xs[2].clone());
339            let dominance = match xs[3].clone() {
340                Tree::One(expr) => Some(expr),
341                Tree::Zero => None,
342                _ => panic!("Expected dominance tree"),
343            };
344            let objective_expr = match xs[4].clone() {
345                Tree::One(expr) => Some(expr),
346                Tree::Zero => None,
347                _ => panic!("Expected objective tree"),
348            };
349
350            let mut self3 = self2.clone();
351
352            let Expression::Root(_, _) = root else {
353                bug!("root expression not root");
354            };
355
356            *self3.root_mut_unchecked() = root;
357            self3.instantiation_conditions = instantiation_conditions
358                .into_iter()
359                .map(|tree| match tree {
360                    Tree::One(expr) => expr,
361                    _ => panic!("Expected instantiation condition expression"),
362                })
363                .collect();
364            *self3.symbols_mut() = symtab;
365            self3.dominance = dominance;
366            self3.objective = match (obj_direction, objective_expr) {
367                (Some(direction), Some(expression)) => Some(Objective {
368                    direction,
369                    expression,
370                }),
371                _ => None,
372            };
373
374            self3
375        });
376
377        (tree, ctx)
378    }
379}
380
381impl Biplate<Atom> for Model {
382    fn biplate(&self) -> (Tree<Atom>, Box<dyn Fn(Tree<Atom>) -> Self>) {
383        let (expression_tree, rebuild_self) = <Model as Biplate<Expression>>::biplate(self);
384        let (expression_list, rebuild_expression_tree) = expression_tree.list();
385
386        let (atom_trees, reconstruct_exprs): (VecDeque<_>, VecDeque<_>) = expression_list
387            .iter()
388            .map(|e| <Expression as Biplate<Atom>>::biplate(e))
389            .unzip();
390
391        let tree = Tree::Many(atom_trees);
392        let ctx = Box::new(move |atom_tree: Tree<Atom>| {
393            let Tree::Many(atoms) = atom_tree else {
394                panic!();
395            };
396
397            assert_eq!(
398                atoms.len(),
399                reconstruct_exprs.len(),
400                "the number of children should not change when using Biplate"
401            );
402
403            let expression_list: VecDeque<Expression> = izip!(atoms, &reconstruct_exprs)
404                .map(|(atom, recons)| recons(atom))
405                .collect();
406
407            let expression_tree = rebuild_expression_tree(expression_list);
408            rebuild_self(expression_tree)
409        });
410
411        (tree, ctx)
412    }
413}
414
415impl Biplate<Comprehension> for Model {
416    fn biplate(
417        &self,
418    ) -> (
419        Tree<Comprehension>,
420        Box<dyn Fn(Tree<Comprehension>) -> Self>,
421    ) {
422        let (f1_tree, f1_ctx) = <_ as Biplate<Comprehension>>::biplate(&self.constraints);
423        let (f2_tree, f2_ctx) = <SymbolTable as Biplate<Comprehension>>::biplate(&self.symbols());
424
425        let tree = Tree::Many(VecDeque::from([f1_tree, f2_tree]));
426        let self2 = self.clone();
427        let ctx = Box::new(move |x| {
428            let Tree::Many(xs) = x else {
429                panic!();
430            };
431
432            let root = f1_ctx(xs[0].clone());
433            let symtab = f2_ctx(xs[1].clone());
434
435            let mut self3 = self2.clone();
436
437            let Expression::Root(_, _) = &*root else {
438                bug!("root expression not root");
439            };
440
441            *self3.symbols_mut() = symtab;
442            self3.constraints = root;
443
444            self3
445        });
446
447        (tree, ctx)
448    }
449}
450
451impl Display for Model {
452    #[allow(clippy::unwrap_used)]
453    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
454        let symbols = self.symbols();
455        for (name, decl) in symbols.iter_local() {
456            match &decl.kind() as &DeclarationKind {
457                DeclarationKind::Find(_) | DeclarationKind::FindAuxiliary(_) => {
458                    writeln!(
459                        f,
460                        "{}",
461                        pretty_variable_declaration(&symbols, name).unwrap()
462                    )?;
463                }
464                DeclarationKind::ValueLetting(_, _) | DeclarationKind::TemporaryValueLetting(_) => {
465                    writeln!(
466                        f,
467                        "{}",
468                        pretty_value_letting_declaration(&symbols, name).unwrap()
469                    )?;
470                }
471                DeclarationKind::DomainLetting(_) => {
472                    writeln!(
473                        f,
474                        "{}",
475                        pretty_domain_letting_declaration(&symbols, name).unwrap()
476                    )?;
477                }
478                DeclarationKind::Given(d) => {
479                    writeln!(f, "given {name}: {d}")?;
480                }
481                DeclarationKind::Quantified(inner) => {
482                    writeln!(f, "quantified {name}: {}", inner.domain())?;
483                }
484                DeclarationKind::QuantifiedExpr(expr) => {
485                    writeln!(f, "quantified expr {name} <- {}", expr)?;
486                }
487            }
488        }
489
490        if !self.constraints().is_empty() {
491            writeln!(f, "\nsuch that\n")?;
492            writeln!(f, "{}", pretty_expressions_as_top_level(self.constraints()))?;
493        }
494
495        if !self.instantiation_conditions.is_empty() {
496            writeln!(f, "\nwhere\n")?;
497            writeln!(
498                f,
499                "{}",
500                pretty_expressions_as_top_level(&self.instantiation_conditions)
501            )?;
502        }
503
504        if !self.clauses().is_empty() {
505            writeln!(f, "\nclauses:\n")?;
506            writeln!(f, "{}", pretty_clauses(self.clauses()))?;
507        }
508        Ok(())
509    }
510}
511
512/// A model that is de/serializable using `serde`.
513///
514/// To turn this into a rewritable model, it needs to be initialised using
515/// [`initialise`](SerdeModel::initialise).
516#[serde_as]
517#[derive(Clone, Debug, Serialize, Deserialize)]
518pub struct SerdeModel {
519    constraints: Moo<Expression>,
520    #[serde(default, skip_serializing_if = "Vec::is_empty")]
521    instantiation_conditions: Vec<Expression>,
522    #[serde_as(as = "PtrAsInner")]
523    symbols: SymbolTablePtr,
524    cnf_clauses: Vec<CnfClause>,
525    search_order: Option<Vec<Name>>,
526    dominance: Option<Expression>,
527    objective: Option<Objective>,
528}
529
530impl SerdeModel {
531    /// Initialises the model for rewriting.
532    pub fn initialise(mut self, context: Arc<RwLock<Context<'static>>>) -> Option<Model> {
533        let mut tables: HashMap<ObjId, SymbolTablePtr> = HashMap::new();
534
535        // Root model symbol table is always definitive.
536        tables.insert(self.symbols.id(), self.symbols.clone());
537
538        let mut exprs: VecDeque<Expression> = self.constraints.universe_bi();
539        exprs.extend(self.instantiation_conditions.clone());
540        if let Some(dominance) = &self.dominance {
541            exprs.push_back(dominance.clone());
542        }
543        if let Some(objective) = &self.objective {
544            exprs.push_back(objective.expression.clone());
545        }
546
547        // Some expressions (e.g. abstract comprehensions) contain additional symbol tables.
548        for table in Biplate::<SymbolTablePtr>::universe_bi(&exprs) {
549            tables.entry(table.id()).or_insert(table);
550        }
551
552        for table in tables.clone().into_values() {
553            let mut table_mut = table.write();
554            let parent_mut = table_mut.parent_mut_unchecked();
555
556            #[allow(clippy::unwrap_used)]
557            if let Some(parent) = parent_mut {
558                let parent_id = parent.id();
559                *parent = tables.get(&parent_id).unwrap().clone();
560            }
561        }
562
563        let mut all_declarations: HashMap<ObjId, DeclarationPtr> = HashMap::new();
564        for table in tables.values() {
565            for (_, decl) in table.read().iter_local() {
566                let id = decl.id();
567                all_declarations.insert(id, decl.clone());
568            }
569        }
570
571        self.constraints = self.constraints.transform_bi(&move |decl: DeclarationPtr| {
572            let id = decl.id();
573            all_declarations
574                .get(&id)
575                .unwrap_or_else(|| {
576                    panic!(
577                        "A declaration used in the expression tree should exist in the symbol table. The missing declaration has id {id}."
578                    )
579                })
580                .clone()
581        });
582
583        Some(Model {
584            constraints: self.constraints,
585            instantiation_conditions: self.instantiation_conditions,
586            symbols: self.symbols,
587            cnf_clauses: self.cnf_clauses,
588            search_order: self.search_order,
589            dominance: self.dominance,
590            objective: self.objective,
591            context,
592        })
593    }
594}
595
596impl From<Model> for SerdeModel {
597    fn from(val: Model) -> Self {
598        SerdeModel {
599            constraints: val.constraints,
600            instantiation_conditions: val.instantiation_conditions,
601            symbols: val.symbols,
602            cnf_clauses: val.cnf_clauses,
603            search_order: val.search_order,
604            dominance: val.dominance,
605            objective: val.objective,
606        }
607    }
608}
609
610impl Display for SerdeModel {
611    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
612        let model = Model {
613            constraints: self.constraints.clone(),
614            instantiation_conditions: self.instantiation_conditions.clone(),
615            symbols: self.symbols.clone(),
616            cnf_clauses: self.cnf_clauses.clone(),
617            search_order: self.search_order.clone(),
618            dominance: self.dominance.clone(),
619            objective: self.objective.clone(),
620            context: default_context(),
621        };
622        std::fmt::Display::fmt(&model, f)
623    }
624}
625
626impl SerdeModel {
627    /// Collects all ObjId values from the model and maps them to stable sequential IDs.
628    pub fn collect_stable_id_mapping(&self) -> HashMap<ObjId, ObjId> {
629        let model = Model {
630            constraints: self.constraints.clone(),
631            instantiation_conditions: self.instantiation_conditions.clone(),
632            symbols: self.symbols.clone(),
633            cnf_clauses: self.cnf_clauses.clone(),
634            search_order: self.search_order.clone(),
635            dominance: self.dominance.clone(),
636            objective: self.objective.clone(),
637            context: default_context(),
638        };
639        model.collect_stable_id_mapping()
640    }
641}
642
643/// A struct for the information about expressions
644#[serde_as]
645#[derive(Serialize)]
646pub struct ExprInfo {
647    pretty: String,
648    domain: Option<Moo<Domain>>,
649    children: Vec<ExprInfo>,
650}
651
652impl ExprInfo {
653    pub fn create(expr: &Expression) -> ExprInfo {
654        let pretty = expr.to_string();
655        let domain = expr.domain_of();
656        let children = expr.children().iter().map(Self::create).collect();
657
658        ExprInfo {
659            pretty,
660            domain,
661            children,
662        }
663    }
664}
conjure_cp_core/ast/model.rs

conjure_cp_core/ast/
model.rs
