conjure_core/ast/
submodel.rs

1use super::{
2    declaration::DeclarationKind,
3    pretty::{
4        pretty_domain_letting_declaration, pretty_expressions_as_top_level,
5        pretty_value_letting_declaration, pretty_variable_declaration,
6    },
7    serde::RcRefCellAsInner,
8    Declaration,
9};
10use serde::{Deserialize, Serialize};
11use serde_with::serde_as;
12use uniplate::{Biplate, Tree, Uniplate};
13
14use crate::{bug, metadata::Metadata};
15use std::{
16    cell::{Ref, RefCell, RefMut},
17    collections::VecDeque,
18    fmt::Display,
19    rc::Rc,
20};
21
22use super::{types::Typeable, Expression, ReturnType, SymbolTable};
23
24/// A sub-model, representing a lexical scope in the model.
25///
26/// Each sub-model contains a symbol table representing its scope, as well as a expression tree.
27///
28/// The expression tree is formed of a root node of type [`Expression::Root`], which contains a
29/// vector of top-level constraints.
30#[serde_as]
31#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
32pub struct SubModel {
33    constraints: Expression,
34    #[serde_as(as = "RcRefCellAsInner")]
35    symbols: Rc<RefCell<SymbolTable>>,
36}
37
38impl SubModel {
39    /// Creates a new [`Submodel`] with no parent scope.
40    ///
41    /// Top level models are represented as [`Model`](super::model): consider using
42    /// [`Model::new`](super::Model::new) instead.
43    #[doc(hidden)]
44    pub(super) fn new_top_level() -> SubModel {
45        SubModel {
46            constraints: Expression::Root(Metadata::new(), vec![]),
47            symbols: Rc::new(RefCell::new(SymbolTable::new())),
48        }
49    }
50
51    /// Creates a new [`Submodel`] as a child scope of `parent`.
52    ///
53    /// `parent` should be the symbol table of the containing scope of this sub-model.
54    pub fn new(parent: Rc<RefCell<SymbolTable>>) -> SubModel {
55        SubModel {
56            constraints: Expression::Root(Metadata::new(), vec![]),
57            symbols: Rc::new(RefCell::new(SymbolTable::with_parent(parent))),
58        }
59    }
60
61    /// The symbol table for this sub-model as a pointer.
62    ///
63    /// The caller should only mutate the returned symbol table if this method was called on a
64    /// mutable model.
65    pub fn symbols_ptr_unchecked(&self) -> &Rc<RefCell<SymbolTable>> {
66        &self.symbols
67    }
68
69    /// The symbol table for this sub-model as a reference.
70    pub fn symbols(&self) -> Ref<SymbolTable> {
71        (*self.symbols).borrow()
72    }
73
74    /// The symbol table for this sub-model as a mutable reference.
75    pub fn symbols_mut(&mut self) -> RefMut<SymbolTable> {
76        (*self.symbols).borrow_mut()
77    }
78
79    /// The root node of this sub-model.
80    ///
81    /// The root node is an [`Expression::Root`] containing a vector of the top level constraints
82    /// in this sub-model.
83    pub fn root(&self) -> &Expression {
84        &self.constraints
85    }
86
87    /// The root node of this sub-model, as a mutable reference.
88    ///
89    /// The caller is responsible for ensuring that the root node remains an [`Expression::Root`].
90    ///
91    fn root_mut_unchecked(&mut self) -> &mut Expression {
92        &mut self.constraints
93    }
94
95    /// Replaces the root node with `new_root`, returning the old root node.
96    ///
97    /// # Panics
98    ///
99    /// - If `new_root` is not an [`Expression::Root`].
100    pub fn replace_root(&mut self, new_root: Expression) -> Expression {
101        let Expression::Root(_, _) = new_root else {
102            panic!("new_root is not an Expression::Root");
103        };
104
105        // INVARIANT: already checked that `new_root` is an [`Expression::Root`]
106        std::mem::replace(self.root_mut_unchecked(), new_root)
107    }
108
109    /// The top-level constraints in this sub-model.
110    pub fn constraints(&self) -> &Vec<Expression> {
111        let Expression::Root(_, constraints) = &self.constraints else {
112            bug!("The top level expression in a submodel should be Expr::Root");
113        };
114
115        constraints
116    }
117
118    /// The top-level constraints in this sub-model as a mutable vector.
119    pub fn constraints_mut(&mut self) -> &mut Vec<Expression> {
120        let Expression::Root(_, constraints) = &mut self.constraints else {
121            bug!("The top level expression in a submodel should be Expr::Root");
122        };
123
124        constraints
125    }
126
127    /// Replaces the top-level constraints with `new_constraints`, returning the old ones.
128    pub fn replace_constraints(&mut self, new_constraints: Vec<Expression>) -> Vec<Expression> {
129        std::mem::replace(self.constraints_mut(), new_constraints)
130    }
131
132    /// Adds a top-level constraint.
133    pub fn add_constraint(&mut self, constraint: Expression) {
134        self.constraints_mut().push(constraint);
135    }
136
137    /// Adds top-level constraints.
138    pub fn add_constraints(&mut self, constraints: Vec<Expression>) {
139        self.constraints_mut().extend(constraints);
140    }
141
142    /// Adds a new symbol to the symbol table
143    /// (Wrapper over `SymbolTable.insert`)
144    pub fn add_symbol(&mut self, sym: Declaration) -> Option<()> {
145        self.symbols_mut().insert(Rc::new(sym))
146    }
147}
148
149impl Typeable for SubModel {
150    fn return_type(&self) -> Option<super::ReturnType> {
151        Some(ReturnType::Bool)
152    }
153}
154
155impl Display for SubModel {
156    #[allow(clippy::unwrap_used)] // [rustdocs]: should only fail iff the formatter fails
157    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
158        for (name, decl) in self.symbols().clone().into_iter_local() {
159            match decl.kind() {
160                DeclarationKind::DecisionVariable(_) => {
161                    writeln!(
162                        f,
163                        "{}",
164                        pretty_variable_declaration(&self.symbols(), &name).unwrap()
165                    )?;
166                }
167                DeclarationKind::ValueLetting(_) => {
168                    writeln!(
169                        f,
170                        "{}",
171                        pretty_value_letting_declaration(&self.symbols(), &name).unwrap()
172                    )?;
173                }
174                DeclarationKind::DomainLetting(_) => {
175                    writeln!(
176                        f,
177                        "{}",
178                        pretty_domain_letting_declaration(&self.symbols(), &name).unwrap()
179                    )?;
180                }
181            }
182        }
183
184        writeln!(f, "\nsuch that\n")?;
185
186        writeln!(f, "{}", pretty_expressions_as_top_level(self.constraints()))?;
187
188        Ok(())
189    }
190}
191
192// Using manual implementations of Uniplate so that we can update the old Rc<RefCell<<>>> with the
193// new value instead of creating a new one. This will keep the parent pointers in sync.
194//
195// I considered adding Rc RefCell shared-mutability to Uniplate, but I think this is unsound in
196// generality: e.g. two pointers to the same object are in our tree, and both get modified in
197// different ways.
198//
199// Shared mutability is probably fine here, as we only have one pointer to each symbol table
200// reachable via Uniplate, the one in its Submodel. The SymbolTable implementation doesn't return
201// or traverse through the parent pointers.
202//
203// -- nd60
204
205impl Uniplate for SubModel {
206    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
207        // Look inside constraint tree and symbol tables.
208
209        let (expr_tree, expr_ctx) = <Expression as Biplate<SubModel>>::biplate(self.root());
210
211        let symtab_ptr = self.symbols();
212        let (symtab_tree, symtab_ctx) = <SymbolTable as Biplate<SubModel>>::biplate(&symtab_ptr);
213
214        let tree = Tree::Many(VecDeque::from([expr_tree, symtab_tree]));
215
216        let self2 = self.clone();
217        let ctx = Box::new(move |x| {
218            let Tree::Many(xs) = x else {
219                panic!();
220            };
221
222            let root = expr_ctx(xs[0].clone());
223            let symtab = symtab_ctx(xs[1].clone());
224
225            let mut self3 = self2.clone();
226
227            let Expression::Root(_, _) = root else {
228                bug!("root expression not root");
229            };
230
231            *self3.root_mut_unchecked() = root;
232
233            *self3.symbols_mut() = symtab;
234
235            self3
236        });
237
238        (tree, ctx)
239    }
240}
241
242impl Biplate<Expression> for SubModel {
243    fn biplate(&self) -> (Tree<Expression>, Box<dyn Fn(Tree<Expression>) -> Self>) {
244        // Return constraints tree and look inside symbol table.
245        let symtab_ptr = self.symbols();
246        let (symtab_tree, symtab_ctx) = <SymbolTable as Biplate<Expression>>::biplate(&symtab_ptr);
247
248        let tree = Tree::Many(VecDeque::from([
249            Tree::One(self.root().clone()),
250            symtab_tree,
251        ]));
252
253        let self2 = self.clone();
254        let ctx = Box::new(move |x| {
255            let Tree::Many(xs) = x else {
256                panic!();
257            };
258
259            let Tree::One(root) = xs[0].clone() else {
260                panic!();
261            };
262
263            let symtab = symtab_ctx(xs[1].clone());
264
265            let mut self3 = self2.clone();
266
267            let Expression::Root(_, _) = root else {
268                bug!("root expression not root");
269            };
270
271            *self3.root_mut_unchecked() = root;
272
273            *self3.symbols_mut() = symtab;
274
275            self3
276        });
277
278        (tree, ctx)
279    }
280}
281
282impl Biplate<SubModel> for SubModel {
283    fn biplate(&self) -> (Tree<SubModel>, Box<dyn Fn(Tree<SubModel>) -> Self>) {
284        (
285            Tree::One(self.clone()),
286            Box::new(move |x| {
287                let Tree::One(x) = x else {
288                    panic!();
289                };
290                x
291            }),
292        )
293    }
294}