1
use conjure_cp::{
2
    ast::{
3
        Atom, DeclarationKind, DeclarationPtr, DomainPtr, Expression, Literal, Metadata, Name,
4
        SymbolTable,
5
        comprehension::{Comprehension, ComprehensionQualifier},
6
        eval_constant,
7
    },
8
    solver::SolverError,
9
};
10
use uniplate::Biplate as _;
11

            
12
use super::via_solver_common::{lift_machine_references_into_parent_scope, simplify_expression};
13

            
14
/// Expands the comprehension without calling an external solver.
15
///
16
/// Algorithm:
17
/// 1. Recurse qualifiers left-to-right.
18
/// 2. For each generator value, temporarily bind the quantified declaration to a
19
///    `TemporaryValueLetting` and recurse.
20
/// 3. For each condition, evaluate and recurse only if true.
21
/// 4. At the leaf, evaluate the return expression under the active bindings.
22
108
pub fn expand_native(
23
108
    comprehension: Comprehension,
24
108
    parent_symbols: &mut SymbolTable,
25
108
) -> Result<Vec<Expression>, SolverError> {
26
108
    let mut expanded = Vec::new();
27
108
    expand_qualifiers(&comprehension, 0, &mut expanded, parent_symbols)?;
28
108
    Ok(expanded)
29
108
}
30

            
31
7956
fn expand_qualifiers(
32
7956
    comprehension: &Comprehension,
33
7956
    qualifier_index: usize,
34
7956
    expanded: &mut Vec<Expression>,
35
7956
    parent_symbols: &mut SymbolTable,
36
7956
) -> Result<(), SolverError> {
37
7956
    if qualifier_index == comprehension.qualifiers.len() {
38
6672
        let child_symbols = comprehension.symbols().clone();
39
6672
        let return_expression =
40
6672
            concretise_resolved_reference_atoms(comprehension.return_expression.clone());
41
6672
        let return_expression = simplify_expression(return_expression);
42
6672
        let return_expression = lift_machine_references_into_parent_scope(
43
6672
            return_expression,
44
6672
            &child_symbols,
45
6672
            parent_symbols,
46
        );
47
6672
        expanded.push(return_expression);
48
6672
        return Ok(());
49
1284
    }
50

            
51
1284
    match &comprehension.qualifiers[qualifier_index] {
52
1164
        ComprehensionQualifier::Generator { ptr } => {
53
1164
            let name = ptr.name().clone();
54
1164
            let domain = ptr.domain().expect("generator declaration has domain");
55
1164
            let values = resolve_generator_values(&name, &domain)?;
56

            
57
7782
            for literal in values {
58
7782
                with_temporary_quantified_binding(ptr, &literal, || {
59
7782
                    expand_qualifiers(comprehension, qualifier_index + 1, expanded, parent_symbols)
60
7782
                })?;
61
            }
62
        }
63
120
        ComprehensionQualifier::Condition(condition) => {
64
120
            if evaluate_guard(condition)? {
65
66
                expand_qualifiers(comprehension, qualifier_index + 1, expanded, parent_symbols)?;
66
54
            }
67
        }
68
    }
69

            
70
1284
    Ok(())
71
7956
}
72

            
73
1164
fn resolve_generator_values(name: &Name, domain: &DomainPtr) -> Result<Vec<Literal>, SolverError> {
74
1164
    let resolved = domain.resolve().ok_or_else(|| {
75
        SolverError::ModelFeatureNotSupported(format!(
76
            "quantified variable '{name}' has unresolved domain after assigning previous generators: {domain}"
77
        ))
78
    })?;
79

            
80
1164
    resolved.values().map(|iter| iter.collect()).map_err(|err| {
81
        SolverError::ModelFeatureNotSupported(format!(
82
            "quantified variable '{name}' has non-enumerable domain: {err}"
83
        ))
84
    })
85
1164
}
86

            
87
7782
fn with_temporary_quantified_binding<T>(
88
7782
    quantified: &DeclarationPtr,
89
7782
    value: &Literal,
90
7782
    f: impl FnOnce() -> Result<T, SolverError>,
91
7782
) -> Result<T, SolverError> {
92
7782
    let mut targets = vec![quantified.clone()];
93
7782
    if let DeclarationKind::Quantified(inner) = &*quantified.kind()
94
7782
        && let Some(generator) = inner.generator()
95
    {
96
        targets.push(generator.clone());
97
7782
    }
98

            
99
7782
    let mut originals = Vec::with_capacity(targets.len());
100
7782
    for mut target in targets {
101
7782
        let old_kind = target.replace_kind(DeclarationKind::TemporaryValueLetting(
102
7782
            Expression::Atomic(Metadata::new(), Atom::Literal(value.clone())),
103
7782
        ));
104
7782
        originals.push((target, old_kind));
105
7782
    }
106

            
107
7782
    let result = f();
108

            
109
7782
    for (mut target, old_kind) in originals.into_iter().rev() {
110
7782
        let _ = target.replace_kind(old_kind);
111
7782
    }
112

            
113
7782
    result
114
7782
}
115

            
116
120
fn evaluate_guard(guard: &Expression) -> Result<bool, SolverError> {
117
120
    let simplified = simplify_expression(guard.clone());
118
120
    match eval_constant(&simplified) {
119
120
        Some(Literal::Bool(value)) => Ok(value),
120
        Some(other) => Err(SolverError::ModelInvalid(format!(
121
            "native comprehension guard must evaluate to Bool, got {other}: {guard}"
122
        ))),
123
        None => Err(SolverError::ModelInvalid(format!(
124
            "native comprehension expansion could not evaluate guard: {guard}"
125
        ))),
126
    }
127
120
}
128

            
129
6672
fn concretise_resolved_reference_atoms(expr: Expression) -> Expression {
130
133962
    expr.transform_bi(&|atom: Atom| match atom {
131
115872
        Atom::Reference(reference) => reference
132
115872
            .resolve_constant()
133
115872
            .map_or_else(|| Atom::Reference(reference), Atom::Literal),
134
18090
        other => other,
135
133962
    })
136
6672
}