conjure_cp_core/ast/
symbol_table.rs

1//! The symbol table.
2//!
3//! See the item documentation for [`SymbolTable`] for more details.
4
5use crate::bug;
6use crate::representation::{Representation, get_repr_rule};
7use std::any::TypeId;
8
9use std::collections::BTreeSet;
10use std::collections::btree_map::Entry;
11use std::collections::{BTreeMap, VecDeque};
12use std::hash::{Hash, Hasher};
13use std::sync::Arc;
14use std::sync::atomic::{AtomicU32, Ordering};
15
16use super::comprehension::Comprehension;
17use super::serde::{AsId, DefaultWithId, HasId, IdPtr, ObjId, PtrAsInner};
18use super::{
19    DeclarationPtr, DomainPtr, Expression, GroundDomain, Moo, Name, ReturnType, SubModel, Typeable,
20};
21use itertools::{Itertools as _, izip};
22use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
23use serde::{Deserialize, Serialize};
24use serde_with::serde_as;
25use tracing::trace;
26use uniplate::{Biplate, Tree, Uniplate};
27
28/// Global counter of symbol tables.
29/// Note that the counter is shared between all threads
30/// Thus, when running multiple models in parallel, IDs may
31/// be different with every run depending on scheduling order
32static SYMBOL_TABLE_ID_COUNTER: AtomicU32 = const { AtomicU32::new(0) };
33
34#[derive(Debug, Clone, PartialEq, Eq, Hash)]
35pub struct SymbolTablePtr
36where
37    Self: Send + Sync,
38{
39    inner: Arc<SymbolTablePtrInner>,
40}
41
42impl SymbolTablePtr {
43    /// Create an empty new [SymbolTable] and return a shared pointer to it
44    pub fn new() -> Self {
45        Self::new_with_data(SymbolTable::new())
46    }
47
48    /// Create an empty new [SymbolTable] with the given parent and return a shared pointer to it
49    pub fn with_parent(symbols: SymbolTablePtr) -> Self {
50        Self::new_with_data(SymbolTable::with_parent(symbols))
51    }
52
53    fn new_with_data(data: SymbolTable) -> Self {
54        let object_id = SYMBOL_TABLE_ID_COUNTER.fetch_add(1, Ordering::Relaxed);
55        let id = ObjId {
56            object_id,
57            type_name: SymbolTablePtr::TYPE_NAME.into(),
58        };
59        Self::new_with_id_and_data(id, data)
60    }
61
62    fn new_with_id_and_data(id: ObjId, data: SymbolTable) -> Self {
63        Self {
64            inner: Arc::new(SymbolTablePtrInner {
65                id,
66                value: RwLock::new(data),
67            }),
68        }
69    }
70
71    /// Read the underlying symbol table.
72    /// This will block the current thread until a read lock can be acquired.
73    ///
74    /// # WARNING
75    ///
76    /// - If the current thread already holds a lock over this table, this may deadlock.
77    pub fn read(&self) -> RwLockReadGuard<'_, SymbolTable> {
78        self.inner.value.read()
79    }
80
81    /// Mutate the underlying symbol table.
82    /// This will block the current thread until an exclusive write lock can be acquired.
83    ///
84    /// # WARNING
85    ///
86    /// - If the current thread already holds a lock over this table, this may deadlock.
87    /// - Trying to acquire any other lock until the write lock is released will cause a deadlock.
88    /// - This will mutate the underlying data, which may be shared between other `SymbolTablePtr`s.
89    ///   Make sure that this is what you want.
90    ///
91    /// To create a separate copy of the table, see [SymbolTablePtr::detach].
92    ///
93    pub fn write(&self) -> RwLockWriteGuard<'_, SymbolTable> {
94        self.inner.value.write()
95    }
96
97    /// Create a new symbol table with the same contents as this one, but a new ID,
98    /// and return a pointer to it.
99    pub fn detach(&self) -> Self {
100        Self::new_with_data(self.read().clone())
101    }
102}
103
104impl Default for SymbolTablePtr {
105    fn default() -> Self {
106        Self::new()
107    }
108}
109
110impl HasId for SymbolTablePtr {
111    const TYPE_NAME: &'static str = "SymbolTable";
112
113    fn id(&self) -> ObjId {
114        self.inner.id.clone()
115    }
116}
117
118impl DefaultWithId for SymbolTablePtr {
119    fn default_with_id(id: ObjId) -> Self {
120        Self::new_with_id_and_data(id, SymbolTable::default())
121    }
122}
123
124impl IdPtr for SymbolTablePtr {
125    type Data = SymbolTable;
126
127    fn get_data(&self) -> Self::Data {
128        self.read().clone()
129    }
130
131    fn with_id_and_data(id: ObjId, data: Self::Data) -> Self {
132        Self::new_with_id_and_data(id, data)
133    }
134}
135
136// TODO: this code is almost exactly copied from [DeclarationPtr].
137//       It should be possible to eliminate the duplication...
138//       Perhaps by merging SymbolTablePtr and DeclarationPtr together?
139//       (Alternatively, a macro?)
140
141impl Uniplate for SymbolTablePtr {
142    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
143        let symtab = self.read();
144        let (tree, recons) = Biplate::<SymbolTablePtr>::biplate(&symtab as &SymbolTable);
145
146        let self2 = self.clone();
147        (
148            tree,
149            Box::new(move |x| {
150                let self3 = self2.clone();
151                *(self3.write()) = recons(x);
152                self3
153            }),
154        )
155    }
156}
157
158impl<To> Biplate<To> for SymbolTablePtr
159where
160    SymbolTable: Biplate<To>,
161    To: Uniplate,
162{
163    fn biplate(&self) -> (Tree<To>, Box<dyn Fn(Tree<To>) -> Self>) {
164        if TypeId::of::<To>() == TypeId::of::<Self>() {
165            unsafe {
166                let self_as_to = std::mem::transmute::<&Self, &To>(self).clone();
167                (
168                    Tree::One(self_as_to),
169                    Box::new(move |x| {
170                        let Tree::One(x) = x else { panic!() };
171
172                        let x_as_self = std::mem::transmute::<&To, &Self>(&x);
173                        x_as_self.clone()
174                    }),
175                )
176            }
177        } else {
178            // call biplate on the enclosed declaration
179            let decl = self.read();
180            let (tree, recons) = Biplate::<To>::biplate(&decl as &SymbolTable);
181
182            let self2 = self.clone();
183            (
184                tree,
185                Box::new(move |x| {
186                    let self3 = self2.clone();
187                    *(self3.write()) = recons(x);
188                    self3
189                }),
190            )
191        }
192    }
193}
194
195#[derive(Debug)]
196struct SymbolTablePtrInner {
197    id: ObjId,
198    value: RwLock<SymbolTable>,
199}
200
201impl Hash for SymbolTablePtrInner {
202    fn hash<H: Hasher>(&self, state: &mut H) {
203        self.id.hash(state);
204    }
205}
206
207impl PartialEq for SymbolTablePtrInner {
208    fn eq(&self, other: &Self) -> bool {
209        self.value.read().eq(&other.value.read())
210    }
211}
212
213impl Eq for SymbolTablePtrInner {}
214
215/// The global symbol table, mapping names to their definitions.
216///
217/// Names in the symbol table are unique, including between different types of object stored in the
218/// symbol table. For example, you cannot have a letting and decision variable with the same name.
219///
220/// # Symbol Kinds
221///
222/// The symbol table tracks the following types of symbol:
223///
224/// ## Decision Variables
225///
226/// ```text
227/// find NAME: DOMAIN
228/// ```
229///
230/// See [`DecisionVariable`](super::DecisionVariable).
231///
232/// ## Lettings
233///
234/// Lettings define constants, of which there are two types:
235///
236///   + **Constant values**: `letting val be A`, where A is an [`Expression`].
237///
238///     A can be any integer, boolean, or matrix expression.
239///     A can include references to other lettings, model parameters, and, unlike Savile Row,
240///     decision variables.
241///
242///   + **Constant domains**: `letting Domain be domain D`, where D is a [`Domain`].
243///
244///     D can include references to other lettings and model parameters, and, unlike Savile Row,
245///     decision variables.
246///
247/// Unless otherwise stated, these follow the semantics specified in section 2.2.2 of the Savile
248/// Row manual (version 1.9.1 at time of writing).
249#[serde_as]
250#[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)]
251pub struct SymbolTable {
252    #[serde_as(as = "Vec<(_,PtrAsInner)>")]
253    table: BTreeMap<Name, DeclarationPtr>,
254
255    #[serde_as(as = "Option<AsId>")]
256    parent: Option<SymbolTablePtr>,
257
258    next_machine_name: i32,
259}
260
261impl SymbolTable {
262    /// Creates an empty symbol table.
263    pub fn new() -> SymbolTable {
264        SymbolTable::new_inner(None)
265    }
266
267    /// Creates an empty symbol table with the given parent.
268    pub fn with_parent(parent: SymbolTablePtr) -> SymbolTable {
269        SymbolTable::new_inner(Some(parent))
270    }
271
272    fn new_inner(parent: Option<SymbolTablePtr>) -> SymbolTable {
273        let id = SYMBOL_TABLE_ID_COUNTER.fetch_add(1, Ordering::Relaxed);
274        trace!(
275            "new symbol table: id = {id}  parent_id = {}",
276            parent
277                .as_ref()
278                .map(|x| x.id().to_string())
279                .unwrap_or(String::from("none"))
280        );
281        SymbolTable {
282            table: BTreeMap::new(),
283            next_machine_name: 0,
284            parent,
285        }
286    }
287
288    /// Looks up the declaration with the given name in the current scope only.
289    ///
290    /// Returns `None` if there is no declaration with that name in the current scope.
291    pub fn lookup_local(&self, name: &Name) -> Option<DeclarationPtr> {
292        self.table.get(name).cloned()
293    }
294
295    /// Looks up the declaration with the given name, checking all enclosing scopes.
296    ///
297    /// Returns `None` if there is no declaration with that name in scope.
298    pub fn lookup(&self, name: &Name) -> Option<DeclarationPtr> {
299        self.lookup_local(name).or_else(|| {
300            self.parent
301                .as_ref()
302                .and_then(|parent| parent.read().lookup(name))
303        })
304    }
305
306    /// Inserts a declaration into the symbol table.
307    ///
308    /// Returns `None` if there is already a symbol with this name in the local scope.
309    pub fn insert(&mut self, declaration: DeclarationPtr) -> Option<()> {
310        let name = declaration.name().clone();
311        if let Entry::Vacant(e) = self.table.entry(name) {
312            e.insert(declaration);
313            Some(())
314        } else {
315            None
316        }
317    }
318
319    /// Updates or adds a declaration in the immediate local scope.
320    pub fn update_insert(&mut self, declaration: DeclarationPtr) {
321        let name = declaration.name().clone();
322        self.table.insert(name, declaration);
323    }
324
325    /// Looks up the return type for name if it has one and is in scope.
326    pub fn return_type(&self, name: &Name) -> Option<ReturnType> {
327        self.lookup(name).map(|x| x.return_type())
328    }
329
330    /// Looks up the return type for name if has one and is in the local scope.
331    pub fn return_type_local(&self, name: &Name) -> Option<ReturnType> {
332        self.lookup_local(name).map(|x| x.return_type())
333    }
334
335    /// Looks up the domain of name if it has one and is in scope.
336    ///
337    /// This method can return domain references: if a ground domain is always required, use
338    /// [`SymbolTable::resolve_domain`].
339    pub fn domain(&self, name: &Name) -> Option<DomainPtr> {
340        if let Name::WithRepresentation(name, _) = name {
341            self.lookup(name)?.domain()
342        } else {
343            self.lookup(name)?.domain()
344        }
345    }
346
347    /// Looks up the domain of name, resolving domain references to ground domains.
348    ///
349    /// See [`SymbolTable::domain`].
350    pub fn resolve_domain(&self, name: &Name) -> Option<Moo<GroundDomain>> {
351        self.domain(name)?.resolve()
352    }
353
354    /// Iterates over entries in the LOCAL symbol table.
355    pub fn into_iter_local(self) -> impl Iterator<Item = (Name, DeclarationPtr)> {
356        self.table.into_iter()
357    }
358
359    /// Iterates over entries in the LOCAL symbol table, by reference.
360    pub fn iter_local(&self) -> impl Iterator<Item = (&Name, &DeclarationPtr)> {
361        self.table.iter()
362    }
363
364    /// Extends the symbol table with the given symbol table, updating the gensym counter if
365    /// necessary.
366    pub fn extend(&mut self, other: SymbolTable) {
367        if other.table.keys().count() > self.table.keys().count() {
368            let new_vars = other.table.keys().collect::<BTreeSet<_>>();
369            let old_vars = self.table.keys().collect::<BTreeSet<_>>();
370
371            for added_var in new_vars.difference(&old_vars) {
372                let next_var = &mut self.next_machine_name;
373                if let Name::Machine(m) = *added_var
374                    && *m >= *next_var
375                {
376                    *next_var = *m + 1;
377                }
378            }
379        }
380
381        self.table.extend(other.table);
382    }
383
384    /// Creates a new variable in this symbol table with a unique name, and returns its
385    /// declaration.
386    pub fn gensym(&mut self, domain: &DomainPtr) -> DeclarationPtr {
387        let num = self.next_machine_name;
388        self.next_machine_name += 1;
389        let decl = DeclarationPtr::new_find(Name::Machine(num), domain.clone());
390        self.insert(decl.clone());
391        decl
392    }
393
394    /// Gets the parent of this symbol table as a mutable reference.
395    ///
396    /// This function provides no sanity checks.
397    pub fn parent_mut_unchecked(&mut self) -> &mut Option<SymbolTablePtr> {
398        &mut self.parent
399    }
400
401    /// Gets the parent of this symbol table.
402    pub fn parent(&self) -> &Option<SymbolTablePtr> {
403        &self.parent
404    }
405
406    /// Gets the representation `representation` for `name`.
407    ///
408    /// # Returns
409    ///
410    /// + `None` if `name` does not exist, is not a decision variable, or does not have that representation.
411    pub fn get_representation(
412        &self,
413        name: &Name,
414        representation: &[&str],
415    ) -> Option<Vec<Box<dyn Representation>>> {
416        // TODO: move representation stuff to declaration / variable to avoid cloning? (we have to
417        // move inside of an rc here, so cannot return borrows)
418        //
419        // Also would prevent constant "does exist" "is var" checks.
420        //
421        // The reason it is not there now is because I'm getting serde issues...
422        //
423        // Also might run into issues putting get_or_add into declaration/variable, as that
424        // requires us to mutably borrow both the symbol table, and the variable inside the symbol
425        // table..
426
427        let decl = self.lookup(name)?;
428        let var = &decl.as_var()?;
429
430        var.representations
431            .iter()
432            .find(|x| &x.iter().map(|r| r.repr_name()).collect_vec()[..] == representation)
433            .cloned()
434    }
435
436    /// Gets all initialised representations for `name`.
437    ///
438    /// # Returns
439    ///
440    /// + `None` if `name` does not exist, or is not a decision variable.
441    pub fn representations_for(&self, name: &Name) -> Option<Vec<Vec<Box<dyn Representation>>>> {
442        let decl = self.lookup(name)?;
443        decl.as_var().map(|x| x.representations.clone())
444    }
445
446    /// Gets the representation `representation` for `name`, creating it if it does not exist.
447    ///
448    /// If the representation does not exist, this method initialises the representation in this
449    /// symbol table, adding the representation to `name`, and the declarations for the represented
450    /// variables to the symbol table.
451    ///
452    /// # Usage
453    ///
454    /// Representations for variable references should be selected and created by the
455    /// `select_representation` rule. Therefore, this method should not be used in other rules.
456    /// Consider using [`get_representation`](`SymbolTable::get_representation`) instead.
457    ///
458    /// # Returns
459    ///
460    /// + `None` if `name` does not exist, is not a decision variable, or cannot be given that
461    ///   representation.
462    pub fn get_or_add_representation(
463        &mut self,
464        name: &Name,
465        representation: &[&str],
466    ) -> Option<Vec<Box<dyn Representation>>> {
467        // Lookup the declaration reference
468        let mut decl = self.lookup(name)?;
469
470        if let Some(var) = decl.as_var()
471            && let Some(existing_reprs) = var
472                .representations
473                .iter()
474                .find(|x| &x.iter().map(|r| r.repr_name()).collect_vec()[..] == representation)
475                .cloned()
476        {
477            return Some(existing_reprs); // Found: return early
478        }
479        // Representation not found
480
481        // TODO: nested representations logic...
482        if representation.len() != 1 {
483            bug!("nested representations not implemented")
484        }
485        let repr_name_str = representation[0];
486        let repr_init_fn = get_repr_rule(repr_name_str)?;
487
488        let reprs = vec![repr_init_fn(name, self)?];
489
490        // Get mutable access to the variable part
491        let mut var = decl.as_var_mut()?;
492
493        for repr_instance in &reprs {
494            repr_instance
495                .declaration_down()
496                .ok()?
497                .into_iter()
498                .for_each(|x| self.update_insert(x));
499        }
500
501        var.representations.push(reprs.clone());
502
503        Some(reprs)
504    }
505}
506
507impl IntoIterator for SymbolTable {
508    type Item = (Name, DeclarationPtr);
509
510    type IntoIter = SymbolTableIter;
511
512    /// Iterates over symbol table entries in scope.
513    fn into_iter(self) -> Self::IntoIter {
514        SymbolTableIter {
515            inner: self.table.into_iter(),
516            parent: self.parent,
517        }
518    }
519}
520
521/// Iterator over all symbol table entries in scope.
522pub struct SymbolTableIter {
523    // iterator over the current scopes' btreemap
524    inner: std::collections::btree_map::IntoIter<Name, DeclarationPtr>,
525
526    // the parent scope
527    parent: Option<SymbolTablePtr>,
528}
529
530impl Iterator for SymbolTableIter {
531    type Item = (Name, DeclarationPtr);
532
533    fn next(&mut self) -> Option<Self::Item> {
534        let mut val = self.inner.next();
535
536        // Go up the tree until we find a parent symbol table with declarations to iterate over.
537        //
538        // Note that the parent symbol table may be empty - this is why this is a loop!
539        while val.is_none() {
540            let parent = self.parent.clone()?;
541
542            let guard = parent.read();
543            self.inner = guard.table.clone().into_iter();
544            self.parent.clone_from(&guard.parent);
545
546            val = self.inner.next();
547        }
548
549        val
550    }
551}
552
553impl Default for SymbolTable {
554    fn default() -> Self {
555        Self::new_inner(None)
556    }
557}
558
559// TODO: if we could override `Uniplate` impl but still derive `Biplate` instances,
560//       we could remove some of this manual code
561impl Uniplate for SymbolTable {
562    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
563        // do not recurse up parents, that would be weird?
564        let self2 = self.clone();
565        (Tree::Zero, Box::new(move |_| self2.clone()))
566    }
567}
568
569impl Biplate<SymbolTablePtr> for SymbolTable {
570    fn biplate(
571        &self,
572    ) -> (
573        Tree<SymbolTablePtr>,
574        Box<dyn Fn(Tree<SymbolTablePtr>) -> Self>,
575    ) {
576        let self2 = self.clone();
577        (Tree::Zero, Box::new(move |_| self2.clone()))
578    }
579}
580
581impl Biplate<Expression> for SymbolTable {
582    fn biplate(&self) -> (Tree<Expression>, Box<dyn Fn(Tree<Expression>) -> Self>) {
583        let (child_trees, ctxs): (VecDeque<_>, Vec<_>) = self
584            .table
585            .values()
586            .map(Biplate::<Expression>::biplate)
587            .unzip();
588
589        let tree = Tree::Many(child_trees);
590
591        let self2 = self.clone();
592        let ctx = Box::new(move |tree| {
593            let Tree::Many(exprs) = tree else {
594                panic!("unexpected children structure");
595            };
596
597            let mut self3 = self2.clone();
598            let self3_iter = self3.table.iter_mut();
599            for (ctx, tree, (_, decl)) in izip!(&ctxs, exprs, self3_iter) {
600                // update declaration inside the pointer instead of creating a new one, so all
601                // things referencing this keep referencing this.
602                *decl = ctx(tree)
603            }
604
605            self3
606        });
607
608        (tree, ctx)
609    }
610}
611
612impl Biplate<Comprehension> for SymbolTable {
613    fn biplate(
614        &self,
615    ) -> (
616        Tree<Comprehension>,
617        Box<dyn Fn(Tree<Comprehension>) -> Self>,
618    ) {
619        let (expr_tree, expr_ctx) = <SymbolTable as Biplate<Expression>>::biplate(self);
620
621        let (exprs, recons_expr_tree) = expr_tree.list();
622
623        let (comprehension_tree, comprehension_ctx) =
624            <VecDeque<Expression> as Biplate<Comprehension>>::biplate(&exprs);
625
626        let ctx = Box::new(move |x| {
627            // 1. turn comprehension tree into a list of expressions
628            let exprs = comprehension_ctx(x);
629
630            // 2. turn list of expressions into an expression tree
631            let expr_tree = recons_expr_tree(exprs);
632
633            // 3. turn expression tree into a symbol table
634            expr_ctx(expr_tree)
635        });
636
637        (comprehension_tree, ctx)
638    }
639}
640
641impl Biplate<SubModel> for SymbolTable {
642    // walk into expressions
643    fn biplate(&self) -> (Tree<SubModel>, Box<dyn Fn(Tree<SubModel>) -> Self>) {
644        let (expr_tree, expr_ctx) = <SymbolTable as Biplate<Expression>>::biplate(self);
645
646        let (exprs, recons_expr_tree) = expr_tree.list();
647
648        let (submodel_tree, submodel_ctx) =
649            <VecDeque<Expression> as Biplate<SubModel>>::biplate(&exprs);
650
651        let ctx = Box::new(move |x| {
652            // 1. turn submodel tree into a list of expressions
653            let exprs = submodel_ctx(x);
654
655            // 2. turn list of expressions into an expression tree
656            let expr_tree = recons_expr_tree(exprs);
657
658            // 3. turn expression tree into a symbol table
659            expr_ctx(expr_tree)
660        });
661        (submodel_tree, ctx)
662    }
663}
conjure_cp_core/ast/symbol_table.rs

conjure_cp_core/ast/
symbol_table.rs
