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
16524
pub fn is_atom(expr: &Expr) -> bool {
12
16524
    matches!(expr, Expr::Atomic(_, _))
13
16524
}
14

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

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

            
36
39746
    true
37
200124
}
38

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

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

            
81
    // No need to put an atom in an aux_var
82
    if is_atom(expr) {
83
        return None;
84
    }
85

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

            
91
    let name = m.gensym();
92

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

            
98
    m.add_variable(name.clone(), DecisionVariable::new(domain.clone()));
99

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

            
109
/// Output data of `to_aux_var`.
110
pub struct ToAuxVarOutput {
111
    aux_name: Name,
112
    aux_decl: Expr,
113
    #[allow(dead_code)] // TODO: aux_domain should be used soon, try removing this pragma
114
    aux_domain: Domain,
115
    new_model: Model,
116
    _unconstructable: (),
117
}
118

            
119
impl ToAuxVarOutput {
120
    /// Returns the new auxiliary variable as an `Atom`.
121
578
    pub fn as_atom(&self) -> Atom {
122
578
        Atom::Reference(self.aux_name())
123
578
    }
124

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

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

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

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