1
use uniplate::{Biplate, Uniplate};
2

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

            
10
/// True iff `expr` is a `Factor`.
11
90
pub fn is_factor(expr: &Expr) -> bool {
12
90
    matches!(expr, Expr::FactorE(_, _))
13
90
}
14

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

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

            
36
16335
    true
37
77640
}
38

            
39
/// Creates a new auxiliary variable using the given expression.
40
///
41
/// # Returns
42
///
43
/// * `None` if `Expr` is a `Factor`, 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
90
pub fn to_aux_var(expr: &Expr, m: &Model) -> Option<ToAuxVarOutput> {
52
90
    let mut m = m.clone();
53
90

            
54
90
    // No need to put a factor in an aux_var
55
90
    if is_factor(expr) {
56
        return None;
57
90
    }
58
90

            
59
90
    let name = m.gensym();
60

            
61
90
    let Some(domain) = expr.domain_of(&m.variables) else {
62
        //bug!("rules::utils::to_aux_var: could not find domain of {expr}");
63
        return None;
64
    };
65

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

            
68
90
    Some(ToAuxVarOutput {
69
90
        aux_name: name.clone(),
70
90
        aux_decl: Expr::AuxDeclaration(Metadata::new(), name, Box::new(expr.clone())),
71
90
        aux_domain: domain,
72
90
        new_model: m,
73
90
        _unconstructable: (),
74
90
    })
75
90
}
76

            
77
/// Output data of `to_aux_var`.
78
pub struct ToAuxVarOutput {
79
    aux_name: Name,
80
    aux_decl: Expr,
81
    aux_domain: Domain,
82
    new_model: Model,
83
    _unconstructable: (),
84
}
85

            
86
impl ToAuxVarOutput {
87
    /// Returns the new auxiliary variable as a `Factor`.
88
180
    pub fn as_factor(&self) -> Factor {
89
180
        Factor::Reference(self.aux_name())
90
180
    }
91

            
92
    /// Returns the new auxiliary variable as an `Expression`.
93
    ///
94
    /// This expression will have default `Metadata`.
95
180
    pub fn as_expr(&self) -> Expr {
96
180
        Expr::FactorE(Metadata::new(), self.as_factor())
97
180
    }
98

            
99
    /// Returns the top level `Expression` to add to the model.
100
    pub fn top_level_expr(&self) -> Expr {
101
        self.aux_decl.clone()
102
    }
103

            
104
    /// Returns the new `Model`, modified to contain this auxiliary variable in the symbol table.
105
    ///
106
    /// Like `Reduction`, this new model does not include the new top level expression. To get
107
    /// this, use [`top_level_expr()`](`ToAuxVarOutput::top_level_expr()`).
108
90
    pub fn model(&self) -> Model {
109
90
        self.new_model.clone()
110
90
    }
111

            
112
    /// Returns the name of the auxiliary variable.
113
180
    pub fn aux_name(&self) -> Name {
114
180
        self.aux_name.clone()
115
180
    }
116
}