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, Model, Moo, Name, ReturnType, 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    /// Iterates over entries in the LOCAL symbol table, by reference.
365    pub fn iter_local_mut(&mut self) -> impl Iterator<Item = (&Name, &mut DeclarationPtr)> {
366        self.table.iter_mut()
367    }
368
369    /// Extends the symbol table with the given symbol table, updating the gensym counter if
370    /// necessary.
371    pub fn extend(&mut self, other: SymbolTable) {
372        if other.table.keys().count() > self.table.keys().count() {
373            let new_vars = other.table.keys().collect::<BTreeSet<_>>();
374            let old_vars = self.table.keys().collect::<BTreeSet<_>>();
375
376            for added_var in new_vars.difference(&old_vars) {
377                let next_var = &mut self.next_machine_name;
378                if let Name::Machine(m) = *added_var
379                    && *m >= *next_var
380                {
381                    *next_var = *m + 1;
382                }
383            }
384        }
385
386        self.table.extend(other.table);
387    }
388
389    /// Creates a new find declaration in this symbol table with a unique name, and returns its
390    /// declaration.
391    pub fn gen_find(&mut self, domain: &DomainPtr) -> DeclarationPtr {
392        let decl = DeclarationPtr::new_find(self.gen_sym(), domain.clone());
393        self.insert(decl.clone());
394        decl
395    }
396
397    // Reserves a unique machine name in the symbol table
398    pub fn gen_sym(&mut self) -> Name {
399        let num = self.next_machine_name;
400        self.next_machine_name += 1;
401        Name::Machine(num)
402    }
403
404    /// Gets the parent of this symbol table as a mutable reference.
405    ///
406    /// This function provides no sanity checks.
407    pub fn parent_mut_unchecked(&mut self) -> &mut Option<SymbolTablePtr> {
408        &mut self.parent
409    }
410
411    /// Gets the parent of this symbol table.
412    pub fn parent(&self) -> &Option<SymbolTablePtr> {
413        &self.parent
414    }
415
416    /// Gets the representation `representation` for `name`.
417    ///
418    /// # Returns
419    ///
420    /// + `None` if `name` does not exist, is not a decision variable, or does not have that representation.
421    pub fn get_representation(
422        &self,
423        name: &Name,
424        representation: &[&str],
425    ) -> Option<Vec<Box<dyn Representation>>> {
426        // TODO: move representation stuff to declaration / variable to avoid cloning? (we have to
427        // move inside of an rc here, so cannot return borrows)
428        //
429        // Also would prevent constant "does exist" "is var" checks.
430        //
431        // The reason it is not there now is because I'm getting serde issues...
432        //
433        // Also might run into issues putting get_or_add into declaration/variable, as that
434        // requires us to mutably borrow both the symbol table, and the variable inside the symbol
435        // table..
436
437        let decl = self.lookup(name)?;
438        let var = &decl.as_find()?;
439
440        var.representations
441            .iter()
442            .find(|x| &x.iter().map(|r| r.repr_name()).collect_vec()[..] == representation)
443            .cloned()
444    }
445
446    /// Gets all initialised representations for `name`.
447    ///
448    /// # Returns
449    ///
450    /// + `None` if `name` does not exist, or is not a decision variable.
451    pub fn representations_for(&self, name: &Name) -> Option<Vec<Vec<Box<dyn Representation>>>> {
452        let decl = self.lookup(name)?;
453        decl.as_find().map(|x| x.representations.clone())
454    }
455
456    /// Gets the representation `representation` for `name`, creating it if it does not exist.
457    ///
458    /// If the representation does not exist, this method initialises the representation in this
459    /// symbol table, adding the representation to `name`, and the declarations for the represented
460    /// variables to the symbol table.
461    ///
462    /// # Usage
463    ///
464    /// Representations for variable references should be selected and created by the
465    /// `select_representation` rule. Therefore, this method should not be used in other rules.
466    /// Consider using [`get_representation`](`SymbolTable::get_representation`) instead.
467    ///
468    /// # Returns
469    ///
470    /// + `None` if `name` does not exist, is not a decision variable, or cannot be given that
471    ///   representation.
472    pub fn get_or_add_representation(
473        &mut self,
474        name: &Name,
475        representation: &[&str],
476    ) -> Option<Vec<Box<dyn Representation>>> {
477        // Lookup the declaration reference
478        let mut decl = self.lookup(name)?;
479
480        if let Some(var) = decl.as_find()
481            && let Some(existing_reprs) = var
482                .representations
483                .iter()
484                .find(|x| &x.iter().map(|r| r.repr_name()).collect_vec()[..] == representation)
485                .cloned()
486        {
487            return Some(existing_reprs); // Found: return early
488        }
489        // Representation not found
490
491        // TODO: nested representations logic...
492        if representation.len() != 1 {
493            bug!("nested representations not implemented")
494        }
495        let repr_name_str = representation[0];
496        let repr_init_fn = get_repr_rule(repr_name_str)?;
497
498        let reprs = vec![repr_init_fn(name, self)?];
499
500        // Get mutable access to the variable part
501        let mut var = decl.as_find_mut()?;
502
503        for repr_instance in &reprs {
504            repr_instance
505                .declaration_down()
506                .ok()?
507                .into_iter()
508                .for_each(|x| self.update_insert(x));
509        }
510
511        var.representations.push(reprs.clone());
512
513        Some(reprs)
514    }
515}
516
517impl IntoIterator for SymbolTable {
518    type Item = (Name, DeclarationPtr);
519
520    type IntoIter = SymbolTableIter;
521
522    /// Iterates over symbol table entries in scope.
523    fn into_iter(self) -> Self::IntoIter {
524        SymbolTableIter {
525            inner: self.table.into_iter(),
526            parent: self.parent,
527        }
528    }
529}
530
531/// Iterator over all symbol table entries in scope.
532pub struct SymbolTableIter {
533    // iterator over the current scopes' btreemap
534    inner: std::collections::btree_map::IntoIter<Name, DeclarationPtr>,
535
536    // the parent scope
537    parent: Option<SymbolTablePtr>,
538}
539
540impl Iterator for SymbolTableIter {
541    type Item = (Name, DeclarationPtr);
542
543    fn next(&mut self) -> Option<Self::Item> {
544        let mut val = self.inner.next();
545
546        // Go up the tree until we find a parent symbol table with declarations to iterate over.
547        //
548        // Note that the parent symbol table may be empty - this is why this is a loop!
549        while val.is_none() {
550            let parent = self.parent.clone()?;
551
552            let guard = parent.read();
553            self.inner = guard.table.clone().into_iter();
554            self.parent.clone_from(&guard.parent);
555
556            val = self.inner.next();
557        }
558
559        val
560    }
561}
562
563impl Default for SymbolTable {
564    fn default() -> Self {
565        Self::new_inner(None)
566    }
567}
568
569// TODO: if we could override `Uniplate` impl but still derive `Biplate` instances,
570//       we could remove some of this manual code
571impl Uniplate for SymbolTable {
572    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
573        // do not recurse up parents, that would be weird?
574        let self2 = self.clone();
575        (Tree::Zero, Box::new(move |_| self2.clone()))
576    }
577}
578
579impl Biplate<SymbolTablePtr> for SymbolTable {
580    fn biplate(
581        &self,
582    ) -> (
583        Tree<SymbolTablePtr>,
584        Box<dyn Fn(Tree<SymbolTablePtr>) -> Self>,
585    ) {
586        let self2 = self.clone();
587        (Tree::Zero, Box::new(move |_| self2.clone()))
588    }
589}
590
591impl Biplate<Expression> for SymbolTable {
592    fn biplate(&self) -> (Tree<Expression>, Box<dyn Fn(Tree<Expression>) -> Self>) {
593        let (child_trees, ctxs): (VecDeque<_>, Vec<_>) = self
594            .table
595            .values()
596            .map(Biplate::<Expression>::biplate)
597            .unzip();
598
599        let tree = Tree::Many(child_trees);
600
601        let self2 = self.clone();
602        let ctx = Box::new(move |tree| {
603            let Tree::Many(exprs) = tree else {
604                panic!("unexpected children structure");
605            };
606
607            let mut self3 = self2.clone();
608            let self3_iter = self3.table.iter_mut();
609            for (ctx, tree, (_, decl)) in izip!(&ctxs, exprs, self3_iter) {
610                // update declaration inside the pointer instead of creating a new one, so all
611                // things referencing this keep referencing this.
612                *decl = ctx(tree)
613            }
614
615            self3
616        });
617
618        (tree, ctx)
619    }
620}
621
622impl Biplate<Comprehension> for SymbolTable {
623    fn biplate(
624        &self,
625    ) -> (
626        Tree<Comprehension>,
627        Box<dyn Fn(Tree<Comprehension>) -> Self>,
628    ) {
629        let (expr_tree, expr_ctx) = <SymbolTable as Biplate<Expression>>::biplate(self);
630
631        let (exprs, recons_expr_tree) = expr_tree.list();
632
633        let (comprehension_tree, comprehension_ctx) =
634            <VecDeque<Expression> as Biplate<Comprehension>>::biplate(&exprs);
635
636        let ctx = Box::new(move |x| {
637            // 1. turn comprehension tree into a list of expressions
638            let exprs = comprehension_ctx(x);
639
640            // 2. turn list of expressions into an expression tree
641            let expr_tree = recons_expr_tree(exprs);
642
643            // 3. turn expression tree into a symbol table
644            expr_ctx(expr_tree)
645        });
646
647        (comprehension_tree, ctx)
648    }
649}
650
651impl Biplate<Model> for SymbolTable {
652    // walk into expressions
653    fn biplate(&self) -> (Tree<Model>, Box<dyn Fn(Tree<Model>) -> Self>) {
654        let (expr_tree, expr_ctx) = <SymbolTable as Biplate<Expression>>::biplate(self);
655
656        let (exprs, recons_expr_tree) = expr_tree.list();
657
658        let (submodel_tree, submodel_ctx) =
659            <VecDeque<Expression> as Biplate<Model>>::biplate(&exprs);
660
661        let ctx = Box::new(move |x| {
662            // 1. turn submodel tree into a list of expressions
663            let exprs = submodel_ctx(x);
664
665            // 2. turn list of expressions into an expression tree
666            let expr_tree = recons_expr_tree(exprs);
667
668            // 3. turn expression tree into a symbol table
669            expr_ctx(expr_tree)
670        });
671        (submodel_tree, ctx)
672    }
673}
conjure_cp_core/ast/symbol_table.rs

conjure_cp_core/ast/
symbol_table.rs
