1
use conjure_cp::ast::{Atom, Expression as Expr, Literal};
2
use conjure_cp::ast::{SATIntEncoding, SymbolTable};
3
use conjure_cp::rule_engine::ApplicationError;
4
use conjure_cp::rule_engine::{
5
    ApplicationError::RuleNotApplicable, ApplicationResult, Reduction, register_rule,
6
};
7

            
8
use crate::sat::boolean::{tseytin_and, tseytin_iff, tseytin_not, tseytin_or};
9
use conjure_cp::ast::Metadata;
10
use conjure_cp::ast::Moo;
11
use conjure_cp::into_matrix_expr;
12

            
13
/// This function confirms that all of the input expressions are order SATInts, and returns vectors for each input of their bits
14
/// This function also normalizes order SATInt operands to a common value range.
15
3744
pub fn validate_order_int_operands(
16
3744
    exprs: Vec<Expr>,
17
3744
) -> Result<(Vec<Vec<Expr>>, i32, i32), ApplicationError> {
18
    // Iterate over all inputs
19
    // Check they are order and calulate a lower and upper bound
20
3744
    let mut global_min: i32 = i32::MAX;
21
3744
    let mut global_max: i32 = i32::MIN;
22

            
23
6822
    for operand in &exprs {
24
4932
        let Expr::SATInt(_, SATIntEncoding::Order, _, (local_min, local_max)) = operand else {
25
1890
            return Err(RuleNotApplicable);
26
        };
27
4932
        global_min = global_min.min(*local_min);
28
4932
        global_max = global_max.max(*local_max);
29
    }
30

            
31
    // build out by iterating over each operand and expanding it to match the new bounds
32
1854
    let out: Vec<Vec<Expr>> = exprs
33
1854
        .into_iter()
34
3654
        .map(|expr| {
35
3654
            let Expr::SATInt(_, SATIntEncoding::Order, inner, (local_min, local_max)) = expr else {
36
                return Err(RuleNotApplicable);
37
            };
38

            
39
3654
            let Some(v) = inner.as_ref().clone().unwrap_list() else {
40
                return Err(RuleNotApplicable);
41
            };
42

            
43
            // calulcate how many trues/falses to prepend/append
44
3654
            let prefix_len = (local_min - global_min) as usize;
45
3654
            let postfix_len = (global_max - local_max) as usize;
46

            
47
3654
            let mut bits = Vec::with_capacity(v.len() + prefix_len + postfix_len);
48

            
49
            // add `true`s to start
50
3654
            bits.extend(std::iter::repeat_n(
51
3654
                Expr::Atomic(Metadata::new(), Atom::Literal(Literal::Bool(true))),
52
3654
                prefix_len,
53
            ));
54

            
55
3654
            bits.extend(v);
56

            
57
            // add `false`s to end
58
3654
            bits.extend(std::iter::repeat_n(
59
3654
                Expr::Atomic(Metadata::new(), Atom::Literal(Literal::Bool(false))),
60
3654
                postfix_len,
61
            ));
62

            
63
3654
            Ok(bits)
64
3654
        })
65
1854
        .collect::<Result<_, _>>()?;
66

            
67
1854
    Ok((out, global_min, global_max))
68
3744
}
69

            
70
/// Encodes a < b for order integers.
71
///
72
/// `x < y` iff `exists i . (NOT x_i AND y_i)`
73
1464
fn sat_order_lt(
74
1464
    a_bits: Vec<Expr>,
75
1464
    b_bits: Vec<Expr>,
76
1464
    clauses: &mut Vec<conjure_cp::ast::CnfClause>,
77
1464
    symbols: &mut SymbolTable,
78
1464
) -> Expr {
79
1464
    let mut result = Expr::Atomic(Metadata::new(), Atom::Literal(Literal::Bool(false)));
80

            
81
10296
    for (a_i, b_i) in a_bits.iter().zip(b_bits.iter()) {
82
        // (NOT a_i AND b_i)
83
10296
        let not_a_i = tseytin_not(a_i.clone(), clauses, symbols);
84
10296
        let current_term = tseytin_and(&vec![not_a_i, b_i.clone()], clauses, symbols);
85
10296

            
86
        // accumulate (NOT a_i AND b_i) into OR term
87
10296
        result = tseytin_or(&vec![result, current_term], clauses, symbols);
88
10296
    }
89
1464
    result
90
1464
}
91

            
92
/// Converts an integer literal to SATInt form
93
///
94
/// ```text
95
///  3
96
///  ~~>
97
///  SATInt([true;int(1..), (3, 3)])
98
///
99
/// ```
100
#[register_rule("SAT_Order", 9500, [Atomic])]
101
279513
fn literal_sat_order_int(expr: &Expr, _: &SymbolTable) -> ApplicationResult {
102
1272
    let value = {
103
12348
        if let Expr::Atomic(_, Atom::Literal(Literal::Int(value))) = expr {
104
1272
            *value
105
        } else {
106
278241
            return Err(RuleNotApplicable);
107
        }
108
    };
109

            
110
1272
    Ok(Reduction::pure(Expr::SATInt(
111
1272
        Metadata::new(),
112
1272
        SATIntEncoding::Order,
113
1272
        Moo::new(into_matrix_expr!(vec![Expr::Atomic(
114
1272
            Metadata::new(),
115
1272
            Atom::Literal(Literal::Bool(true)),
116
1272
        )])),
117
1272
        (value, value),
118
1272
    )))
119
279513
}
120

            
121
/// Builds CNF for equality between two order SATInt bit-vectors.
122
/// This function is used by both eq and neq rules, with the output negated for neq.
123
/// Returns (expr, clauses, symbols).
124
336
fn sat_order_eq_expr(
125
336
    lhs_bits: &[Expr],
126
336
    rhs_bits: &[Expr],
127
336
    symbols: &SymbolTable,
128
336
) -> (Expr, Vec<conjure_cp::ast::CnfClause>, SymbolTable) {
129
336
    let bit_count = lhs_bits.len();
130

            
131
336
    let mut output = true.into();
132
336
    let mut new_symbols = symbols.clone();
133
336
    let mut new_clauses = vec![];
134

            
135
1788
    for i in 0..bit_count {
136
1788
        let comparison = tseytin_iff(
137
1788
            lhs_bits[i].clone(),
138
1788
            rhs_bits[i].clone(),
139
1788
            &mut new_clauses,
140
1788
            &mut new_symbols,
141
1788
        );
142
1788
        output = tseytin_and(
143
1788
            &vec![comparison, output],
144
1788
            &mut new_clauses,
145
1788
            &mut new_symbols,
146
1788
        );
147
1788
    }
148

            
149
336
    (output, new_clauses, new_symbols)
150
336
}
151

            
152
/// Converts a = expression between two order SATInts to a boolean expression in cnf
153
///
154
/// ```text
155
/// SATInt(a) = SATInt(b) ~> Bool
156
/// ```
157
#[register_rule("SAT_Order", 9100, [Eq])]
158
79653
fn eq_sat_order(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
159
79653
    let Expr::Eq(_, lhs, rhs) = expr else {
160
79281
        return Err(RuleNotApplicable);
161
    };
162

            
163
312
    let (binding, _, _) =
164
372
        validate_order_int_operands(vec![lhs.as_ref().clone(), rhs.as_ref().clone()])?;
165
312
    let [lhs_bits, rhs_bits] = binding.as_slice() else {
166
        return Err(RuleNotApplicable);
167
    };
168

            
169
312
    let (output, new_clauses, new_symbols) = sat_order_eq_expr(lhs_bits, rhs_bits, symbols);
170

            
171
312
    Ok(Reduction::cnf(output, new_clauses, new_symbols))
172
79653
}
173

            
174
/// Converts a != expression between two order SATInts to a boolean expression in cnf
175
#[register_rule("SAT_Order", 9100, [Neq])]
176
79653
fn neq_sat_order(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
177
79653
    let Expr::Neq(_, lhs, rhs) = expr else {
178
79617
        return Err(RuleNotApplicable);
179
    }; // considered covered
180

            
181
24
    let (binding, _, _) =
182
36
        validate_order_int_operands(vec![lhs.as_ref().clone(), rhs.as_ref().clone()])?;
183
24
    let [lhs_bits, rhs_bits] = binding.as_slice() else {
184
        return Err(RuleNotApplicable); // consider covered
185
    };
186

            
187
24
    let (mut output, mut new_clauses, mut new_symbols) =
188
24
        sat_order_eq_expr(lhs_bits, rhs_bits, symbols);
189

            
190
24
    output = tseytin_not(output, &mut new_clauses, &mut new_symbols);
191

            
192
24
    Ok(Reduction::cnf(output, new_clauses, new_symbols))
193
79653
}
194

            
195
/// Converts a </>/<=/>= expression between two order SATInts to a boolean expression in cnf
196
///
197
/// ```text
198
/// SATInt(a) </>/<=/>= SATInt(b) ~> Bool
199
///
200
/// ```
201
/// Note: < and <= are rewritten by swapping operands to reuse lt logic.
202
#[register_rule("SAT_Order", 9100, [Lt, Gt, Leq, Geq])]
203
79653
fn ineq_sat_order(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
204
79653
    let (lhs, rhs, negate) = match expr {
205
        // A < B -> sat_order_lt(A, B)
206
54
        Expr::Lt(_, x, y) => (x, y, false),
207
        // A > B -> sat_order_lt(B, A)
208
12
        Expr::Gt(_, x, y) => (y, x, false),
209
        // A <= B -> NOT (B < A)
210
1986
        Expr::Leq(_, x, y) => (y, x, true),
211
        // A >= B -> NOT (A < B)
212
1218
        Expr::Geq(_, x, y) => (x, y, true),
213
76383
        _ => return Err(RuleNotApplicable),
214
    };
215

            
216
1464
    let (binding, _, _) =
217
3270
        validate_order_int_operands(vec![lhs.as_ref().clone(), rhs.as_ref().clone()])?;
218
1464
    let [lhs_bits, rhs_bits] = binding.as_slice() else {
219
        return Err(RuleNotApplicable);
220
    };
221

            
222
1464
    let mut new_symbols = symbols.clone();
223
1464
    let mut new_clauses = vec![];
224

            
225
1464
    let mut output = sat_order_lt(
226
1464
        lhs_bits.clone(),
227
1464
        rhs_bits.clone(),
228
1464
        &mut new_clauses,
229
1464
        &mut new_symbols,
230
    );
231

            
232
1464
    if negate {
233
1404
        output = tseytin_not(output, &mut new_clauses, &mut new_symbols);
234
1404
    }
235

            
236
1464
    Ok(Reduction::cnf(output, new_clauses, new_symbols))
237
79653
}
238

            
239
/// Converts a - expression for a SATInt to a new SATInt
240
///
241
/// ```text
242
/// -SATInt(a) ~> SATInt(b)
243
///
244
/// ```
245
#[register_rule("SAT_Order", 9100, [Neg])]
246
79653
fn neg_sat_order(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
247
79653
    let Expr::Neg(_, value) = expr else {
248
79587
        return Err(RuleNotApplicable);
249
    };
250

            
251
66
    let (binding, old_min, old_max) = validate_order_int_operands(vec![value.as_ref().clone()])?;
252
54
    let [val_bits] = binding.as_slice() else {
253
        return Err(RuleNotApplicable); // consider covered
254
    };
255

            
256
54
    let new_min = -old_max;
257
54
    let new_max = -old_min;
258

            
259
54
    let n = val_bits.len();
260
54
    let mut out: Vec<Expr> = Vec::with_capacity(n);
261

            
262
54
    let mut new_symbols = symbols.clone();
263
54
    let mut new_clauses = vec![];
264

            
265
54
    let ff = Expr::Atomic(Metadata::new(), Atom::Literal(Literal::Bool(false)));
266
54
    out.push(tseytin_not(ff, &mut new_clauses, &mut new_symbols));
267

            
268
324
    for i in 1..n {
269
324
        let src = val_bits[n - i].clone();
270
324
        let neg_bit = tseytin_not(src, &mut new_clauses, &mut new_symbols);
271
324
        out.push(neg_bit);
272
324
    }
273

            
274
54
    Ok(Reduction::cnf(
275
54
        Expr::SATInt(
276
54
            Metadata::new(),
277
54
            SATIntEncoding::Order,
278
54
            Moo::new(into_matrix_expr!(out)),
279
54
            (new_min, new_max),
280
54
        ),
281
54
        new_clauses,
282
54
        new_symbols,
283
54
    ))
284
79653
}