conjure_core/ast/

symbol_table.rs

//! The symbol table.
//!
//! See the item documentation for [`SymbolTable`] for more details.

use crate::bug;
use crate::representation::{get_repr_rule, Representation};

use super::comprehension::Comprehension;
use super::serde::{RcRefCellAsId, RcRefCellAsInner};
use std::cell::RefCell;
use std::collections::btree_map::Entry;
use std::collections::BTreeSet;
use std::collections::{BTreeMap, VecDeque};
use std::rc::Rc;
use std::sync::atomic::{AtomicU32, Ordering};

use super::declaration::Declaration;
use super::serde::{DefaultWithId, HasId, ObjId};
use super::types::Typeable;
use itertools::{izip, Itertools as _};
use serde::{Deserialize, Serialize};
use serde_with::serde_as;
use uniplate::Tree;
use uniplate::{Biplate, Uniplate};

use super::name::Name;
use super::{Domain, Expression, ReturnType, SubModel};
use derivative::Derivative;

// Count symbol tables per thread / model.
//
// We run tests in parallel and having the id's be per thread keeps them more deterministic in the
// JSON output. If this were not thread local, ids would be given to symbol tables differently in
// each test run (depending on how the threads were scheduled). These id changes would result in
// all the generated tests "changing" each time `ACCEPT=true cargo test` is ran.
//
// SAFETY: Symbol tables use Rc<RefCell<<>>, so a model is not thread-safe anyways.
thread_local! {
static ID_COUNTER: AtomicU32 = const { AtomicU32::new(0) };
}

/// The global symbol table, mapping names 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.
///
/// # Symbol Kinds
///
/// The symbol table tracks the following types of symbol:
///
/// ## Decision Variables
///
/// ```text
/// find NAME: DOMAIN
/// ```
///
/// See [`DecisionVariable`](super::DecisionVariable).
///
/// ## Lettings
///
/// Lettings define constants, of which there are two types:
///
///   + **Constant values**: `letting val be A`, where A is an [`Expression`].
///
///     A can be any integer, boolean, or matrix expression.
///     A can include references to other lettings, model parameters, and, unlike Savile Row,
///     decision variables.
///
///   + **Constant domains**: `letting Domain be domain D`, where D is a [`Domain`].
///
///     D can include references to other lettings and model parameters, and, unlike Savile Row,
///     decision variables.
///
/// Unless otherwise stated, these follow the semantics specified in section 2.2.2 of the Savile
/// Row manual (version 1.9.1 at time of writing).
#[derive(Derivative)]
#[derivative(PartialEq)]
#[derive(Debug, Eq)]
#[serde_as]
#[derive(Serialize, Deserialize)]
pub struct SymbolTable {
    #[serde_as(as = "Vec<(_,RcRefCellAsInner)>")]
    table: BTreeMap<Name, Rc<Declaration>>,

    /// A unique id for this symbol table, for serialisation and debugging.
    #[derivative(PartialEq = "ignore")] // eq by value not id.
    id: ObjId,

    #[serde_as(as = "Option<RcRefCellAsId>")]
    parent: Option<Rc<RefCell<SymbolTable>>>,

    next_machine_name: RefCell<i32>,
}

impl SymbolTable {
    /// Creates an empty symbol table.
    pub fn new() -> SymbolTable {
        SymbolTable::new_inner(None)
    }

    /// Creates an empty symbol table with the given parent.
    pub fn with_parent(parent: Rc<RefCell<SymbolTable>>) -> SymbolTable {
        SymbolTable::new_inner(Some(parent))
    }

    fn new_inner(parent: Option<Rc<RefCell<SymbolTable>>>) -> SymbolTable {
        let id = ID_COUNTER.with(|x| x.fetch_add(1, Ordering::Relaxed));
        SymbolTable {
            id,
            table: BTreeMap::new(),
            next_machine_name: RefCell::new(0),
            parent,
        }
    }

    /// Looks up the declaration with the given name in the current scope only.
    ///
    /// Returns `None` if there is no declaration with that name in the current scope.
    pub fn lookup_local(&self, name: &Name) -> Option<Rc<Declaration>> {
        self.table.get(name).cloned()
    }

    /// Looks up the declaration with the given name, checking all enclosing scopes.
    ///
    /// Returns `None` if there is no declaration with that name in scope.
    pub fn lookup(&self, name: &Name) -> Option<Rc<Declaration>> {
        self.lookup_local(name).or_else(|| {
            self.parent
                .as_ref()
                .and_then(|parent| (*parent).borrow().lookup(name))
        })
    }

    /// Inserts a declaration into the symbol table.
    ///
    /// Returns `None` if there is already a symbol with this name in the local scope.
    pub fn insert(&mut self, declaration: Rc<Declaration>) -> Option<()> {
        let name = declaration.name().clone();
        if let Entry::Vacant(e) = self.table.entry(name) {
            e.insert(declaration);
            Some(())
        } else {
            None
        }
    }

    /// Updates or adds a declaration in the immediate local scope.
    pub fn update_insert(&mut self, declaration: Rc<Declaration>) {
        let name = declaration.name().clone();
        self.table.insert(name, declaration);
    }

    /// Looks up the return type for name if it has one and is in scope.
    pub fn return_type(&self, name: &Name) -> Option<ReturnType> {
        self.lookup(name).and_then(|x| x.return_type())
    }

    /// Looks up the return type for name if has one and is in the local scope.
    pub fn return_type_local(&self, name: &Name) -> Option<ReturnType> {
        self.lookup_local(name).and_then(|x| x.return_type())
    }

    /// Looks up the domain of name if it has one and is in scope.
    ///
    /// This method can return domain references: if a ground domain is always required, use
    /// [`SymbolTable::resolve_domain`].
    pub fn domain(&self, name: &Name) -> Option<Domain> {
        // TODO: do not clone here: in the future, we might want to wrap all domains in Rc's to get
        // clone-on-write behaviour (saving memory in scenarios such as matrix decomposition where
        // a lot of the domains would be the same).

        if let Name::WithRepresentation(name, _) = name {
            let decl = self.lookup(name.as_ref())?;
            decl.domain().cloned()
        } else {
            let decl = self.lookup(name)?;
            decl.domain().cloned()
        }
    }

    /// Looks up the domain of name, resolving domain references to ground domains.
    ///
    /// See [`SymbolTable::domain`].
    pub fn resolve_domain(&self, name: &Name) -> Option<Domain> {
        self.domain(name).map(|domain| domain.resolve(self))
    }

    /// Iterates over entries in the local symbol table only.
    pub fn into_iter_local(self) -> LocalIntoIter {
        LocalIntoIter {
            inner: self.table.into_iter(),
        }
    }

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

            for added_var in new_vars.difference(&old_vars) {
                let mut next_var = self.next_machine_name.borrow_mut();
                if let Name::Machine(m) = *added_var {
                    if *m >= *next_var {
                        *next_var = *m + 1;
                    }
                }
            }
        }

        self.table.extend(other.table);
    }

    /// 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::Machine(num) // incremented when inserted
    }

    /// Gets the parent of this symbol table as a mutable reference.
    ///
    /// This function provides no sanity checks.
    pub fn parent_mut_unchecked(&mut self) -> &mut Option<Rc<RefCell<SymbolTable>>> {
        &mut self.parent
    }

    /// Gets the representation `representation` for `name`.
    ///
    /// # Returns
    ///
    /// + `None` if `name` does not exist, is not a decision variable, or does not have that representation.
    pub fn get_representation(
        &self,
        name: &Name,
        representation: &[&str],
    ) -> Option<Vec<Box<dyn Representation>>> {
        // TODO: move representation stuff to declaration / variable to avoid cloning? (we have to
        // move inside of an rc here, so cannot return borrows)
        //
        // Also would prevent constant "does exist" "is var" checks.
        //
        // The reason it is not there now is because I'm getting serde issues...
        //
        // Also might run into issues putting get_or_add into declaration/variable, as that
        // requires us to mutably borrow both the symbol table, and the variable inside the symbol
        // table..

        let decl = self.lookup(name)?;
        let var = decl.as_var()?;

        var.representations
            .iter()
            .find(|x| &x.iter().map(|r| r.repr_name()).collect_vec()[..] == representation)
            .cloned()
            .clone()
    }

    /// Gets all initialised representations for `name`.
    ///
    /// # Returns
    ///
    /// + `None` if `name` does not exist, or is not a decision variable.
    pub fn representations_for(&self, name: &Name) -> Option<Vec<Vec<Box<dyn Representation>>>> {
        let decl = self.lookup(name)?;
        decl.as_var().map(|x| x.representations.clone())
    }

    /// Gets the representation `representation` for `name`, creating it if it does not exist.
    ///
    /// If the representation does not exist, this method initialises the representation in this
    /// symbol table, adding the representation to `name`, and the declarations for the represented
    /// variables to the symbol table.
    ///
    /// # Usage
    ///
    /// Representations for variable references should be selected and created by the
    /// `select_representation` rule. Therefore, this method should not be used in other rules.
    /// Consider using [`get_representation`](`SymbolTable::get_representation`) instead.
    ///
    /// # Returns
    ///
    /// + `None` if `name` does not exist, is not a decision variable, or cannot be given that
    ///   representation.
    pub fn get_or_add_representation(
        &mut self,
        name: &Name,
        representation: &[&str],
    ) -> Option<Vec<Box<dyn Representation>>> {
        let decl = self.lookup(name)?;
        let var = decl.as_var()?;

        var.representations
            .iter()
            .find(|x| &x.iter().map(|r| r.repr_name()).collect_vec()[..] == representation)
            .cloned()
            .clone()
            .or_else(|| {
                // add represented name declaration to the same scope as the name they are
                // representing
                if let Some(decl) = self.lookup_local(name) {
                    let mut decl: Declaration = Rc::unwrap_or_clone(decl);
                    let var = decl.as_var_mut()?;

                    // TODO: no idea how to deal with initialising representations for nested things..

                    if representation.len() != 1 {
                        bug!("nested representations not implemented")
                    }

                    let repr_rule = get_repr_rule(representation[0])?;

                    let reprs = vec![repr_rule(name, self)?];
                    for repr in &reprs {
                        repr.declaration_down()
                            .ok()
                            .unwrap()
                            .iter()
                            .for_each(|x| self.update_insert(Rc::new(x.clone())));
                    }
                    var.representations.push(reprs.clone());

                    self.update_insert(Rc::new(decl));
                    Some(reprs)
                } else {
                    None

                    // declaration is in some parent table
                    //
                    // FIXME: this is broken. for now, just add representations for variables in
                    // our scope. This is slower, in particular, when we want to represent a
                    // decision variable referenced inside a comprehension, but it is correct
                    // nonetheless.
                    //
                    //
                    // let mut parent_table = self.parent.clone();
                    // let mut decl = None;
                    // while let Some(table) = parent_table.clone() {
                    //     let table_ref= (*table).borrow();
                    //     if let Some(d) = table_ref.lookup_local(name) {
                    //         decl = Some(d);
                    //         break;
                    //     }
                    //     parent_table = table_ref.parent.clone();
                    // }

                    // let mut decl: Declaration = Rc::unwrap_or_clone(
                    //     decl.expect("declaration should exist: we used it earlier!"),
                    // );

                    // let parent_table = parent_table.unwrap();
                    // let mut table = parent_table.borrow_mut();

                    // let var = decl.as_var_mut()?;

                    // // TODO: no idea how to deal with initialising representations for nested things..

                    // if representation.len() != 1 {
                    //     bug!("nested representations not implemented")
                    // }

                    // let repr_rule = get_repr_rule(representation[0])?;

                    // let reprs = vec![repr_rule(name, &table)?];
                    // for repr in &reprs {
                    //     repr.declaration_down()
                    //         .ok()
                    //         .unwrap()
                    //         .iter()
                    //         .for_each(|x| table.update_insert(Rc::new(x.clone())));
                    // }
                    // var.representations.push(reprs.clone());

                    // table.update_insert(Rc::new(decl));
                    // dbg!(table);
                    // Some(reprs)
                }
            })
    }
}

impl IntoIterator for SymbolTable {
    type Item = (Name, Rc<Declaration>);

    type IntoIter = IntoIter;

    /// Iterates over symbol table entries in scope.
    fn into_iter(self) -> Self::IntoIter {
        IntoIter {
            inner: self.table.into_iter(),
            parent: self.parent,
        }
    }
}

/// Iterator over symbol table entries in the current scope only.
pub struct LocalIntoIter {
    // iterator over the btreemap
    inner: std::collections::btree_map::IntoIter<Name, Rc<Declaration>>,
}

impl Iterator for LocalIntoIter {
    type Item = (Name, Rc<Declaration>);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }
}

/// Iterator over all symbol table entries in scope.
pub struct IntoIter {
    // iterator over the current scopes' btreemap
    inner: std::collections::btree_map::IntoIter<Name, Rc<Declaration>>,

    // the parent scope
    parent: Option<Rc<RefCell<SymbolTable>>>,
}

impl Iterator for IntoIter {
    type Item = (Name, Rc<Declaration>);

    fn next(&mut self) -> Option<Self::Item> {
        let mut val = self.inner.next();

        // Go up the tree until we find a parent symbol table with declarations to iterate over.
        //
        // Note that the parent symbol table may be empty - this is why this is a loop!
        while val.is_none() {
            let parent = self.parent.clone()?;
            let parent_ref = (*parent).borrow();
            self.parent = parent_ref.parent.clone();
            self.inner = parent_ref.table.clone().into_iter();

            val = self.inner.next();
        }

        val
    }
}

impl HasId for SymbolTable {
    fn id(&self) -> ObjId {
        self.id
    }
}

impl DefaultWithId for SymbolTable {
    fn default_with_id(id: ObjId) -> Self {
        Self {
            table: BTreeMap::new(),
            id,
            parent: None,
            next_machine_name: RefCell::new(0),
        }
    }
}

impl Clone for SymbolTable {
    fn clone(&self) -> Self {
        Self {
            table: self.table.clone(),
            id: ID_COUNTER.with(|x| x.fetch_add(1, Ordering::Relaxed)),
            parent: self.parent.clone(),
            next_machine_name: self.next_machine_name.clone(),
        }
    }
}

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

impl Uniplate for SymbolTable {
    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
        // do not recurse up parents, that would be weird?
        let self2 = self.clone();
        (Tree::Zero, Box::new(move |_| self2.clone()))
    }
}

impl Biplate<Expression> for SymbolTable {
    fn biplate(&self) -> (Tree<Expression>, Box<dyn Fn(Tree<Expression>) -> Self>) {
        let (child_trees, ctxs): (VecDeque<_>, Vec<_>) = self
            .table
            .values()
            .map(|decl| <Declaration as Biplate<Expression>>::biplate(decl.as_ref()))
            .unzip();

        let tree = Tree::Many(child_trees);

        let self2 = self.clone();
        let ctx = Box::new(move |tree| {
            let Tree::Many(exprs) = tree else {
                panic!("unexpected children structure");
            };

            let mut self3 = self2.clone();
            let self3_iter = self3.table.iter_mut();
            for (ctx, tree, (_, decl)) in izip!(&ctxs, exprs, self3_iter) {
                // update declaration inside the pointer instead of creating a new one, so all
                // things referencing this keep referencing this.
                *decl = Rc::new(ctx(tree));
            }

            self3
        });

        (tree, ctx)
    }
}

impl Biplate<Comprehension> for SymbolTable {
    fn biplate(
        &self,
    ) -> (
        Tree<Comprehension>,
        Box<dyn Fn(Tree<Comprehension>) -> Self>,
    ) {
        let (expr_tree, expr_ctx) = <SymbolTable as Biplate<Expression>>::biplate(self);
        let (comprehension_tree, comprehension_ctx) =
            <VecDeque<Expression> as Biplate<Comprehension>>::biplate(
                &expr_tree.into_iter().collect(),
            );

        let ctx = Box::new(move |x| {
            expr_ctx(Tree::Many(
                comprehension_ctx(x).into_iter().map(Tree::One).collect(),
            ))
        });

        (comprehension_tree, ctx)
    }
}

impl Biplate<SubModel> for SymbolTable {
    // walk into expressions
    fn biplate(&self) -> (Tree<SubModel>, Box<dyn Fn(Tree<SubModel>) -> Self>) {
        let (exprs, exprs_ctx) = <SymbolTable as Biplate<Expression>>::biplate(self);
        let (submodel_tree, submodel_ctx) =
            <VecDeque<Expression> as Biplate<SubModel>>::biplate(&exprs.into_iter().collect());

        let ctx = Box::new(move |x| {
            exprs_ctx(Tree::Many(
                submodel_ctx(x).into_iter().map(Tree::One).collect(),
            ))
        });
        (submodel_tree, ctx)
    }
}