1
//! The symbol table.
2
//!
3
//! See the item documentation for [`SymbolTable`] for more details.
4

            
5
use super::serde::{RcRefCellAsId, RcRefCellAsInner};
6
use std::cell::RefCell;
7
use std::collections::btree_map::Entry;
8
use std::collections::BTreeSet;
9
use std::collections::{BTreeMap, VecDeque};
10
use std::rc::Rc;
11
use std::sync::atomic::{AtomicU32, Ordering};
12

            
13
use super::declaration::Declaration;
14
use super::serde::{DefaultWithId, HasId, ObjId};
15
use super::types::Typeable;
16
use itertools::izip;
17
use serde::{Deserialize, Serialize};
18
use serde_with::serde_as;
19
use uniplate::Tree;
20
use uniplate::{Biplate, Uniplate};
21

            
22
use super::name::Name;
23
use super::{Domain, Expression, ReturnType};
24
use derivative::Derivative;
25

            
26
// Count symbol tables per thread / model.
27
//
28
// We run tests in parallel and having the id's be per thread keeps them more deterministic in the
29
// JSON output. If this were not thread local, ids would be given to symbol tables differently in
30
// each test run (depending on how the threads were scheduled). These id changes would result in
31
// all the generated tests "changing" each time `ACCEPT=true cargo test` is ran.
32
//
33
// SAFETY: Symbol tables use Rc<RefCell<<>>, so a model is not thread-safe anyways.
34
thread_local! {
35
static ID_COUNTER: AtomicU32 = const { AtomicU32::new(0) };
36
}
37

            
38
/// The global symbol table, mapping names to their definitions.
39
///
40
/// Names in the symbol table are unique, including between different types of object stored in the
41
/// symbol table. For example, you cannot have a letting and decision variable with the same name.
42
///
43
/// # Symbol Kinds
44
///
45
/// The symbol table tracks the following types of symbol:
46
///
47
/// ## Decision Variables
48
///
49
/// ```text
50
/// find NAME: DOMAIN
51
/// ```
52
///
53
/// See [`DecisionVariable`](super::DecisionVariable).
54
///
55
/// ## Lettings
56
///
57
/// Lettings define constants, of which there are two types:
58
///
59
///   + **Constant values**: `letting val be A`, where A is an [`Expression`].
60
///
61
///     A can be any integer, boolean, or matrix expression.
62
///     A can include references to other lettings, model parameters, and, unlike Savile Row,
63
///     decision variables.
64
///
65
///   + **Constant domains**: `letting Domain be domain D`, where D is a [`Domain`].
66
///
67
///     D can include references to other lettings and model parameters, and, unlike Savile Row,
68
///     decision variables.
69
///
70
/// Unless otherwise stated, these follow the semantics specified in section 2.2.2 of the Savile
71
/// Row manual (version 1.9.1 at time of writing).
72
#[derive(Derivative)]
73
#[derivative(PartialEq)]
74
#[derive(Debug, Eq)]
75
#[serde_as]
76
#[derive(Serialize, Deserialize)]
77
pub struct SymbolTable {
78
    #[serde_as(as = "Vec<(_,RcRefCellAsInner)>")]
79
    table: BTreeMap<Name, Rc<Declaration>>,
80

            
81
    /// A unique id for this symbol table, for serialisation and debugging.
82
    #[derivative(PartialEq = "ignore")] // eq by value not id.
83
    id: ObjId,
84

            
85
    #[serde_as(as = "Option<RcRefCellAsId>")]
86
    parent: Option<Rc<RefCell<SymbolTable>>>,
87

            
88
    next_machine_name: RefCell<i32>,
89
}
90

            
91
impl SymbolTable {
92
    /// Creates an empty symbol table.
93
14984
    pub fn new() -> SymbolTable {
94
14984
        SymbolTable::new_inner(None)
95
14984
    }
96

            
97
    /// Creates an empty symbol table with the given parent.
98
    pub fn with_parent(parent: Rc<RefCell<SymbolTable>>) -> SymbolTable {
99
        SymbolTable::new_inner(Some(parent))
100
    }
101

            
102
18061
    fn new_inner(parent: Option<Rc<RefCell<SymbolTable>>>) -> SymbolTable {
103
18061
        let id = ID_COUNTER.with(|x| x.fetch_add(1, Ordering::Relaxed));
104
18061
        SymbolTable {
105
18061
            id,
106
18061
            table: BTreeMap::new(),
107
18061
            next_machine_name: RefCell::new(0),
108
18061
            parent,
109
18061
        }
110
18061
    }
111

            
112
    /// Looks up the declaration with the given name in the current scope only.
113
    ///
114
    /// Returns `None` if there is no declaration with that name in the current scope.
115
162407
    pub fn lookup_local(&self, name: &Name) -> Option<Rc<Declaration>> {
116
162407
        self.table.get(name).cloned()
117
162407
    }
118

            
119
    /// Looks up the declaration with the given name, checking all enclosing scopes.
120
    ///
121
    /// Returns `None` if there is no declaration with that name in scope.
122
162407
    pub fn lookup(&self, name: &Name) -> Option<Rc<Declaration>> {
123
162407
        self.lookup_local(name).or_else(|| {
124
120
            self.parent
125
120
                .as_ref()
126
120
                .and_then(|parent| (*parent).borrow().lookup(name))
127
162407
        })
128
162407
    }
129

            
130
    /// Inserts a declaration into the symbol table.
131
    ///
132
    /// Returns `None` if there is already a symbol with this name in the local scope.
133
10696
    pub fn insert(&mut self, declaration: Rc<Declaration>) -> Option<()> {
134
10696
        let name = declaration.name().clone();
135
10696
        if let Entry::Vacant(e) = self.table.entry(name) {
136
10696
            e.insert(declaration);
137
10696
            Some(())
138
        } else {
139
            None
140
        }
141
10696
    }
142

            
143
    /// Updates or adds a declaration in the immediate local scope.
144
3768747
    pub fn update_insert(&mut self, declaration: Rc<Declaration>) {
145
3768747
        let name = declaration.name().clone();
146
3768747
        self.table.insert(name, declaration);
147
3768747
    }
148

            
149
    /// Looks up the return type for name if it has one and is in scope.
150
    pub fn return_type(&self, name: &Name) -> Option<ReturnType> {
151
        self.lookup(name).and_then(|x| x.return_type())
152
    }
153

            
154
    /// Looks up the return type for name if has one and is in the local scope.
155
    pub fn return_type_local(&self, name: &Name) -> Option<ReturnType> {
156
        self.lookup_local(name).and_then(|x| x.return_type())
157
    }
158

            
159
    /// Looks up the domain of name if it has one and is in scope.
160
3321
    pub fn domain(&self, name: &Name) -> Option<Domain> {
161
        // TODO: do not clone here: in the future, we might want to wrap all domains in Rc's to get
162
        // clone-on-write behaviour (saving memory in scenarios such as matrix decomposition where
163
        // a lot of the domains would be the same).
164
3321
        let decl = self.lookup(name)?;
165

            
166
3320
        decl.domain().cloned()
167
3321
    }
168

            
169
    /// Iterates over entries in the local symbol table only.
170
41497
    pub fn into_iter_local(self) -> LocalIntoIter {
171
41497
        LocalIntoIter {
172
41497
            inner: self.table.into_iter(),
173
41497
        }
174
41497
    }
175

            
176
    /// Extends the symbol table with the given symbol table, updating the gensym counter if
177
    /// necessary.
178
12410
    pub fn extend(&mut self, other: SymbolTable) {
179
12410
        if other.table.keys().count() > self.table.keys().count() {
180
1241
            let new_vars = other.table.keys().collect::<BTreeSet<_>>();
181
1241
            let old_vars = self.table.keys().collect::<BTreeSet<_>>();
182

            
183
1360
            for added_var in new_vars.difference(&old_vars) {
184
1360
                let mut next_var = self.next_machine_name.borrow_mut();
185
1360
                match *added_var {
186
68
                    Name::UserName(_) => {}
187
1292
                    Name::MachineName(m) => {
188
1292
                        if *m >= *next_var {
189
1292
                            *next_var = *m + 1;
190
1292
                        }
191
                    }
192
                }
193
            }
194
11169
        }
195

            
196
12410
        self.table.extend(other.table);
197
12410
    }
198

            
199
    /// Returns an arbitrary variable name that is not in the symbol table.
200
1972
    pub fn gensym(&self) -> Name {
201
1972
        let num = *self.next_machine_name.borrow();
202
1972
        *(self.next_machine_name.borrow_mut()) += 1;
203
1972
        Name::MachineName(num) // incremented when inserted
204
1972
    }
205
}
206

            
207
impl IntoIterator for SymbolTable {
208
    type Item = (Name, Rc<Declaration>);
209

            
210
    type IntoIter = IntoIter;
211

            
212
    /// Iterates over symbol table entries in scope.
213
206601
    fn into_iter(self) -> Self::IntoIter {
214
206601
        IntoIter {
215
206601
            inner: self.table.into_iter(),
216
206601
            parent: self.parent,
217
206601
        }
218
206601
    }
219
}
220

            
221
/// Iterator over symbol table entries in the current scope only.
222
pub struct LocalIntoIter {
223
    // iterator over the btreemap
224
    inner: std::collections::btree_map::IntoIter<Name, Rc<Declaration>>,
225
}
226

            
227
impl Iterator for LocalIntoIter {
228
    type Item = (Name, Rc<Declaration>);
229

            
230
247180
    fn next(&mut self) -> Option<Self::Item> {
231
247180
        self.inner.next()
232
247180
    }
233
}
234

            
235
/// Iterator over all symbol table entries in scope.
236
pub struct IntoIter {
237
    // iterator over the current scopes' btreemap
238
    inner: std::collections::btree_map::IntoIter<Name, Rc<Declaration>>,
239

            
240
    // the parent scope
241
    parent: Option<Rc<RefCell<SymbolTable>>>,
242
}
243

            
244
impl Iterator for IntoIter {
245
    type Item = (Name, Rc<Declaration>);
246

            
247
3976623
    fn next(&mut self) -> Option<Self::Item> {
248
3976623
        let mut val = self.inner.next();
249

            
250
        // Go up the tree until we find a parent symbol table with declarations to iterate over.
251
        //
252
        // Note that the parent symbol table may be empty - this is why this is a loop!
253
3976623
        while val.is_none() {
254
206601
            let parent = self.parent.clone()?;
255
            let parent_ref = (*parent).borrow();
256
            self.parent = parent_ref.parent.clone();
257
            self.inner = parent_ref.table.clone().into_iter();
258

            
259
            val = self.inner.next();
260
        }
261

            
262
3770022
        val
263
3976623
    }
264
}
265

            
266
impl HasId for SymbolTable {
267
    fn id(&self) -> ObjId {
268
        self.id
269
    }
270
}
271

            
272
impl DefaultWithId for SymbolTable {
273
    fn default_with_id(id: ObjId) -> Self {
274
        Self {
275
            table: BTreeMap::new(),
276
            id,
277
            parent: None,
278
            next_machine_name: RefCell::new(0),
279
        }
280
    }
281
}
282

            
283
impl Clone for SymbolTable {
284
623628
    fn clone(&self) -> Self {
285
623628
        Self {
286
623628
            table: self.table.clone(),
287
623628
            id: ID_COUNTER.with(|x| x.fetch_add(1, Ordering::Relaxed)),
288
623628
            parent: self.parent.clone(),
289
623628
            next_machine_name: self.next_machine_name.clone(),
290
623628
        }
291
623628
    }
292
}
293

            
294
impl Default for SymbolTable {
295
3077
    fn default() -> Self {
296
3077
        Self::new_inner(None)
297
3077
    }
298
}
299

            
300
impl Uniplate for SymbolTable {
301
    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
302
        // do not recurse up parents, that would be weird?
303
        let self2 = self.clone();
304
        (Tree::Zero, Box::new(move |_| self2.clone()))
305
    }
306
}
307

            
308
impl Biplate<Expression> for SymbolTable {
309
102646
    fn biplate(&self) -> (Tree<Expression>, Box<dyn Fn(Tree<Expression>) -> Self>) {
310
102646
        let (child_trees, ctxs): (VecDeque<_>, Vec<_>) = self
311
102646
            .table
312
102646
            .values()
313
728178
            .map(|decl| <Declaration as Biplate<Expression>>::biplate(decl.as_ref()))
314
102646
            .unzip();
315
102646

            
316
102646
        let tree = Tree::Many(child_trees);
317
102646

            
318
102646
        let self2 = self.clone();
319
102646
        let ctx = Box::new(move |tree| {
320
12308
            let Tree::Many(exprs) = tree else {
321
                panic!("unexpected children structure");
322
            };
323

            
324
12308
            let mut self3 = self2.clone();
325
12308
            let self3_iter = self3.table.iter_mut();
326
88876
            for (ctx, tree, (_, decl)) in izip!(&ctxs, exprs, self3_iter) {
327
88876
                // update declaration inside the pointer instead of creating a new one, so all
328
88876
                // things referencing this keep referencing this.
329
88876
                *decl = Rc::new(ctx(tree));
330
88876
            }
331

            
332
12308
            self3
333
102646
        });
334
102646

            
335
102646
        (tree, ctx)
336
102646
    }
337
}