conjure_core/ast/
symbol_table.rs

1//! The symbol table.
2//!
3//! See the item documentation for [`SymbolTable`] for more details.
4
5use super::serde::{RcRefCellAsId, RcRefCellAsInner};
6use std::cell::RefCell;
7use std::collections::btree_map::Entry;
8use std::collections::BTreeSet;
9use std::collections::{BTreeMap, VecDeque};
10use std::rc::Rc;
11use std::sync::atomic::{AtomicU32, Ordering};
12
13use super::declaration::Declaration;
14use super::serde::{DefaultWithId, HasId, ObjId};
15use super::types::Typeable;
16use itertools::izip;
17use serde::{Deserialize, Serialize};
18use serde_with::serde_as;
19use uniplate::Tree;
20use uniplate::{Biplate, Uniplate};
21
22use super::name::Name;
23use super::{Domain, Expression, ReturnType, SubModel};
24use 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.
34thread_local! {
35static 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)]
77pub 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
91impl SymbolTable {
92    /// Creates an empty symbol table.
93    pub fn new() -> SymbolTable {
94        SymbolTable::new_inner(None)
95    }
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    fn new_inner(parent: Option<Rc<RefCell<SymbolTable>>>) -> SymbolTable {
103        let id = ID_COUNTER.with(|x| x.fetch_add(1, Ordering::Relaxed));
104        SymbolTable {
105            id,
106            table: BTreeMap::new(),
107            next_machine_name: RefCell::new(0),
108            parent,
109        }
110    }
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    pub fn lookup_local(&self, name: &Name) -> Option<Rc<Declaration>> {
116        self.table.get(name).cloned()
117    }
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    pub fn lookup(&self, name: &Name) -> Option<Rc<Declaration>> {
123        self.lookup_local(name).or_else(|| {
124            self.parent
125                .as_ref()
126                .and_then(|parent| (*parent).borrow().lookup(name))
127        })
128    }
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    pub fn insert(&mut self, declaration: Rc<Declaration>) -> Option<()> {
134        let name = declaration.name().clone();
135        if let Entry::Vacant(e) = self.table.entry(name) {
136            e.insert(declaration);
137            Some(())
138        } else {
139            None
140        }
141    }
142
143    /// Updates or adds a declaration in the immediate local scope.
144    pub fn update_insert(&mut self, declaration: Rc<Declaration>) {
145        let name = declaration.name().clone();
146        self.table.insert(name, declaration);
147    }
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    ///
161    /// This method can return domain references: if a ground domain is always required, use
162    /// [`SymbolTable::resolve_domain`].
163    pub fn domain(&self, name: &Name) -> Option<Domain> {
164        // TODO: do not clone here: in the future, we might want to wrap all domains in Rc's to get
165        // clone-on-write behaviour (saving memory in scenarios such as matrix decomposition where
166        // a lot of the domains would be the same).
167        let decl = self.lookup(name)?;
168
169        decl.domain().cloned()
170    }
171
172    /// Looks up the domain of name, resolving domain references to ground domains.
173    ///
174    /// See [`SymbolTable::domain`].
175    pub fn resolve_domain(&self, name: &Name) -> Option<Domain> {
176        match self.domain(name) {
177            Some(Domain::DomainReference(name)) => self
178                .lookup(&name)
179                .and_then(|decl| decl.as_domain_letting().cloned()),
180            result => result,
181        }
182    }
183
184    /// Iterates over entries in the local symbol table only.
185    pub fn into_iter_local(self) -> LocalIntoIter {
186        LocalIntoIter {
187            inner: self.table.into_iter(),
188        }
189    }
190
191    /// Extends the symbol table with the given symbol table, updating the gensym counter if
192    /// necessary.
193    pub fn extend(&mut self, other: SymbolTable) {
194        if other.table.keys().count() > self.table.keys().count() {
195            let new_vars = other.table.keys().collect::<BTreeSet<_>>();
196            let old_vars = self.table.keys().collect::<BTreeSet<_>>();
197
198            for added_var in new_vars.difference(&old_vars) {
199                let mut next_var = self.next_machine_name.borrow_mut();
200                match *added_var {
201                    Name::UserName(_) => {}
202                    Name::MachineName(m) => {
203                        if *m >= *next_var {
204                            *next_var = *m + 1;
205                        }
206                    }
207                }
208            }
209        }
210
211        self.table.extend(other.table);
212    }
213
214    /// Returns an arbitrary variable name that is not in the symbol table.
215    pub fn gensym(&self) -> Name {
216        let num = *self.next_machine_name.borrow();
217        *(self.next_machine_name.borrow_mut()) += 1;
218        Name::MachineName(num) // incremented when inserted
219    }
220
221    /// Gets the parent of this symbol table as a mutable reference.
222    ///
223    /// This function provides no sanity checks.
224    pub fn parent_mut_unchecked(&mut self) -> &mut Option<Rc<RefCell<SymbolTable>>> {
225        &mut self.parent
226    }
227}
228
229impl IntoIterator for SymbolTable {
230    type Item = (Name, Rc<Declaration>);
231
232    type IntoIter = IntoIter;
233
234    /// Iterates over symbol table entries in scope.
235    fn into_iter(self) -> Self::IntoIter {
236        IntoIter {
237            inner: self.table.into_iter(),
238            parent: self.parent,
239        }
240    }
241}
242
243/// Iterator over symbol table entries in the current scope only.
244pub struct LocalIntoIter {
245    // iterator over the btreemap
246    inner: std::collections::btree_map::IntoIter<Name, Rc<Declaration>>,
247}
248
249impl Iterator for LocalIntoIter {
250    type Item = (Name, Rc<Declaration>);
251
252    fn next(&mut self) -> Option<Self::Item> {
253        self.inner.next()
254    }
255}
256
257/// Iterator over all symbol table entries in scope.
258pub struct IntoIter {
259    // iterator over the current scopes' btreemap
260    inner: std::collections::btree_map::IntoIter<Name, Rc<Declaration>>,
261
262    // the parent scope
263    parent: Option<Rc<RefCell<SymbolTable>>>,
264}
265
266impl Iterator for IntoIter {
267    type Item = (Name, Rc<Declaration>);
268
269    fn next(&mut self) -> Option<Self::Item> {
270        let mut val = self.inner.next();
271
272        // Go up the tree until we find a parent symbol table with declarations to iterate over.
273        //
274        // Note that the parent symbol table may be empty - this is why this is a loop!
275        while val.is_none() {
276            let parent = self.parent.clone()?;
277            let parent_ref = (*parent).borrow();
278            self.parent = parent_ref.parent.clone();
279            self.inner = parent_ref.table.clone().into_iter();
280
281            val = self.inner.next();
282        }
283
284        val
285    }
286}
287
288impl HasId for SymbolTable {
289    fn id(&self) -> ObjId {
290        self.id
291    }
292}
293
294impl DefaultWithId for SymbolTable {
295    fn default_with_id(id: ObjId) -> Self {
296        Self {
297            table: BTreeMap::new(),
298            id,
299            parent: None,
300            next_machine_name: RefCell::new(0),
301        }
302    }
303}
304
305impl Clone for SymbolTable {
306    fn clone(&self) -> Self {
307        Self {
308            table: self.table.clone(),
309            id: ID_COUNTER.with(|x| x.fetch_add(1, Ordering::Relaxed)),
310            parent: self.parent.clone(),
311            next_machine_name: self.next_machine_name.clone(),
312        }
313    }
314}
315
316impl Default for SymbolTable {
317    fn default() -> Self {
318        Self::new_inner(None)
319    }
320}
321
322impl Uniplate for SymbolTable {
323    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
324        // do not recurse up parents, that would be weird?
325        let self2 = self.clone();
326        (Tree::Zero, Box::new(move |_| self2.clone()))
327    }
328}
329
330impl Biplate<Expression> for SymbolTable {
331    fn biplate(&self) -> (Tree<Expression>, Box<dyn Fn(Tree<Expression>) -> Self>) {
332        let (child_trees, ctxs): (VecDeque<_>, Vec<_>) = self
333            .table
334            .values()
335            .map(|decl| <Declaration as Biplate<Expression>>::biplate(decl.as_ref()))
336            .unzip();
337
338        let tree = Tree::Many(child_trees);
339
340        let self2 = self.clone();
341        let ctx = Box::new(move |tree| {
342            let Tree::Many(exprs) = tree else {
343                panic!("unexpected children structure");
344            };
345
346            let mut self3 = self2.clone();
347            let self3_iter = self3.table.iter_mut();
348            for (ctx, tree, (_, decl)) in izip!(&ctxs, exprs, self3_iter) {
349                // update declaration inside the pointer instead of creating a new one, so all
350                // things referencing this keep referencing this.
351                *decl = Rc::new(ctx(tree));
352            }
353
354            self3
355        });
356
357        (tree, ctx)
358    }
359}
360
361impl Biplate<SubModel> for SymbolTable {
362    // walk into expressions
363    fn biplate(&self) -> (Tree<SubModel>, Box<dyn Fn(Tree<SubModel>) -> Self>) {
364        let (exprs, exprs_ctx) = <SymbolTable as Biplate<Expression>>::biplate(self);
365        let (submodel_tree, submodel_ctx) =
366            <VecDeque<Expression> as Biplate<SubModel>>::biplate(&exprs.into_iter().collect());
367
368        let ctx = Box::new(move |x| {
369            exprs_ctx(Tree::Many(
370                submodel_ctx(x).into_iter().map(Tree::One).collect(),
371            ))
372        });
373        (submodel_tree, ctx)
374    }
375}
conjure_core/ast/symbol_table.rs

conjure_core/ast/
symbol_table.rs
