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

conjure_cp_core/ast/
symbol_table.rs
