1
use std::sync::{Arc, Mutex};
2

            
3
use conjure_cp::{
4
    ast::{
5
        DecisionVariable, DeclarationKind, Expression, Model, Name, comprehension::Comprehension,
6
    },
7
    rule_engine::resolve_rule_sets,
8
    settings::{SolverFamily, current_rewriter},
9
    solver::{Solver, SolverError, adaptors::Minion},
10
};
11

            
12
use super::via_solver_common::{
13
    instantiate_return_expressions_from_values, retain_quantified_solution_values,
14
    rewrite_model_with_configured_rewriter, with_temporary_model,
15
};
16

            
17
/// Expands the comprehension by solving quantified variables with Minion.
18
///
19
/// This returns one expression per assignment to quantified variables that satisfies the static
20
/// guards of the comprehension.
21
126
pub fn expand_via_solver(comprehension: Comprehension) -> Result<Vec<Expression>, SolverError> {
22
126
    let minion = Solver::new(Minion::new());
23
126
    let quantified_vars = comprehension.quantified_vars();
24

            
25
126
    let mut generator_model = comprehension.to_generator_model();
26
126
    materialise_quantified_finds(&mut generator_model, &quantified_vars);
27

            
28
    // only branch on the quantified variables.
29
126
    let generator_model = with_temporary_model(generator_model, Some(quantified_vars.clone()));
30
126
    tracing::trace!(
31
        target: "rule_engine_human",
32
        "Via-solver generator model before rewriting:\n\n{}\n--\n",
33
        generator_model
34
    );
35

            
36
    // call rewrite here as well as in expand_via_solver_ac, just to be consistent
37
126
    let extra_rule_sets = &["Base", "Constant", "Bubble"];
38

            
39
126
    let rule_sets = resolve_rule_sets(SolverFamily::Minion, extra_rule_sets).unwrap();
40
126
    let configured_rewriter = current_rewriter();
41

            
42
126
    let generator_model =
43
126
        rewrite_model_with_configured_rewriter(generator_model, &rule_sets, configured_rewriter);
44

            
45
    // Call the rewriter to rewrite inside the comprehension
46
    //
47
    // The original idea was to let the top level rewriter rewrite the return expression model
48
    // and the generator model. The comprehension wouldn't be expanded until the generator
49
    // model is in valid minion that can be ran, at which point the return expression model
50
    // should also be in valid minion.
51
    //
52
    // By calling the rewriter inside the rule, we no longer wait for the generator model to be
53
    // valid Minion, so we don't get the simplified return model either...
54
    //
55
    // We need to do this as we want to modify the generator model (add the dummy Z's) then
56
    // solve and return in one go.
57
    //
58
    // Comprehensions need a big rewrite soon, as theres lots of sharp edges such as this in
59
    // my original implementation, and I don't think we can fit our new optimisation into it.
60
    // If we wanted to avoid calling the rewriter, we would need to run the first half the rule
61
    // up to adding the return expr to the generator model, yield, then come back later to
62
    // actually solve it?
63

            
64
    // Keep return expressions unreduced until quantified assignments are substituted.
65
    // Rewriting before substitution can change guard structure in ways that are unsafe for
66
    // constant evaluation after instantiation.
67
126
    let return_expression_model =
68
126
        with_temporary_model(comprehension.to_return_expression_model(), None);
69

            
70
126
    let values = {
71
126
        let solver_model = generator_model.clone();
72

            
73
126
        let minion = minion.load_model(solver_model)?;
74

            
75
126
        let values = Arc::new(Mutex::new(Vec::new()));
76
126
        let values_ptr = Arc::clone(&values);
77
126
        let quantified_vars_for_solution = quantified_vars.clone();
78

            
79
        tracing::debug!(model=%generator_model,comprehension=%comprehension,"Minion solving comprehension (solver mode)");
80
6462
        minion.solve(Box::new(move |sols| {
81
            // Only keep quantified assignments; discard solver auxiliaries/locals.
82
6462
            let values = &mut *values_ptr.lock().unwrap();
83
6462
            values.push(retain_quantified_solution_values(
84
6462
                sols,
85
6462
                &quantified_vars_for_solution,
86
            ));
87
6462
            true
88
6462
        }))?;
89

            
90
126
        values.lock().unwrap().clone()
91
    };
92
126
    Ok(instantiate_return_expressions_from_values(
93
126
        values,
94
126
        &return_expression_model,
95
126
        &quantified_vars,
96
126
    ))
97
126
}
98

            
99
126
fn materialise_quantified_finds(model: &mut Model, quantified_vars: &[Name]) {
100
126
    let symbols = model.symbols().clone();
101

            
102
162
    for name in quantified_vars {
103
162
        let Some(mut decl) = symbols.lookup_local(name) else {
104
            continue;
105
        };
106
162
        let Some(domain) = decl.domain() else {
107
            continue;
108
        };
109

            
110
162
        let _ = decl.replace_kind(DeclarationKind::Find(DecisionVariable::new(domain)));
111
    }
112
126
}