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

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

            
10
/// True iff `expr` is an `Atom`.
11
25755
pub fn is_atom(expr: &Expr) -> bool {
12
25755
    matches!(expr, Expr::Atomic(_, _))
13
25755
}
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
652800
pub fn is_all_constant(expression: &Expr) -> bool {
27
657033
    for atom in expression.universe_bi() {
28
657033
        match atom {
29
133263
            Atom::Literal(_) => {}
30
            Atom::Reference(_) => {
31
523770
                return false;
32
            }
33
        }
34
    }
35

            
36
129030
    true
37
652800
}
38

            
39
/// Converts a vector of expressions to a vector of atoms.
40
///
41
/// # Returns
42
///
43
/// `Some(Vec<Atom>)` if the vectors direct children expressions are all atomic, otherwise `None`.
44
#[allow(dead_code)]
45
pub fn expressions_to_atoms(exprs: &Vec<Expr>) -> Option<Vec<Atom>> {
46
    let mut atoms: Vec<Atom> = vec![];
47
    for expr in exprs {
48
        let Expr::Atomic(_, atom) = expr else {
49
            return None;
50
        };
51
        atoms.push(atom.clone());
52
    }
53

            
54
    Some(atoms)
55
}
56

            
57
/// Creates a new auxiliary variable using the given expression.
58
///
59
/// # Returns
60
///
61
/// * `None` if `Expr` is a `Atom`, or `Expr` does not have a domain (for example, if it is a `Bubble`).
62
///
63
/// * `Some(ToAuxVarOutput)` if successful, containing:
64
///     
65
///     + A new model, modified to include the auxiliary variable in the symbol table.
66
///     + A new top level expression, containing the declaration of the auxiliary variable.
67
///     + A reference to the auxiliary variable to replace the existing expression with.
68
///
69
#[instrument]
70
pub fn to_aux_var(expr: &Expr, m: &Model) -> Option<ToAuxVarOutput> {
71
    let mut m = m.clone();
72

            
73
    // No need to put an atom in an aux_var
74
    if is_atom(expr) {
75
        return None;
76
    }
77

            
78
    // Anything that should be bubbled, bubble
79
    if !expr.is_safe() {
80
        return None;
81
    }
82

            
83
    let name = m.gensym();
84

            
85
    let Some(domain) = expr.domain_of(m.symbols()) else {
86
        tracing::trace!("could not find domain of {}", expr);
87
        return None;
88
    };
89

            
90
    m.add_variable(name.clone(), DecisionVariable::new(domain.clone()));
91

            
92
    Some(ToAuxVarOutput {
93
        aux_name: name.clone(),
94
        aux_decl: Expr::AuxDeclaration(Metadata::new(), name, Box::new(expr.clone())),
95
        aux_domain: domain,
96
        new_model: m,
97
        _unconstructable: (),
98
    })
99
}
100

            
101
/// Output data of `to_aux_var`.
102
pub struct ToAuxVarOutput {
103
    aux_name: Name,
104
    aux_decl: Expr,
105
    #[allow(dead_code)] // TODO: aux_domain should be used soon, try removing this pragma
106
    aux_domain: Domain,
107
    new_model: Model,
108
    _unconstructable: (),
109
}
110

            
111
impl ToAuxVarOutput {
112
    /// Returns the new auxiliary variable as an `Atom`.
113
1768
    pub fn as_atom(&self) -> Atom {
114
1768
        Atom::Reference(self.aux_name())
115
1768
    }
116

            
117
    /// Returns the new auxiliary variable as an `Expression`.
118
    ///
119
    /// This expression will have default `Metadata`.
120
1581
    pub fn as_expr(&self) -> Expr {
121
1581
        Expr::Atomic(Metadata::new(), self.as_atom())
122
1581
    }
123

            
124
    /// Returns the top level `Expression` to add to the model.
125
1768
    pub fn top_level_expr(&self) -> Expr {
126
1768
        self.aux_decl.clone()
127
1768
    }
128

            
129
    /// Returns the new `Model`, modified to contain this auxiliary variable in the symbol table.
130
    ///
131
    /// Like `Reduction`, this new model does not include the new top level expression. To get
132
    /// this, use [`top_level_expr()`](`ToAuxVarOutput::top_level_expr()`).
133
1768
    pub fn model(&self) -> Model {
134
1768
        self.new_model.clone()
135
1768
    }
136

            
137
    /// Returns the name of the auxiliary variable.
138
1768
    pub fn aux_name(&self) -> Name {
139
1768
        self.aux_name.clone()
140
1768
    }
141
}