Skip to main content

conjure_cp_core/ast/
model.rs

1use std::collections::{HashMap, VecDeque};
2use std::fmt::{Debug, Display};
3use std::sync::{Arc, RwLock};
4
5use derivative::Derivative;
6use indexmap::IndexSet;
7use serde::{Deserialize, Serialize};
8use uniplate::{Biplate, Tree, Uniplate};
9
10use crate::context::Context;
11
12use super::serde::{HasId, ObjId};
13use super::{DeclarationPtr, Expression, Name, ReturnType, SubModel, SymbolTablePtr, Typeable};
14
15/// An Essence model.
16///
17/// - This type wraps a [`Submodel`] containing the top-level lexical scope. To manipulate the
18///   model's constraints or symbols, first convert it to a [`Submodel`] using
19///   [`as_submodel`](Model::as_submodel) / [`as_submodel_mut`](Model::as_submodel_mut).
20///
21/// - To de/serialise a model using `serde`, see [`SerdeModel`].
22#[derive(Derivative, Clone, Debug)]
23#[derivative(PartialEq, Eq)]
24pub struct Model {
25    submodel: SubModel,
26    pub search_order: Option<Vec<Name>>,
27    pub dominance: Option<Expression>,
28    #[derivative(PartialEq = "ignore")]
29    pub context: Arc<RwLock<Context<'static>>>,
30}
31
32impl Model {
33    pub fn from_submodel(submodel: SubModel) -> Model {
34        Model {
35            submodel,
36            ..Default::default()
37        }
38    }
39
40    /// Creates a new model from the given context.
41    pub fn new(context: Arc<RwLock<Context<'static>>>) -> Model {
42        Model {
43            submodel: SubModel::new_top_level(),
44            dominance: None,
45            context,
46            search_order: None,
47        }
48    }
49
50    /// Returns this model as a [`Submodel`].
51    pub fn as_submodel(&self) -> &SubModel {
52        &self.submodel
53    }
54
55    /// Returns this model as a mutable [`Submodel`].
56    pub fn as_submodel_mut(&mut self) -> &mut SubModel {
57        &mut self.submodel
58    }
59
60    /// Replaces the model contents with `new_submodel`, returning the old contents.
61    pub fn replace_submodel(&mut self, new_submodel: SubModel) -> SubModel {
62        std::mem::replace(self.as_submodel_mut(), new_submodel)
63    }
64}
65
66impl Default for Model {
67    fn default() -> Self {
68        Model {
69            submodel: SubModel::new_top_level(),
70            dominance: None,
71            context: Arc::new(RwLock::new(Context::default())),
72            search_order: None,
73        }
74    }
75}
76
77impl Typeable for Model {
78    fn return_type(&self) -> ReturnType {
79        ReturnType::Bool
80    }
81}
82
83// At time of writing (03/02/2025), the Uniplate derive macro doesn't like the lifetimes inside
84// context, and we do not yet have a way of ignoring this field.
85impl Uniplate for Model {
86    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
87        // Model contains no sub-models.
88        let self2 = self.clone();
89        (Tree::Zero, Box::new(move |_| self2.clone()))
90    }
91}
92
93impl Biplate<Expression> for Model {
94    fn biplate(&self) -> (Tree<Expression>, Box<dyn Fn(Tree<Expression>) -> Self>) {
95        // walk into submodel
96        let submodel = self.as_submodel();
97        let (expr_tree, expr_ctx) = <SubModel as Biplate<Expression>>::biplate(submodel);
98
99        // walk into dominance relation if it exists
100        let dom_tree = match &self.dominance {
101            Some(expr) => Tree::One(expr.clone()),
102            None => Tree::Zero,
103        };
104        let tree = Tree::<Expression>::Many(VecDeque::from([expr_tree, dom_tree]));
105
106        let self2 = self.clone();
107        let ctx = Box::new(move |x| match x {
108            Tree::Many(xs) => {
109                if xs.len() != 2 {
110                    panic!("Expected a tree with two children");
111                }
112                let submodel_tree = xs[0].clone();
113                let dom_tree = xs[1].clone();
114
115                // reconstruct the submodel
116                let submodel = expr_ctx(submodel_tree);
117                // reconstruct the dominance relation
118                let dominance = match dom_tree {
119                    Tree::One(expr) => Some(expr),
120                    Tree::Zero => None,
121                    _ => panic!("Expected a tree with two children"),
122                };
123
124                let mut self3 = self2.clone();
125                self3.replace_submodel(submodel);
126                self3.dominance = dominance;
127                self3
128            }
129            _ => {
130                panic!("Expected a tree with two children");
131            }
132        });
133
134        (tree, ctx)
135    }
136}
137
138impl Biplate<SubModel> for Model {
139    fn biplate(&self) -> (Tree<SubModel>, Box<dyn Fn(Tree<SubModel>) -> Self>) {
140        let submodel = self.as_submodel().clone();
141
142        let self2 = self.clone();
143        let ctx = Box::new(move |x| {
144            let Tree::One(submodel) = x else {
145                panic!();
146            };
147
148            let mut self3 = self2.clone();
149            self3.replace_submodel(submodel);
150            self3
151        });
152
153        (Tree::One(submodel), ctx)
154    }
155}
156
157impl Display for Model {
158    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
159        std::fmt::Display::fmt(self.as_submodel(), f)
160    }
161}
162
163/// A model that is de/serializable using `serde`.
164///
165/// To turn this into a rewritable model, it needs to be initialised using [`initialise`](SerdeModel::initialise).
166///
167/// To deserialise a [`Model`], use `.into()` to convert it into a `SerdeModel` first.
168#[derive(Clone, Debug, Serialize, Deserialize)]
169pub struct SerdeModel {
170    #[serde(flatten)]
171    submodel: SubModel,
172    search_order: Option<Vec<Name>>, // TODO: make this a [expressions]
173    dominance: Option<Expression>,
174}
175
176impl SerdeModel {
177    /// Initialises the model for rewriting.
178    ///
179    /// This swizzles the pointers to symbol tables and declarations using the stored ids.
180    pub fn initialise(mut self, context: Arc<RwLock<Context<'static>>>) -> Option<Model> {
181        // The definitive versions of each symbol table are stored in the submodels. Parent
182        // pointers store dummy values with the correct ids, but nothing else. We need to replace
183        // these dummy values with pointers to the actual parent symbol tables, using the ids to
184        // know which tables should be equal.
185        //
186        // See super::serde::RcRefCellToInner, super::serde::RcRefCellToId.
187
188        // Store the definitive versions of all symbol tables by id.
189        let mut tables: HashMap<ObjId, SymbolTablePtr> = HashMap::new();
190
191        // Find the definitive versions by traversing the sub-models.
192        for submodel in self.submodel.universe() {
193            let table_ptr = submodel.symbols_ptr_unchecked().clone();
194            let id = table_ptr.id();
195
196            // ids should be unique!
197            assert_eq!(tables.insert(id, table_ptr), None);
198        }
199
200        // Restore parent pointers using `tables`.
201        for table in tables.clone().into_values() {
202            let mut table_mut = table.write();
203            let parent_mut = table_mut.parent_mut_unchecked();
204
205            #[allow(clippy::unwrap_used)]
206            if let Some(parent) = parent_mut {
207                let parent_id = parent.id();
208                *parent = tables.get(&parent_id).unwrap().clone();
209            }
210        }
211
212        // The definitive versions of declarations are stored in the symbol table. References store
213        // dummy values with the correct ids, but nothing else.
214
215        // Store the definitive version of all declarations by id.
216        let mut all_declarations: HashMap<ObjId, DeclarationPtr> = HashMap::new();
217        for table in tables.values() {
218            for (_, decl) in table.read().iter_local() {
219                let id = decl.id();
220                all_declarations.insert(id, decl.clone());
221            }
222        }
223
224        // Swizzle declaration pointers in expressions (references, auxdecls) using their ids and `all_declarations`.
225        *self.submodel.constraints_mut() = self.submodel.constraints().transform_bi(&move |decl: DeclarationPtr| {
226                let id = decl.id();
227                        all_declarations
228                            .get(&id)
229                            .unwrap_or_else(|| panic!("A declaration used in the expression tree should exist in the symbol table. The missing declaration has id {id}."))
230                            .clone()
231        });
232
233        Some(Model {
234            submodel: self.submodel,
235            dominance: self.dominance,
236            context,
237            search_order: self.search_order,
238        })
239    }
240}
241
242impl From<Model> for SerdeModel {
243    fn from(val: Model) -> Self {
244        SerdeModel {
245            submodel: val.submodel,
246            dominance: val.dominance,
247            search_order: val.search_order,
248        }
249    }
250}
251
252impl Display for SerdeModel {
253    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
254        std::fmt::Display::fmt(&self.submodel, f)
255    }
256}
257
258impl SerdeModel {
259    /// Collects all ObjId values from the model using uniplate traversal.
260    ///
261    /// Traverses the model structure using `universe_bi` to collect IDs from:
262    /// - All SubModels (to get SymbolTable IDs via `HasId`)
263    /// - All DeclarationPtrs (to get declaration IDs via `HasId`)
264    ///
265    /// Returns a mapping from original ID to stable sequential ID (0, 1, 2, ...).
266    /// IDs are assigned in the order they are encountered during traversal, ensuring
267    /// stability across identical model structures.
268    pub fn collect_stable_id_mapping(&self) -> HashMap<ObjId, ObjId> {
269        fn visit_symbol_table(symbol_table: SymbolTablePtr, id_list: &mut IndexSet<ObjId>) {
270            // If we have seen this table before, all its local declarations were already handled.
271            if !id_list.insert(symbol_table.id()) {
272                return;
273            }
274
275            let table_ref = symbol_table.read();
276            table_ref.iter_local().for_each(|(_, decl)| {
277                id_list.insert(decl.id());
278            });
279        }
280
281        // Using an IndexSet here, we maintain insertion order while deduplicating IDs.
282        let mut id_list: IndexSet<ObjId> = IndexSet::new();
283
284        // Collect SymbolTable IDs by traversing all SubModels
285        for submodel in self.submodel.universe() {
286            visit_symbol_table(submodel.symbols_ptr_unchecked().clone(), &mut id_list);
287        }
288
289        // Collect remaining IDs by traversing the expression tree
290        // (Some expressions, e.g. AbstractComprehensions, contain SymbolTable's not
291        // contained in submodels).
292        let mut exprs: VecDeque<Expression> = self.submodel.universe_bi();
293        if let Some(dominance) = &self.dominance {
294            exprs.push_back(dominance.clone());
295        }
296        for symbol_table in Biplate::<SymbolTablePtr>::universe_bi(&exprs) {
297            visit_symbol_table(symbol_table, &mut id_list);
298        }
299        for declaration in Biplate::<DeclarationPtr>::universe_bi(&exprs) {
300            id_list.insert(declaration.id());
301        }
302
303        // Create stable mapping: original_id -> stable_id
304        let mut id_map = HashMap::new();
305        for (stable_id, original_id) in id_list.into_iter().enumerate() {
306            let type_name = original_id.type_name;
307            id_map.insert(
308                original_id,
309                ObjId {
310                    object_id: stable_id as u32,
311                    type_name,
312                },
313            );
314        }
315
316        id_map
317    }
318}