1
use tracing::instrument;
2
use uniplate::{Biplate, Uniplate};
3

            
4
use crate::{
5
    ast::{Atom, DecisionVariable, Domain, Expression as Expr, Name, ReturnType},
6
    metadata::Metadata,
7
    Model,
8
};
9

            
10
/// True iff `expr` is an `Atom`.
11
pub fn is_atom(expr: &Expr) -> bool {
12
    matches!(expr, Expr::Atomic(_, _))
13
}
14

            
15
/// True if `expr` is flat; i.e. it only contains atoms.
16
pub fn is_flat(expr: &Expr) -> bool {
17
    for e in expr.children() {
18
        if !is_atom(&e) {
19
            return false;
20
        }
21
    }
22
    true
23
}
24

            
25
/// True if the entire AST is constants.
26
36
pub fn is_all_constant(expression: &Expr) -> bool {
27
36
    for atom in <Expr as Biplate<Atom>>::universe_bi(expression) {
28
18
        match atom {
29
18
            Atom::Literal(_) => {}
30
            Atom::Reference(_) => {
31
                return false;
32
            }
33
        }
34
    }
35

            
36
36
    true
37
36
}
38

            
39
/// Creates a new auxiliary variable using the given expression.
40
///
41
/// # Returns
42
///
43
/// * `None` if `Expr` is a `Atom`, or `Expr` does not have a domain (for example, if it is a `Bubble`).
44
///
45
/// * `Some(ToAuxVarOutput)` if successful, containing:
46
///     
47
///     + A new model, modified to include the auxiliary variable in the symbol table.
48
///     + A new top level expression, containing the declaration of the auxiliary variable.
49
///     + A reference to the auxiliary variable to replace the existing expression with.
50
///
51
#[instrument]
52
pub fn to_aux_var(expr: &Expr, m: &Model) -> Option<ToAuxVarOutput> {
53
    let mut m = m.clone();
54

            
55
    // No need to put an atom in an aux_var
56
    if is_atom(expr) {
57
        return None;
58
    }
59

            
60
    // Anything that should be bubbled, bubble
61
    if expr.can_be_undefined() {
62
        return None;
63
    }
64

            
65
    let name = m.gensym();
66

            
67
    let Some(domain) = expr.domain_of(&m.variables) else {
68
        tracing::trace!("could not find domain of {}", expr);
69
        return None;
70
    };
71

            
72
    m.add_variable(name.clone(), DecisionVariable::new(domain.clone()));
73

            
74
    Some(ToAuxVarOutput {
75
        aux_name: name.clone(),
76
        aux_decl: Expr::AuxDeclaration(Metadata::new(), name, Box::new(expr.clone())),
77
        aux_domain: domain,
78
        new_model: m,
79
        _unconstructable: (),
80
    })
81
}
82

            
83
/// Output data of `to_aux_var`.
84
pub struct ToAuxVarOutput {
85
    aux_name: Name,
86
    aux_decl: Expr,
87
    #[allow(dead_code)] // TODO: aux_domain should be used soon, try removing this pragma
88
    aux_domain: Domain,
89
    new_model: Model,
90
    _unconstructable: (),
91
}
92

            
93
impl ToAuxVarOutput {
94
    /// Returns the new auxiliary variable as an `Atom`.
95
    pub fn as_atom(&self) -> Atom {
96
        Atom::Reference(self.aux_name())
97
    }
98

            
99
    /// Returns the new auxiliary variable as an `Expression`.
100
    ///
101
    /// This expression will have default `Metadata`.
102
    pub fn as_expr(&self) -> Expr {
103
        Expr::Atomic(Metadata::new(), self.as_atom())
104
    }
105

            
106
    /// Returns the top level `Expression` to add to the model.
107
    pub fn top_level_expr(&self) -> Expr {
108
        self.aux_decl.clone()
109
    }
110

            
111
    /// Returns the new `Model`, modified to contain this auxiliary variable in the symbol table.
112
    ///
113
    /// Like `Reduction`, this new model does not include the new top level expression. To get
114
    /// this, use [`top_level_expr()`](`ToAuxVarOutput::top_level_expr()`).
115
    pub fn model(&self) -> Model {
116
        self.new_model.clone()
117
    }
118

            
119
    /// Returns the name of the auxiliary variable.
120
    pub fn aux_name(&self) -> Name {
121
        self.aux_name.clone()
122
    }
123
}