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

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

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

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

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

            
37
19057
    true
38
93075
}
39

            
40
/// Creates a single `Expr` from a `Vec<Expr>` by forming the conjunction.
41
///
42
/// If it contains a single element, that element is returned, otherwise the conjunction of the
43
/// elements is returned.
44
///
45
/// # Returns
46
///
47
/// - Some: the expression list is non empty, and all expressions are booleans (so can be
48
/// conjuncted).
49
///
50
/// - None: the expression list is empty, or not all expressions are booleans.
51
170
pub fn exprs_to_conjunction(exprs: &Vec<Expr>) -> Option<Expr> {
52
170
    match exprs.as_slice() {
53
170
        [] => None,
54
170
        [a] if a.return_type() == Some(ReturnType::Bool) => Some(a.clone()),
55
        [_, _, ..] => {
56
            for expr in exprs {
57
                if expr.return_type() != Some(ReturnType::Bool) {
58
                    return None;
59
                }
60
            }
61
            Some(Expr::And(Metadata::new(), exprs.clone()))
62
        }
63
        _ => None,
64
    }
65
170
}
66

            
67
/// Creates a new auxiliary variable using the given expression.
68
///
69
/// # Returns
70
///
71
/// * `None` if `Expr` is a `Factor`, or `Expr` does not have a domain (for example, if it is a `Bubble`).
72
///
73
/// * `Some(ToAuxVarOutput)` if successful, containing:
74
///     
75
///     + A new model, modified to include the auxiliary variable in the symbol table.
76
///     + A new top level expression, containing the declaration of the auxiliary variable.
77
///     + A reference to the auxiliary variable to replace the existing expression with.
78
///
79
#[instrument]
80
pub fn to_aux_var(expr: &Expr, m: &Model) -> Option<ToAuxVarOutput> {
81
    let mut m = m.clone();
82

            
83
    // No need to put a factor in an aux_var
84
    if is_factor(expr) {
85
        return None;
86
    }
87

            
88
    // Anything that should be bubbled, bubble
89
    if (expr.can_be_undefined()) {
90
        return None;
91
    }
92

            
93
    let name = m.gensym();
94

            
95
    let Some(domain) = expr.domain_of(&m.variables) else {
96
        tracing::trace!("could not find domain of {}", expr);
97
        return None;
98
    };
99

            
100
    m.add_variable(name.clone(), DecisionVariable::new(domain.clone()));
101

            
102
    Some(ToAuxVarOutput {
103
        aux_name: name.clone(),
104
        aux_decl: Expr::AuxDeclaration(Metadata::new(), name, Box::new(expr.clone())),
105
        aux_domain: domain,
106
        new_model: m,
107
        _unconstructable: (),
108
    })
109
}
110

            
111
/// Output data of `to_aux_var`.
112
pub struct ToAuxVarOutput {
113
    aux_name: Name,
114
    aux_decl: Expr,
115
    aux_domain: Domain,
116
    new_model: Model,
117
    _unconstructable: (),
118
}
119

            
120
impl ToAuxVarOutput {
121
    /// Returns the new auxiliary variable as a `Factor`.
122
289
    pub fn as_factor(&self) -> Factor {
123
289
        Factor::Reference(self.aux_name())
124
289
    }
125

            
126
    /// Returns the new auxiliary variable as an `Expression`.
127
    ///
128
    /// This expression will have default `Metadata`.
129
289
    pub fn as_expr(&self) -> Expr {
130
289
        Expr::FactorE(Metadata::new(), self.as_factor())
131
289
    }
132

            
133
    /// Returns the top level `Expression` to add to the model.
134
289
    pub fn top_level_expr(&self) -> Expr {
135
289
        self.aux_decl.clone()
136
289
    }
137

            
138
    /// Returns the new `Model`, modified to contain this auxiliary variable in the symbol table.
139
    ///
140
    /// Like `Reduction`, this new model does not include the new top level expression. To get
141
    /// this, use [`top_level_expr()`](`ToAuxVarOutput::top_level_expr()`).
142
289
    pub fn model(&self) -> Model {
143
289
        self.new_model.clone()
144
289
    }
145

            
146
    /// Returns the name of the auxiliary variable.
147
289
    pub fn aux_name(&self) -> Name {
148
289
        self.aux_name.clone()
149
289
    }
150
}