1
use 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
};
10
use serde::{Deserialize, Serialize};
11
use serde_with::serde_as;
12
use uniplate::{Biplate, Tree, Uniplate};
13

            
14
use crate::{bug, metadata::Metadata};
15
use std::{
16
    cell::{Ref, RefCell, RefMut},
17
    collections::VecDeque,
18
    fmt::Display,
19
    rc::Rc,
20
};
21

            
22
use 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)]
32
pub struct SubModel {
33
    constraints: Expression,
34
    #[serde_as(as = "RcRefCellAsInner")]
35
    symbols: Rc<RefCell<SymbolTable>>,
36
}
37

            
38
impl 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
2196
    pub(super) fn new_top_level() -> SubModel {
45
2196
        SubModel {
46
2196
            constraints: Expression::Root(Metadata::new(), vec![]),
47
2196
            symbols: Rc::new(RefCell::new(SymbolTable::new())),
48
2196
        }
49
2196
    }
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
5382
    pub fn symbols_ptr_unchecked(&self) -> &Rc<RefCell<SymbolTable>> {
66
5382
        &self.symbols
67
5382
    }
68

            
69
    /// The symbol table for this sub-model as a reference.
70
18692694
    pub fn symbols(&self) -> Ref<SymbolTable> {
71
18692694
        (*self.symbols).borrow()
72
18692694
    }
73

            
74
    /// The symbol table for this sub-model as a mutable reference.
75
23382
    pub fn symbols_mut(&mut self) -> RefMut<SymbolTable> {
76
23382
        (*self.symbols).borrow_mut()
77
23382
    }
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
162180
    pub fn root(&self) -> &Expression {
84
162180
        &self.constraints
85
162180
    }
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
17784
    fn root_mut_unchecked(&mut self) -> &mut Expression {
92
17784
        &mut self.constraints
93
17784
    }
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
17784
    pub fn replace_root(&mut self, new_root: Expression) -> Expression {
101
17784
        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
17784
        std::mem::replace(self.root_mut_unchecked(), new_root)
107
17784
    }
108

            
109
    /// The top-level constraints in this sub-model.
110
5526
    pub fn constraints(&self) -> &Vec<Expression> {
111
5526
        let Expression::Root(_, constraints) = &self.constraints else {
112
            bug!("The top level expression in a submodel should be Expr::Root");
113
        };
114

            
115
5526
        constraints
116
5526
    }
117

            
118
    /// The top-level constraints in this sub-model as a mutable vector.
119
19980
    pub fn constraints_mut(&mut self) -> &mut Vec<Expression> {
120
19980
        let Expression::Root(_, constraints) = &mut self.constraints else {
121
            bug!("The top level expression in a submodel should be Expr::Root");
122
        };
123

            
124
19980
        constraints
125
19980
    }
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
198
    pub fn add_constraint(&mut self, constraint: Expression) {
134
198
        self.constraints_mut().push(constraint);
135
198
    }
136

            
137
    /// Adds top-level constraints.
138
19710
    pub fn add_constraints(&mut self, constraints: Vec<Expression>) {
139
19710
        self.constraints_mut().extend(constraints);
140
19710
    }
141

            
142
    /// Adds a new symbol to the symbol table
143
    /// (Wrapper over `SymbolTable.insert`)
144
306
    pub fn add_symbol(&mut self, sym: Declaration) -> Option<()> {
145
306
        self.symbols_mut().insert(Rc::new(sym))
146
306
    }
147
}
148

            
149
impl Typeable for SubModel {
150
    fn return_type(&self) -> Option<super::ReturnType> {
151
        Some(ReturnType::Bool)
152
    }
153
}
154

            
155
impl Display for SubModel {
156
    #[allow(clippy::unwrap_used)] // [rustdocs]: should only fail iff the formatter fails
157
3564
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
158
10944
        for (name, decl) in self.symbols().clone().into_iter_local() {
159
10944
            match decl.kind() {
160
                DeclarationKind::DecisionVariable(_) => {
161
10548
                    writeln!(
162
10548
                        f,
163
10548
                        "{}",
164
10548
                        pretty_variable_declaration(&self.symbols(), &name).unwrap()
165
10548
                    )?;
166
                }
167
                DeclarationKind::ValueLetting(_) => {
168
216
                    writeln!(
169
216
                        f,
170
216
                        "{}",
171
216
                        pretty_value_letting_declaration(&self.symbols(), &name).unwrap()
172
216
                    )?;
173
                }
174
                DeclarationKind::DomainLetting(_) => {
175
180
                    writeln!(
176
180
                        f,
177
180
                        "{}",
178
180
                        pretty_domain_letting_declaration(&self.symbols(), &name).unwrap()
179
180
                    )?;
180
                }
181
            }
182
        }
183

            
184
3564
        writeln!(f, "\nsuch that\n")?;
185

            
186
3564
        writeln!(f, "{}", pretty_expressions_as_top_level(self.constraints()))?;
187

            
188
3564
        Ok(())
189
3564
    }
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

            
205
impl Uniplate for SubModel {
206
24984
    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
207
24984
        // Look inside constraint tree and symbol tables.
208
24984

            
209
24984
        let (expr_tree, expr_ctx) = <Expression as Biplate<SubModel>>::biplate(self.root());
210
24984

            
211
24984
        let symtab_ptr = self.symbols();
212
24984
        let (symtab_tree, symtab_ctx) = <SymbolTable as Biplate<SubModel>>::biplate(&symtab_ptr);
213
24984

            
214
24984
        let tree = Tree::Many(VecDeque::from([expr_tree, symtab_tree]));
215
24984

            
216
24984
        let self2 = self.clone();
217
24984
        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
24984
        });
237
24984

            
238
24984
        (tree, ctx)
239
24984
    }
240
}
241

            
242
impl 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

            
282
impl 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
}