conjure_cp_core/ast/
model.rs

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

conjure_cp_core/ast/
model.rs
