conjure_core/ast/
symbol_table.rs

Help
1//! The symbol table.
2//!
3//! See the item documentation for [`SymbolTable`] for more details.
4
5use crate::bug;
6use crate::representation::{get_repr_rule, Representation};
7
8use super::comprehension::Comprehension;
9use super::serde::{RcRefCellAsId, RcRefCellAsInner};
10use std::cell::RefCell;
11use std::collections::btree_map::Entry;
12use std::collections::BTreeSet;
13use std::collections::{BTreeMap, VecDeque};
14use std::rc::Rc;
15use std::sync::atomic::{AtomicU32, Ordering};
16
17use super::declaration::Declaration;
18use super::serde::{DefaultWithId, HasId, ObjId};
19use super::types::Typeable;
20use itertools::{izip, Itertools as _};
21use serde::{Deserialize, Serialize};
22use serde_with::serde_as;
23use uniplate::Tree;
24use uniplate::{Biplate, Uniplate};
25
26use super::name::Name;
27use super::{Domain, Expression, ReturnType, SubModel};
28use derivative::Derivative;
29
30// Count symbol tables per thread / model.
31//
32// We run tests in parallel and having the id's be per thread keeps them more deterministic in the
33// JSON output. If this were not thread local, ids would be given to symbol tables differently in
34// each test run (depending on how the threads were scheduled). These id changes would result in
35// all the generated tests "changing" each time `ACCEPT=true cargo test` is ran.
36//
37// SAFETY: Symbol tables use Rc<RefCell<<>>, so a model is not thread-safe anyways.
38thread_local! {
39static ID_COUNTER: AtomicU32 = const { AtomicU32::new(0) };
40}
41
42/// The global symbol table, mapping names to their definitions.
43///
44/// Names in the symbol table are unique, including between different types of object stored in the
45/// symbol table. For example, you cannot have a letting and decision variable with the same name.
46///
47/// # Symbol Kinds
48///
49/// The symbol table tracks the following types of symbol:
50///
51/// ## Decision Variables
52///
53/// ```text
54/// find NAME: DOMAIN
55/// ```
56///
57/// See [`DecisionVariable`](super::DecisionVariable).
58///
59/// ## Lettings
60///
61/// Lettings define constants, of which there are two types:
62///
63///   + **Constant values**: `letting val be A`, where A is an [`Expression`].
64///
65///     A can be any integer, boolean, or matrix expression.
66///     A can include references to other lettings, model parameters, and, unlike Savile Row,
67///     decision variables.
68///
69///   + **Constant domains**: `letting Domain be domain D`, where D is a [`Domain`].
70///
71///     D can include references to other lettings and model parameters, and, unlike Savile Row,
72///     decision variables.
73///
74/// Unless otherwise stated, these follow the semantics specified in section 2.2.2 of the Savile
75/// Row manual (version 1.9.1 at time of writing).
76#[derive(Derivative)]
77#[derivative(PartialEq)]
78#[derive(Debug, Eq)]
79#[serde_as]
80#[derive(Serialize, Deserialize)]
81pub struct SymbolTable {
82    #[serde_as(as = "Vec<(_,RcRefCellAsInner)>")]
83    table: BTreeMap<Name, Rc<Declaration>>,
84
85    /// A unique id for this symbol table, for serialisation and debugging.
86    #[derivative(PartialEq = "ignore")] // eq by value not id.
87    id: ObjId,
88
89    #[serde_as(as = "Option<RcRefCellAsId>")]
90    parent: Option<Rc<RefCell<SymbolTable>>>,
91
92    next_machine_name: RefCell<i32>,
93}
94
95impl SymbolTable {
96    /// Creates an empty symbol table.
97    pub fn new() -> SymbolTable {
98        SymbolTable::new_inner(None)
99    }
100
101    /// Creates an empty symbol table with the given parent.
102    pub fn with_parent(parent: Rc<RefCell<SymbolTable>>) -> SymbolTable {
103        SymbolTable::new_inner(Some(parent))
104    }
105
106    fn new_inner(parent: Option<Rc<RefCell<SymbolTable>>>) -> SymbolTable {
107        let id = ID_COUNTER.with(|x| x.fetch_add(1, Ordering::Relaxed));
108        SymbolTable {
109            id,
110            table: BTreeMap::new(),
111            next_machine_name: RefCell::new(0),
112            parent,
113        }
114    }
115
116    /// Looks up the declaration with the given name in the current scope only.
117    ///
118    /// Returns `None` if there is no declaration with that name in the current scope.
119    pub fn lookup_local(&self, name: &Name) -> Option<Rc<Declaration>> {
120        self.table.get(name).cloned()
121    }
122
123    /// Looks up the declaration with the given name, checking all enclosing scopes.
124    ///
125    /// Returns `None` if there is no declaration with that name in scope.
126    pub fn lookup(&self, name: &Name) -> Option<Rc<Declaration>> {
127        self.lookup_local(name).or_else(|| {
128            self.parent
129                .as_ref()
130                .and_then(|parent| (*parent).borrow().lookup(name))
131        })
132    }
133
134    /// Inserts a declaration into the symbol table.
135    ///
136    /// Returns `None` if there is already a symbol with this name in the local scope.
137    pub fn insert(&mut self, declaration: Rc<Declaration>) -> Option<()> {
138        let name = declaration.name().clone();
139        if let Entry::Vacant(e) = self.table.entry(name) {
140            e.insert(declaration);
141            Some(())
142        } else {
143            None
144        }
145    }
146
147    /// Updates or adds a declaration in the immediate local scope.
148    pub fn update_insert(&mut self, declaration: Rc<Declaration>) {
149        let name = declaration.name().clone();
150        self.table.insert(name, declaration);
151    }
152
153    /// Looks up the return type for name if it has one and is in scope.
154    pub fn return_type(&self, name: &Name) -> Option<ReturnType> {
155        self.lookup(name).and_then(|x| x.return_type())
156    }
157
158    /// Looks up the return type for name if has one and is in the local scope.
159    pub fn return_type_local(&self, name: &Name) -> Option<ReturnType> {
160        self.lookup_local(name).and_then(|x| x.return_type())
161    }
162
163    /// Looks up the domain of name if it has one and is in scope.
164    ///
165    /// This method can return domain references: if a ground domain is always required, use
166    /// [`SymbolTable::resolve_domain`].
167    pub fn domain(&self, name: &Name) -> Option<Domain> {
168        // TODO: do not clone here: in the future, we might want to wrap all domains in Rc's to get
169        // clone-on-write behaviour (saving memory in scenarios such as matrix decomposition where
170        // a lot of the domains would be the same).
171
172        if let Name::WithRepresentation(name, _) = name {
173            let decl = self.lookup(name.as_ref())?;
174            decl.domain().cloned()
175        } else {
176            let decl = self.lookup(name)?;
177            decl.domain().cloned()
178        }
179    }
180
181    /// Looks up the domain of name, resolving domain references to ground domains.
182    ///
183    /// See [`SymbolTable::domain`].
184    pub fn resolve_domain(&self, name: &Name) -> Option<Domain> {
185        self.domain(name).map(|domain| domain.resolve(self))
186    }
187
188    /// Iterates over entries in the local symbol table only.
189    pub fn into_iter_local(self) -> LocalIntoIter {
190        LocalIntoIter {
191            inner: self.table.into_iter(),
192        }
193    }
194
195    /// Extends the symbol table with the given symbol table, updating the gensym counter if
196    /// necessary.
197    pub fn extend(&mut self, other: SymbolTable) {
198        if other.table.keys().count() > self.table.keys().count() {
199            let new_vars = other.table.keys().collect::<BTreeSet<_>>();
200            let old_vars = self.table.keys().collect::<BTreeSet<_>>();
201
202            for added_var in new_vars.difference(&old_vars) {
203                let mut next_var = self.next_machine_name.borrow_mut();
204                if let Name::MachineName(m) = *added_var {
205                    if *m >= *next_var {
206                        *next_var = *m + 1;
207                    }
208                }
209            }
210        }
211
212        self.table.extend(other.table);
213    }
214
215    /// Returns an arbitrary variable name that is not in the symbol table.
216    pub fn gensym(&self) -> Name {
217        let num = *self.next_machine_name.borrow();
218        *(self.next_machine_name.borrow_mut()) += 1;
219        Name::MachineName(num) // incremented when inserted
220    }
221
222    /// Gets the parent of this symbol table as a mutable reference.
223    ///
224    /// This function provides no sanity checks.
225    pub fn parent_mut_unchecked(&mut self) -> &mut Option<Rc<RefCell<SymbolTable>>> {
226        &mut self.parent
227    }
228
229    /// Gets the representation `representation` for `name`.
230    ///
231    /// # Returns
232    ///
233    /// + `None` if `name` does not exist, is not a decision variable, or does not have that representation.
234    pub fn get_representation(
235        &self,
236        name: &Name,
237        representation: &[&str],
238    ) -> Option<Vec<Box<dyn Representation>>> {
239        // TODO: move representation stuff to declaration / variable to avoid cloning? (we have to
240        // move inside of an rc here, so cannot return borrows)
241        //
242        // Also would prevent constant "does exist" "is var" checks.
243        //
244        // The reason it is not there now is because I'm getting serde issues...
245        //
246        // Also might run into issues putting get_or_add into declaration/variable, as that
247        // requires us to mutably borrow both the symbol table, and the variable inside the symbol
248        // table..
249
250        let decl = self.lookup(name)?;
251        let var = decl.as_var()?;
252
253        var.representations
254            .iter()
255            .find(|x| &x.iter().map(|r| r.repr_name()).collect_vec()[..] == representation)
256            .cloned()
257            .clone()
258    }
259
260    /// Gets all initialised representations for `name`.
261    ///
262    /// # Returns
263    ///
264    /// + `None` if `name` does not exist, or is not a decision variable.
265    pub fn representations_for(&self, name: &Name) -> Option<Vec<Vec<Box<dyn Representation>>>> {
266        let decl = self.lookup(name)?;
267        decl.as_var().map(|x| x.representations.clone())
268    }
269
270    /// Gets the representation `representation` for `name`, creating it if it does not exist.
271    ///
272    /// If the representation does not exist, this method initialises the representation in this
273    /// symbol table, adding the representation to `name`, and the declarations for the represented
274    /// variables to the symbol table.
275    ///
276    /// # Usage
277    ///
278    /// Representations for variable references should be selected and created by the
279    /// `select_representation` rule. Therefore, this method should not be used in other rules.
280    /// Consider using [`get_representation`](`SymbolTable::get_representation`) instead.
281    ///
282    /// # Returns
283    ///
284    /// + `None` if `name` does not exist, is not a decision variable, or cannot be given that
285    ///   representation.
286    pub fn get_or_add_representation(
287        &mut self,
288        name: &Name,
289        representation: &[&str],
290    ) -> Option<Vec<Box<dyn Representation>>> {
291        let decl = self.lookup(name)?;
292        let var = decl.as_var()?;
293
294        var.representations
295            .iter()
296            .find(|x| &x.iter().map(|r| r.repr_name()).collect_vec()[..] == representation)
297            .cloned()
298            .clone()
299            .or_else(|| {
300                let mut decl: Declaration = Rc::unwrap_or_clone(decl);
301                let var = decl.as_var_mut()?;
302
303                // TODO: no idea how to deal with initialising representations for nested things..
304
305                if representation.len() != 1 {
306                    bug!("nested representations not implemented")
307                }
308
309                let repr_rule = get_repr_rule(representation[0])?;
310
311                let reprs = vec![repr_rule(name, self)?];
312                for repr in &reprs {
313                    repr.declaration_down()
314                        .ok()
315                        .unwrap()
316                        .iter()
317                        .for_each(|x| self.update_insert(Rc::new(x.clone())));
318                }
319                var.representations.push(reprs.clone());
320
321                self.update_insert(Rc::new(decl));
322                Some(reprs)
323            })
324    }
325}
326
327impl IntoIterator for SymbolTable {
328    type Item = (Name, Rc<Declaration>);
329
330    type IntoIter = IntoIter;
331
332    /// Iterates over symbol table entries in scope.
333    fn into_iter(self) -> Self::IntoIter {
334        IntoIter {
335            inner: self.table.into_iter(),
336            parent: self.parent,
337        }
338    }
339}
340
341/// Iterator over symbol table entries in the current scope only.
342pub struct LocalIntoIter {
343    // iterator over the btreemap
344    inner: std::collections::btree_map::IntoIter<Name, Rc<Declaration>>,
345}
346
347impl Iterator for LocalIntoIter {
348    type Item = (Name, Rc<Declaration>);
349
350    fn next(&mut self) -> Option<Self::Item> {
351        self.inner.next()
352    }
353}
354
355/// Iterator over all symbol table entries in scope.
356pub struct IntoIter {
357    // iterator over the current scopes' btreemap
358    inner: std::collections::btree_map::IntoIter<Name, Rc<Declaration>>,
359
360    // the parent scope
361    parent: Option<Rc<RefCell<SymbolTable>>>,
362}
363
364impl Iterator for IntoIter {
365    type Item = (Name, Rc<Declaration>);
366
367    fn next(&mut self) -> Option<Self::Item> {
368        let mut val = self.inner.next();
369
370        // Go up the tree until we find a parent symbol table with declarations to iterate over.
371        //
372        // Note that the parent symbol table may be empty - this is why this is a loop!
373        while val.is_none() {
374            let parent = self.parent.clone()?;
375            let parent_ref = (*parent).borrow();
376            self.parent = parent_ref.parent.clone();
377            self.inner = parent_ref.table.clone().into_iter();
378
379            val = self.inner.next();
380        }
381
382        val
383    }
384}
385
386impl HasId for SymbolTable {
387    fn id(&self) -> ObjId {
388        self.id
389    }
390}
391
392impl DefaultWithId for SymbolTable {
393    fn default_with_id(id: ObjId) -> Self {
394        Self {
395            table: BTreeMap::new(),
396            id,
397            parent: None,
398            next_machine_name: RefCell::new(0),
399        }
400    }
401}
402
403impl Clone for SymbolTable {
404    fn clone(&self) -> Self {
405        Self {
406            table: self.table.clone(),
407            id: ID_COUNTER.with(|x| x.fetch_add(1, Ordering::Relaxed)),
408            parent: self.parent.clone(),
409            next_machine_name: self.next_machine_name.clone(),
410        }
411    }
412}
413
414impl Default for SymbolTable {
415    fn default() -> Self {
416        Self::new_inner(None)
417    }
418}
419
420impl Uniplate for SymbolTable {
421    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
422        // do not recurse up parents, that would be weird?
423        let self2 = self.clone();
424        (Tree::Zero, Box::new(move |_| self2.clone()))
425    }
426}
427
428impl Biplate<Expression> for SymbolTable {
429    fn biplate(&self) -> (Tree<Expression>, Box<dyn Fn(Tree<Expression>) -> Self>) {
430        let (child_trees, ctxs): (VecDeque<_>, Vec<_>) = self
431            .table
432            .values()
433            .map(|decl| <Declaration as Biplate<Expression>>::biplate(decl.as_ref()))
434            .unzip();
435
436        let tree = Tree::Many(child_trees);
437
438        let self2 = self.clone();
439        let ctx = Box::new(move |tree| {
440            let Tree::Many(exprs) = tree else {
441                panic!("unexpected children structure");
442            };
443
444            let mut self3 = self2.clone();
445            let self3_iter = self3.table.iter_mut();
446            for (ctx, tree, (_, decl)) in izip!(&ctxs, exprs, self3_iter) {
447                // update declaration inside the pointer instead of creating a new one, so all
448                // things referencing this keep referencing this.
449                *decl = Rc::new(ctx(tree));
450            }
451
452            self3
453        });
454
455        (tree, ctx)
456    }
457}
458
459impl Biplate<Comprehension> for SymbolTable {
460    fn biplate(
461        &self,
462    ) -> (
463        Tree<Comprehension>,
464        Box<dyn Fn(Tree<Comprehension>) -> Self>,
465    ) {
466        let (expr_tree, expr_ctx) = <SymbolTable as Biplate<Expression>>::biplate(self);
467        let (comprehension_tree, comprehension_ctx) =
468            <VecDeque<Expression> as Biplate<Comprehension>>::biplate(
469                &expr_tree.into_iter().collect(),
470            );
471
472        let ctx = Box::new(move |x| {
473            expr_ctx(Tree::Many(
474                comprehension_ctx(x).into_iter().map(Tree::One).collect(),
475            ))
476        });
477
478        (comprehension_tree, ctx)
479    }
480}
481
482impl Biplate<SubModel> for SymbolTable {
483    // walk into expressions
484    fn biplate(&self) -> (Tree<SubModel>, Box<dyn Fn(Tree<SubModel>) -> Self>) {
485        let (exprs, exprs_ctx) = <SymbolTable as Biplate<Expression>>::biplate(self);
486        let (submodel_tree, submodel_ctx) =
487            <VecDeque<Expression> as Biplate<SubModel>>::biplate(&exprs.into_iter().collect());
488
489        let ctx = Box::new(move |x| {
490            exprs_ctx(Tree::Many(
491                submodel_ctx(x).into_iter().map(Tree::One).collect(),
492            ))
493        });
494        (submodel_tree, ctx)
495    }
496}
conjure_core/ast/symbol_table.rs

conjure_core/ast/
symbol_table.rs
