1
//! The (global) symbol table ([`SymbolTable`]), mapping identifiers (of type [`Name`]) to their
2
//! definitions.
3
//!
4
//! This only contains one type of definition at present, [`DecisionVariable`].
5

            
6
use std::collections::BTreeSet;
7
use std::fmt::Display;
8
use std::{cell::RefCell, collections::BTreeMap};
9

            
10
use serde::{Deserialize, Serialize};
11
use serde_with::serde_as;
12

            
13
use crate::ast::variables::DecisionVariable;
14

            
15
use super::{Domain, ReturnType};
16
use derivative::Derivative;
17

            
18
/// A reference to an object stored in the [`SymbolTable`].
19
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)]
20
pub enum Name {
21
    /// A name given in the input model.
22
    UserName(String),
23
    /// A name generated by Conjure-Oxide.
24
    MachineName(i32),
25
}
26

            
27
uniplate::derive_unplateable!(Name);
28

            
29
impl Display for Name {
30
524467
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31
524467
        match self {
32
367710
            Name::UserName(s) => write!(f, "{}", s),
33
156757
            Name::MachineName(i) => write!(f, "__{}", i),
34
        }
35
524467
    }
36
}
37

            
38
/// The global symbol table. Maps [`Names`](Name) 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
#[derive(Derivative)]
43
#[derivative(Hash, PartialEq)]
44
#[serde_as]
45
#[derive(Clone, Debug, Eq, Serialize, Deserialize)]
46
pub struct SymbolTable {
47
    // doing some information hiding here to allow for future refactoring (in particular adding
48
    // lettings, removing variables).
49
    #[serde_as(as = "Vec<(_, _)>")]
50
    variables: BTreeMap<Name, DecisionVariable>,
51

            
52
    #[derivative(Hash = "ignore")]
53
    #[derivative(PartialEq = "ignore")]
54
    next_machine_name: RefCell<i32>,
55
}
56

            
57
impl SymbolTable {
58
    /// Creates an empty symbol table.
59
16871
    pub fn new() -> Self {
60
16871
        SymbolTable {
61
16871
            variables: BTreeMap::new(),
62
16871
            next_machine_name: RefCell::new(0),
63
16871
        }
64
16871
    }
65

            
66
    /*****************************/
67
    /*        get entries        */
68
    /*****************************/
69

            
70
    /// Returns an iterator over the names in the symbol table.
71
    ///
72
    /// Alias of [`names`].
73
    pub fn keys(&self) -> impl Iterator<Item = &Name> {
74
        self.names()
75
    }
76

            
77
    /// Returns an iterator over the names in the symbol table.
78
52088
    pub fn names(&self) -> impl Iterator<Item = &Name> {
79
52088
        self.variables.keys()
80
52088
    }
81

            
82
    /// Returns an iterator over the names and defintions of all decision variables in the symbol
83
    /// table.
84
1462
    pub fn iter_var(&self) -> impl Iterator<Item = (&Name, &DecisionVariable)> {
85
1462
        self.variables.iter()
86
1462
    }
87

            
88
    /// Returns a reference to the decision variable with the given name.
89
    ///
90
    /// Returns `None` if:
91
    ///
92
    /// + There is no decision variable with that name.
93
    ///
94
    /// + The decision variable with that name has been deleted.
95
    ///
96
    /// + The object with that name is not a decision variable.
97
10302
    pub fn get_var(&self, name: &Name) -> Option<&DecisionVariable> {
98
10302
        self.variables.get(name)
99
10302
    }
100

            
101
    /// Returns a mutable reference to the decision variable with the given name.
102
    ///
103
    /// Returns `None` if:
104
    ///
105
    /// + There is no decision variable with that name.
106
    ///
107
    /// + The decision variable with that name has been deleted.
108
    ///
109
    /// + The object with that name is not a decision variable.
110
17
    pub fn get_var_mut(&mut self, name: &Name) -> Option<&mut DecisionVariable> {
111
17
        self.variables.get_mut(name)
112
17
    }
113

            
114
    /********************************/
115
    /*        mutate entries        */
116
    /********************************/
117

            
118
    /// Adds a decision variable to the symbol table as `name`.
119
    ///
120
    /// Returns `None` if there is a decision variable or other object with that name in the symbol
121
    /// table.
122
8350
    pub fn add_var(&mut self, name: Name, var: DecisionVariable) -> Option<()> {
123
8350
        if let std::collections::btree_map::Entry::Vacant(e) = self.variables.entry(name) {
124
8350
            e.insert(var);
125
8350
            Some(())
126
        } else {
127
            None
128
        }
129
8350
    }
130

            
131
    /// Updates a decision variable to the symbol table as `name`, or adds it.
132
    ///
133
    /// Returns `None` if `name` refers to an object that is not a decision variable.
134
    pub fn update_add_var(&mut self, name: Name, var: DecisionVariable) -> Option<()> {
135
        self.variables.insert(name, var);
136
        Some(())
137
    }
138

            
139
    /// Extends the symbol table with the given symbol table, updating the gensym counter if
140
    /// necessary.
141
11815
    pub fn extend(&mut self, other: SymbolTable) {
142
11815
        if other.names().count() > self.names().count() {
143
986
            let new_vars = other.names().collect::<BTreeSet<_>>();
144
986
            let old_vars = self.names().collect::<BTreeSet<_>>();
145

            
146
1105
            for added_var in new_vars.difference(&old_vars) {
147
1105
                let mut next_var = self.next_machine_name.borrow_mut();
148
1105
                match *added_var {
149
                    Name::UserName(_) => {}
150
1105
                    Name::MachineName(m) => {
151
1105
                        if *m >= *next_var {
152
1105
                            *next_var = *m + 1;
153
1105
                        }
154
                    }
155
                }
156
            }
157
10829
        }
158
11815
        self.variables.extend(other.variables);
159
11815
    }
160

            
161
    /****************************************/
162
    /*        get info about symbols        */
163
    /****************************************/
164

            
165
    /// Gets the domain of `name` if it exists and has a domain.
166
3355
    pub fn domain_of(&self, name: &Name) -> Option<&Domain> {
167
3355
        Some(&self.variables.get(name)?.domain)
168
3355
    }
169

            
170
    /// Gets the domain of `name` as a mutable reference if it exists and has a domain.
171
    pub fn domain_of_mut(&mut self, name: &Name) -> Option<&mut Domain> {
172
        Some(&mut self.variables.get_mut(name)?.domain)
173
    }
174

            
175
    /// Gets the type of `name` if it exists and has a type.
176
    pub fn type_of(&self, name: &Name) -> Option<ReturnType> {
177
        match self.domain_of(name)? {
178
            Domain::BoolDomain => Some(ReturnType::Bool),
179
            Domain::IntDomain(_) => Some(ReturnType::Int),
180
        }
181
    }
182

            
183
    /// Returns an arbitrary variable name that is not in the symbol table.
184
1972
    pub fn gensym(&self) -> Name {
185
1972
        let num = *self.next_machine_name.borrow();
186
1972
        *(self.next_machine_name.borrow_mut()) += 1;
187
1972
        Name::MachineName(num) // incremented when inserted
188
1972
    }
189
}
190

            
191
impl Default for SymbolTable {
192
6239
    fn default() -> Self {
193
6239
        Self::new()
194
6239
    }
195
}