1
//! Comprehension expansion rules
2

            
3
mod expand_native;
4
mod expand_via_solver;
5
mod expand_via_solver_ac;
6
mod via_solver_common;
7

            
8
pub use expand_native::expand_native;
9
pub use expand_via_solver::expand_via_solver;
10
pub use expand_via_solver_ac::expand_via_solver_ac;
11

            
12
use conjure_cp::{
13
    ast::{
14
        DeclarationPtr, Domain, DomainPtr, Expression as Expr, IntVal, Moo, Name, Range, Reference,
15
        SymbolTable, UnresolvedDomain,
16
        comprehension::{Comprehension, ComprehensionQualifier},
17
        serde::{HasId, ObjId},
18
    },
19
    bug, into_matrix_expr,
20
    rule_engine::{
21
        ApplicationError::RuleNotApplicable, ApplicationResult, Reduction, register_rule,
22
    },
23
    settings::{QuantifiedExpander, comprehension_expander},
24
};
25
use std::collections::HashMap;
26
use uniplate::{Biplate, Uniplate};
27

            
28
/// Rewrite top-level `exists` comprehensions into constraints over fresh machine `find`s.
29
///
30
/// `exists` is represented as `or([comprehension])`.
31
#[register_rule("Base", 2003, [Root])]
32
325010
fn exists_quantified_to_finds(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
33
325010
    let Expr::Root(metadata, constraints) = expr else {
34
312577
        return Err(RuleNotApplicable);
35
    };
36

            
37
12433
    let mut new_symbols = symbols.clone();
38
12433
    let mut new_constraints = Vec::with_capacity(constraints.len());
39
12433
    let mut changed = false;
40

            
41
57293
    for constraint in constraints {
42
57293
        let Some(comprehension) = as_exists_comprehension(constraint) else {
43
57053
            new_constraints.push(constraint.clone());
44
57053
            continue;
45
        };
46

            
47
240
        let Some(new_constraints_for_exists) =
48
240
            rewrite_exists_comprehension_to_constraints(&comprehension, &mut new_symbols)
49
        else {
50
            new_constraints.push(constraint.clone());
51
            continue;
52
        };
53

            
54
240
        new_constraints.extend(new_constraints_for_exists);
55
240
        changed = true;
56
    }
57

            
58
12433
    if changed {
59
192
        Ok(Reduction::with_symbols(
60
192
            Expr::Root(metadata.clone(), new_constraints),
61
192
            new_symbols,
62
192
        ))
63
    } else {
64
12241
        Err(RuleNotApplicable)
65
    }
66
325010
}
67

            
68
/// Expand comprehensions using `--comprehension-expander native`.
69
#[register_rule("Base", 2000, [Comprehension])]
70
169872
fn expand_comprehension_native(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
71
169872
    if comprehension_expander() != QuantifiedExpander::Native {
72
116300
        return Err(RuleNotApplicable);
73
53572
    }
74

            
75
53572
    let Expr::Comprehension(_, comprehension) = expr else {
76
53146
        return Err(RuleNotApplicable);
77
    };
78

            
79
426
    let comprehension = comprehension.as_ref().clone();
80

            
81
744
    for qual in &comprehension.qualifiers {
82
744
        if let ComprehensionQualifier::ExpressionGenerator { .. } = qual {
83
30
            return Err(RuleNotApplicable);
84
714
        }
85
    }
86

            
87
396
    let mut symbols = symbols.clone();
88
396
    let results = expand_native(comprehension, &mut symbols)
89
        .unwrap_or_else(|e| bug!("native comprehension expansion failed: {e}"));
90
396
    Ok(Reduction::with_symbols(into_matrix_expr!(results), symbols))
91
169872
}
92

            
93
/// Expand comprehensions using `--comprehension-expander via-solver`.
94
///
95
/// Algorithm sketch:
96
/// 1. Match one comprehension node.
97
/// 2. Build a temporary generator submodel from its qualifiers/guards.
98
/// 3. Materialise quantified declarations as temporary `find` declarations.
99
/// 4. Wrap that submodel as a standalone temporary model, with search order restricted to the
100
///    quantified names.
101
/// 5. Rewrite the temporary model using the configured rewriter and Minion-oriented rules.
102
/// 6. Solve the rewritten temporary model with Minion and keep only quantified assignments from
103
///    each solution.
104
/// 7. Instantiate the original return expression under each quantified assignment.
105
/// 8. Replace the comprehension by a matrix literal containing all instantiated return values.
106
#[register_rule("Base", 2000, [Comprehension])]
107
169872
fn expand_comprehension_via_solver(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
108
53572
    if !matches!(
109
169872
        comprehension_expander(),
110
        QuantifiedExpander::ViaSolver | QuantifiedExpander::ViaSolverAc
111
    ) {
112
53572
        return Err(RuleNotApplicable);
113
116300
    }
114

            
115
116300
    let Expr::Comprehension(_, comprehension) = expr else {
116
116108
        return Err(RuleNotApplicable);
117
    };
118

            
119
192
    let comprehension = comprehension.as_ref().clone();
120

            
121
264
    for qual in &comprehension.qualifiers {
122
264
        if let ComprehensionQualifier::ExpressionGenerator { .. } = qual {
123
6
            return Err(RuleNotApplicable);
124
258
        }
125
    }
126

            
127
186
    let results = expand_via_solver(comprehension)
128
        .unwrap_or_else(|e| bug!("via-solver comprehension expansion failed: {e}"));
129
186
    Ok(Reduction::with_symbols(
130
186
        into_matrix_expr!(results),
131
186
        symbols.clone(),
132
186
    ))
133
169872
}
134

            
135
/// Expand comprehensions inside AC operators using `--comprehension-expander via-solver-ac`.
136
///
137
/// Algorithm sketch:
138
/// 1. Match an AC operator whose single child is a comprehension.
139
/// 2. Build a temporary generator submodel from the comprehension qualifiers/guards.
140
/// 3. Add a derived constraint from the return expression to this generator model:
141
///    localise non-local references, and replace non-quantified fragments with dummy variables so
142
///    the constraint depends only on locally solvable symbols.
143
/// 4. Materialise quantified declarations as temporary `find` declarations in the temporary model.
144
/// 5. Rewrite and solve the temporary model with Minion; keep only quantified assignments.
145
/// 6. Instantiate the original return expression under those assignments.
146
/// 7. Rebuild the same AC operator around the instantiated matrix literal.
147
#[register_rule("Base", 2002)]
148
316427
fn expand_comprehension_via_solver_ac(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
149
316427
    if comprehension_expander() != QuantifiedExpander::ViaSolverAc {
150
129189
        return Err(RuleNotApplicable);
151
187238
    }
152

            
153
    // Is this an ac expression?
154
187238
    let ac_operator_kind = expr.to_ac_operator_kind().ok_or(RuleNotApplicable)?;
155

            
156
11474
    let children = expr.children();
157
11474
    debug_assert_eq!(
158
11474
        children.len(),
159
        1,
160
        "AC expressions should have exactly one child."
161
    );
162

            
163
11474
    let comprehension = children
164
11474
        .front()
165
11474
        .and_then(as_single_comprehension)
166
11474
        .ok_or(RuleNotApplicable)?;
167

            
168
675
    for qual in &comprehension.qualifiers {
169
675
        if let ComprehensionQualifier::ExpressionGenerator { .. } = qual {
170
6
            return Err(RuleNotApplicable);
171
669
        }
172
    }
173

            
174
447
    let results =
175
603
        expand_via_solver_ac(comprehension, ac_operator_kind).or(Err(RuleNotApplicable))?;
176

            
177
447
    let new_expr = ac_operator_kind.as_expression(into_matrix_expr!(results));
178
447
    Ok(Reduction::with_symbols(new_expr, symbols.clone()))
179
316427
}
180

            
181
21946
fn as_single_comprehension(expr: &Expr) -> Option<Comprehension> {
182
21946
    if let Expr::Comprehension(_, comprehension) = expr {
183
849
        return Some(comprehension.as_ref().clone());
184
21097
    }
185

            
186
21097
    let exprs = expr.clone().unwrap_list()?;
187
17212
    let [Expr::Comprehension(_, comprehension)] = exprs.as_slice() else {
188
17212
        return None;
189
    };
190

            
191
    Some(comprehension.as_ref().clone())
192
21946
}
193

            
194
57293
fn as_exists_comprehension(expr: &Expr) -> Option<Comprehension> {
195
57293
    let Expr::Or(_, or_child) = expr else {
196
46821
        return None;
197
    };
198

            
199
10472
    as_single_comprehension(or_child.as_ref())
200
57293
}
201

            
202
240
fn rewrite_exists_comprehension_to_constraints(
203
240
    comprehension: &Comprehension,
204
240
    symbols: &mut SymbolTable,
205
240
) -> Option<Vec<Expr>> {
206
240
    let quantified_declarations = quantified_declarations(comprehension)?;
207

            
208
240
    let mut replacements_by_id: HashMap<ObjId, DeclarationPtr> = HashMap::new();
209
240
    let mut replacements_by_name: HashMap<Name, DeclarationPtr> = HashMap::new();
210

            
211
336
    for decl in quantified_declarations {
212
336
        let domain = decl.domain()?;
213
336
        let rewritten_domain =
214
336
            replace_declaration_ptrs_in_domain(domain, &replacements_by_id, &replacements_by_name);
215
336
        let fresh_decl = symbols.gen_find(&rewritten_domain);
216
336
        replacements_by_id.insert(decl.id(), fresh_decl.clone());
217
336
        replacements_by_name.insert(decl.name().clone(), fresh_decl);
218
    }
219

            
220
240
    let mut conjuncts = Vec::new();
221
336
    for qualifier in &comprehension.qualifiers {
222
336
        if let ComprehensionQualifier::Condition(condition) = qualifier {
223
            conjuncts.push(replace_declaration_ptrs_in_expr(
224
                condition.clone(),
225
                &replacements_by_id,
226
                &replacements_by_name,
227
            ));
228
336
        }
229
    }
230
240
    conjuncts.push(replace_declaration_ptrs_in_expr(
231
240
        comprehension.return_expression.clone(),
232
240
        &replacements_by_id,
233
240
        &replacements_by_name,
234
    ));
235

            
236
240
    Some(conjuncts)
237
240
}
238

            
239
240
fn quantified_declarations(comprehension: &Comprehension) -> Option<Vec<DeclarationPtr>> {
240
240
    let quantified_names = comprehension.quantified_vars();
241
240
    let symbols = comprehension.symbols();
242
240
    quantified_names
243
240
        .into_iter()
244
336
        .map(|name| symbols.lookup_local(&name))
245
240
        .collect()
246
240
}
247

            
248
336
fn replace_declaration_ptrs_in_expr(
249
336
    expr: Expr,
250
336
    replacements_by_id: &HashMap<ObjId, DeclarationPtr>,
251
336
    replacements_by_name: &HashMap<Name, DeclarationPtr>,
252
336
) -> Expr {
253
828
    expr.transform_bi(&|decl: DeclarationPtr| {
254
828
        if let Some(replacement) = replacements_by_id.get(&decl.id()) {
255
504
            return replacement.clone();
256
324
        }
257

            
258
324
        let name = decl.name().clone();
259
324
        replacements_by_name.get(&name).cloned().unwrap_or(decl)
260
828
    })
261
336
}
262

            
263
336
fn replace_declaration_ptrs_in_domain(
264
336
    domain: DomainPtr,
265
336
    replacements_by_id: &HashMap<ObjId, DeclarationPtr>,
266
336
    replacements_by_name: &HashMap<Name, DeclarationPtr>,
267
336
) -> DomainPtr {
268
336
    let mut rewritten = domain
269
336
        .transform_bi(&|expr: Expr| {
270
            replace_declaration_ptrs_in_expr(expr, replacements_by_id, replacements_by_name)
271
        })
272
336
        .transform_bi(&|reference: Reference| {
273
            replace_reference(reference, replacements_by_id, replacements_by_name)
274
        });
275

            
276
    // `Range<T>` does not participate in the generic biplate traversal, so recurse through
277
    // unresolved domain structure once to rewrite symbolic integer bounds.
278
336
    rewrite_int_ranges_in_domain_ptr(&mut rewritten, replacements_by_id, replacements_by_name);
279

            
280
336
    rewritten
281
336
}
282

            
283
fn replace_reference(
284
    reference: Reference,
285
    replacements_by_id: &HashMap<ObjId, DeclarationPtr>,
286
    replacements_by_name: &HashMap<Name, DeclarationPtr>,
287
) -> Reference {
288
    let replacement = replacements_by_id
289
        .get(&reference.ptr().id())
290
        .cloned()
291
        .or_else(|| {
292
            let name = reference.name().clone();
293
            replacements_by_name.get(&name).cloned()
294
        });
295

            
296
    replacement.map(Reference::new).unwrap_or(reference)
297
}
298

            
299
336
fn rewrite_int_ranges_in_domain_ptr(
300
336
    domain: &mut DomainPtr,
301
336
    replacements_by_id: &HashMap<ObjId, DeclarationPtr>,
302
336
    replacements_by_name: &HashMap<Name, DeclarationPtr>,
303
336
) {
304
336
    let mut rewritten = domain.as_ref().clone();
305
336
    rewrite_int_ranges_in_domain(&mut rewritten, replacements_by_id, replacements_by_name);
306
336
    *domain = Moo::new(rewritten);
307
336
}
308

            
309
336
fn rewrite_int_ranges_in_domain(
310
336
    domain: &mut Domain,
311
336
    replacements_by_id: &HashMap<ObjId, DeclarationPtr>,
312
336
    replacements_by_name: &HashMap<Name, DeclarationPtr>,
313
336
) {
314
336
    let Domain::Unresolved(unresolved) = domain else {
315
240
        return;
316
    };
317

            
318
96
    rewrite_int_ranges_in_unresolved_domain(
319
96
        Moo::make_mut(unresolved),
320
96
        replacements_by_id,
321
96
        replacements_by_name,
322
    );
323
336
}
324

            
325
96
fn rewrite_int_ranges_in_unresolved_domain(
326
96
    unresolved: &mut UnresolvedDomain,
327
96
    replacements_by_id: &HashMap<ObjId, DeclarationPtr>,
328
96
    replacements_by_name: &HashMap<Name, DeclarationPtr>,
329
96
) {
330
96
    match unresolved {
331
96
        UnresolvedDomain::Int(ranges) => {
332
96
            for range in ranges {
333
96
                rewrite_int_range(range, replacements_by_id, replacements_by_name);
334
96
            }
335
        }
336
        UnresolvedDomain::Set(attr, inner) => {
337
            rewrite_int_range(&mut attr.size, replacements_by_id, replacements_by_name);
338
            rewrite_int_ranges_in_domain_ptr(inner, replacements_by_id, replacements_by_name);
339
        }
340
        UnresolvedDomain::MSet(attr, inner) => {
341
            rewrite_int_range(&mut attr.size, replacements_by_id, replacements_by_name);
342
            rewrite_int_range(
343
                &mut attr.occurrence,
344
                replacements_by_id,
345
                replacements_by_name,
346
            );
347
            rewrite_int_ranges_in_domain_ptr(inner, replacements_by_id, replacements_by_name);
348
        }
349
        UnresolvedDomain::Matrix(inner, index_domains) => {
350
            rewrite_int_ranges_in_domain_ptr(inner, replacements_by_id, replacements_by_name);
351
            for index_domain in index_domains {
352
                rewrite_int_ranges_in_domain_ptr(
353
                    index_domain,
354
                    replacements_by_id,
355
                    replacements_by_name,
356
                );
357
            }
358
        }
359
        UnresolvedDomain::Tuple(inner_domains) => {
360
            for inner_domain in inner_domains {
361
                rewrite_int_ranges_in_domain_ptr(
362
                    inner_domain,
363
                    replacements_by_id,
364
                    replacements_by_name,
365
                );
366
            }
367
        }
368
        UnresolvedDomain::Partition(attr, inner) => {
369
            rewrite_int_range(
370
                &mut attr.num_parts,
371
                replacements_by_id,
372
                replacements_by_name,
373
            );
374
            rewrite_int_range(&mut attr.part_len, replacements_by_id, replacements_by_name);
375
            rewrite_int_ranges_in_domain_ptr(inner, replacements_by_id, replacements_by_name);
376
        }
377
        UnresolvedDomain::Sequence(attr, inner) => {
378
            rewrite_int_range(&mut attr.size, replacements_by_id, replacements_by_name);
379
            rewrite_int_ranges_in_domain_ptr(inner, replacements_by_id, replacements_by_name);
380
        }
381
        UnresolvedDomain::Reference(_) => {}
382
        UnresolvedDomain::Record(entries) => {
383
            for entry in entries {
384
                rewrite_int_ranges_in_domain_ptr(
385
                    &mut entry.value,
386
                    replacements_by_id,
387
                    replacements_by_name,
388
                );
389
            }
390
        }
391
        UnresolvedDomain::Function(attr, domain, codomain) => {
392
            rewrite_int_range(&mut attr.size, replacements_by_id, replacements_by_name);
393
            rewrite_int_ranges_in_domain_ptr(domain, replacements_by_id, replacements_by_name);
394
            rewrite_int_ranges_in_domain_ptr(codomain, replacements_by_id, replacements_by_name);
395
        }
396
        UnresolvedDomain::Variant(entries) => {
397
            for entry in entries {
398
                rewrite_int_ranges_in_domain_ptr(
399
                    &mut entry.value,
400
                    replacements_by_id,
401
                    replacements_by_name,
402
                );
403
            }
404
        }
405
        UnresolvedDomain::Relation(attr, domains) => {
406
            rewrite_int_range(&mut attr.size, replacements_by_id, replacements_by_name);
407
            for domain in domains {
408
                rewrite_int_ranges_in_domain_ptr(domain, replacements_by_id, replacements_by_name);
409
            }
410
        }
411
    }
412
96
}
413

            
414
96
fn rewrite_int_range(
415
96
    range: &mut Range<IntVal>,
416
96
    replacements_by_id: &HashMap<ObjId, DeclarationPtr>,
417
96
    replacements_by_name: &HashMap<Name, DeclarationPtr>,
418
96
) {
419
96
    match range {
420
        Range::Single(value) | Range::UnboundedL(value) | Range::UnboundedR(value) => {
421
            rewrite_int_value(value, replacements_by_id, replacements_by_name);
422
        }
423
96
        Range::Bounded(lower, upper) => {
424
96
            rewrite_int_value(lower, replacements_by_id, replacements_by_name);
425
96
            rewrite_int_value(upper, replacements_by_id, replacements_by_name);
426
96
        }
427
        Range::Unbounded => {}
428
    }
429
96
}
430

            
431
192
fn rewrite_int_value(
432
192
    int_val: &mut IntVal,
433
192
    replacements_by_id: &HashMap<ObjId, DeclarationPtr>,
434
192
    replacements_by_name: &HashMap<Name, DeclarationPtr>,
435
192
) {
436
192
    if let IntVal::Expr(expr) = int_val {
437
96
        let rewritten = replace_declaration_ptrs_in_expr(
438
96
            (**expr).clone(),
439
96
            replacements_by_id,
440
96
            replacements_by_name,
441
96
        );
442
96
        *expr = Moo::new(rewritten);
443
96
    }
444
192
}