conjure_core/ast/

symbol_table.rs

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

use std::collections::BTreeSet;
use std::fmt::Display;
use std::{cell::RefCell, collections::BTreeMap};

use serde::{Deserialize, Serialize};
use serde_with::serde_as;

use crate::ast::variables::DecisionVariable;

use super::{Domain, ReturnType};
use derivative::Derivative;

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

uniplate::derive_unplateable!(Name);

impl Display for Name {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Name::UserName(s) => write!(f, "{}", s),
            Name::MachineName(i) => write!(f, "__{}", i),
        }
    }
}

/// The global symbol table. Maps [`Names`](Name) to their definitions.
///
/// Names in the symbol table are unique, including between different types of object stored in the
/// symbol table. For example, you cannot have a letting and decision variable with the same name.
#[derive(Derivative)]
#[derivative(Hash, PartialEq)]
#[serde_as]
#[derive(Clone, Debug, Eq, Serialize, Deserialize)]
pub struct SymbolTable {
    // doing some information hiding here to allow for future refactoring (in particular adding
    // lettings, removing variables).
    #[serde_as(as = "Vec<(_, _)>")]
    variables: BTreeMap<Name, DecisionVariable>,

    #[derivative(Hash = "ignore")]
    #[derivative(PartialEq = "ignore")]
    next_machine_name: RefCell<i32>,
}

impl SymbolTable {
    /// Creates an empty symbol table.
    pub fn new() -> Self {
        SymbolTable {
            variables: BTreeMap::new(),
            next_machine_name: RefCell::new(0),
        }
    }

    /*****************************/
    /*        get entries        */
    /*****************************/

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

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

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

    /// Returns a reference to the decision variable with the given name.
    ///
    /// Returns `None` if:
    ///
    /// + There is no decision variable with that name.
    ///
    /// + The decision variable with that name has been deleted.
    ///
    /// + The object with that name is not a decision variable.
    pub fn get_var(&self, name: &Name) -> Option<&DecisionVariable> {
        self.variables.get(name)
    }

    /// Returns a mutable reference to the decision variable with the given name.
    ///
    /// Returns `None` if:
    ///
    /// + There is no decision variable with that name.
    ///
    /// + The decision variable with that name has been deleted.
    ///
    /// + The object with that name is not a decision variable.
    pub fn get_var_mut(&mut self, name: &Name) -> Option<&mut DecisionVariable> {
        self.variables.get_mut(name)
    }

    /********************************/
    /*        mutate entries        */
    /********************************/

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

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

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

            for added_var in new_vars.difference(&old_vars) {
                let mut next_var = self.next_machine_name.borrow_mut();
                match *added_var {
                    Name::UserName(_) => {}
                    Name::MachineName(m) => {
                        if *m >= *next_var {
                            *next_var = *m + 1;
                        }
                    }
                }
            }
        }
        self.variables.extend(other.variables);
    }

    /****************************************/
    /*        get info about symbols        */
    /****************************************/

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

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

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

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

impl Default for SymbolTable {
    fn default() -> Self {
        Self::new()
    }
}