1
use crate::errors::{FatalParseError, RecoverableParseError};
2
use crate::expression::parse_expression;
3
use crate::field;
4
use crate::parser::ParseContext;
5
use crate::util::{TypecheckingContext, named_children};
6
use conjure_cp_core::ast::{
7
    Atom, DeclarationKind, Expression, Metadata, Moo, ReturnType, Typeable,
8
};
9
use conjure_cp_core::into_matrix_expr;
10
use tree_sitter::Node;
11
use uniplate::Uniplate;
12

            
13
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14
enum ParetoDirection {
15
    Minimising,
16
    Maximising,
17
}
18

            
19
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20
enum ReferenceRewriteAction {
21
    LeaveAsIs,
22
    ExpandValueLetting,
23
    WrapInFromSolution,
24
}
25

            
26
14835
pub fn parse_dominance_relation(
27
14835
    ctx: &mut ParseContext,
28
14835
    node: &Node,
29
14835
) -> Result<Option<Expression>, FatalParseError> {
30
14835
    if ctx.root.kind() == "dominance_relation" {
31
        ctx.record_error(RecoverableParseError::new(
32
            "Nested dominance relations are not allowed".to_string(),
33
            Some(node.range()),
34
        ));
35
        return Ok(None);
36
14835
    }
37

            
38
14835
    let Some(inner_node) = field!(recover, ctx, node, "expression") else {
39
        return Ok(None);
40
    };
41

            
42
    // Create a nested context so downstream parsing knows it is inside a dominance relation.
43
14835
    let mut inner_ctx = ParseContext {
44
14835
        source_code: ctx.source_code,
45
14835
        root: node,
46
14835
        symbols: ctx.symbols.clone(),
47
14835
        errors: ctx.errors,
48
14835
        source_map: &mut *ctx.source_map,
49
14835
        decl_spans: ctx.decl_spans,
50
14835
        typechecking_context: TypecheckingContext::Unknown,
51
14835
        inner_typechecking_context: TypecheckingContext::Unknown,
52
14835
    };
53

            
54
14835
    let Some(inner) = parse_expression(&mut inner_ctx, inner_node)? else {
55
        return Ok(None);
56
    };
57

            
58
14835
    Ok(Some(Expression::DominanceRelation(
59
14835
        Metadata::new(),
60
14835
        Moo::new(inner),
61
14835
    )))
62
14835
}
63

            
64
9027
pub fn parse_pareto_expression(
65
9027
    ctx: &mut ParseContext,
66
9027
    node: &Node,
67
9027
) -> Result<Option<Expression>, FatalParseError> {
68
9027
    if ctx.root.kind() != "dominance_relation" {
69
        ctx.record_error(RecoverableParseError::new(
70
            "pareto(...) only allowed inside dominance relations".to_string(),
71
            Some(node.range()),
72
        ));
73
        return Ok(None);
74
9027
    }
75

            
76
9027
    let mut non_worsening = Vec::new();
77
9027
    let mut strict_improvements = Vec::new();
78
9027
    let components = field!(node, "components");
79

            
80
9027
    if components.kind() != "pareto_items" {
81
        return Err(FatalParseError::internal_error(
82
            format!("Unexpected pareto component list: '{}'", components.kind()),
83
            Some(components.range()),
84
        ));
85
9027
    }
86

            
87
9027
    for item_node in named_children(&components) {
88
18052
        let direction_node = field!(item_node, "direction");
89
        let direction_str =
90
18052
            &ctx.source_code[direction_node.start_byte()..direction_node.end_byte()];
91
18052
        let direction = match direction_str {
92
18051
            "minimising" => ParetoDirection::Minimising,
93
1
            "maximising" => ParetoDirection::Maximising,
94
            _ => {
95
                return Err(FatalParseError::internal_error(
96
                    format!("Unexpected pareto direction: '{direction_str}'"),
97
                    Some(direction_node.range()),
98
                ));
99
            }
100
        };
101

            
102
18052
        let component_node = field!(item_node, "expression");
103
18052
        let Some(component_expr) = parse_pareto_component(ctx, &component_node)? else {
104
            return Ok(None);
105
        };
106
        let Some((non_worse, strict)) =
107
18052
            build_pareto_constraints(ctx, &component_node, component_expr, direction)
108
        else {
109
            return Ok(None);
110
        };
111
18052
        non_worsening.push(non_worse);
112
18052
        strict_improvements.push(strict);
113
    }
114

            
115
9027
    let mut conjuncts = non_worsening;
116
9027
    conjuncts.push(combine_with_and_or(strict_improvements, true));
117

            
118
9027
    Ok(Some(combine_with_and_or(conjuncts, false)))
119
9027
}
120

            
121
18052
fn parse_pareto_component(
122
18052
    ctx: &mut ParseContext,
123
18052
    node: &Node,
124
18052
) -> Result<Option<Expression>, FatalParseError> {
125
18052
    let saved_context = ctx.typechecking_context;
126
18052
    let saved_inner_context = ctx.inner_typechecking_context;
127
18052
    ctx.typechecking_context = TypecheckingContext::Unknown;
128
18052
    ctx.inner_typechecking_context = TypecheckingContext::Unknown;
129

            
130
18052
    let parsed = parse_expression(ctx, *node);
131

            
132
18052
    ctx.typechecking_context = saved_context;
133
18052
    ctx.inner_typechecking_context = saved_inner_context;
134
18052
    parsed
135
18052
}
136

            
137
18052
fn build_pareto_constraints(
138
18052
    ctx: &mut ParseContext,
139
18052
    node: &Node,
140
18052
    component: Expression,
141
18052
    direction: ParetoDirection,
142
18052
) -> Option<(Expression, Expression)> {
143
18052
    if component
144
18052
        .universe()
145
18052
        .iter()
146
114031
        .any(|expr| matches!(expr, Expression::FromSolution(_, _)))
147
    {
148
        ctx.record_error(RecoverableParseError::new(
149
            "pareto(...) components cannot contain fromSolution(...) explicitly".to_string(),
150
            Some(node.range()),
151
        ));
152
        return None;
153
18052
    }
154

            
155
18052
    let current = expand_value_lettings(&component);
156
18052
    let previous = lift_to_previous_solution(&current);
157

            
158
18052
    match current.return_type() {
159
        ReturnType::Bool => Some(match direction {
160
            ParetoDirection::Minimising => (
161
                Expression::Imply(
162
                    Metadata::new(),
163
                    Moo::new(current.clone()),
164
                    Moo::new(previous.clone()),
165
                ),
166
                combine_with_and_or(
167
                    vec![
168
                        Expression::Not(Metadata::new(), Moo::new(current)),
169
                        previous,
170
                    ],
171
                    false,
172
                ),
173
            ),
174
            ParetoDirection::Maximising => (
175
                Expression::Imply(
176
                    Metadata::new(),
177
                    Moo::new(previous.clone()),
178
                    Moo::new(current.clone()),
179
                ),
180
                combine_with_and_or(
181
                    vec![
182
                        current,
183
                        Expression::Not(Metadata::new(), Moo::new(previous)),
184
                    ],
185
                    false,
186
                ),
187
            ),
188
        }),
189
18052
        ReturnType::Int => Some(match direction {
190
18051
            ParetoDirection::Minimising => (
191
18051
                Expression::Leq(
192
18051
                    Metadata::new(),
193
18051
                    Moo::new(current.clone()),
194
18051
                    Moo::new(previous.clone()),
195
18051
                ),
196
18051
                Expression::Lt(Metadata::new(), Moo::new(current), Moo::new(previous)),
197
18051
            ),
198
1
            ParetoDirection::Maximising => (
199
1
                Expression::Geq(
200
1
                    Metadata::new(),
201
1
                    Moo::new(current.clone()),
202
1
                    Moo::new(previous.clone()),
203
1
                ),
204
1
                Expression::Gt(Metadata::new(), Moo::new(current), Moo::new(previous)),
205
1
            ),
206
        }),
207
        found => {
208
            ctx.record_error(RecoverableParseError::new(
209
                format!(
210
                    "pareto(...) only supports int or bool components, found '{}'",
211
                    found
212
                ),
213
                Some(node.range()),
214
            ));
215
            None
216
        }
217
    }
218
18052
}
219

            
220
18052
fn expand_value_lettings(expr: &Expression) -> Expression {
221
18052
    rewrite_references(expr, false)
222
18052
}
223

            
224
18052
fn lift_to_previous_solution(expr: &Expression) -> Expression {
225
18052
    rewrite_references(expr, true)
226
18052
}
227

            
228
36104
fn rewrite_references(expr: &Expression, to_previous_solution: bool) -> Expression {
229
36104
    let mut lifted = expr.clone();
230

            
231
    loop {
232
342093
        let next = lifted.rewrite(&|subexpr| match subexpr {
233
            Expression::Atomic(_, Atom::Reference(ref reference)) => {
234
                let action = {
235
66824
                    let kind = reference.ptr.kind();
236
66824
                    match &*kind {
237
                        DeclarationKind::Find(_) | DeclarationKind::FindAuxiliary(_)
238
66824
                            if to_previous_solution =>
239
                        {
240
33412
                            ReferenceRewriteAction::WrapInFromSolution
241
                        }
242
                        DeclarationKind::Find(_) | DeclarationKind::FindAuxiliary(_) => {
243
33412
                            ReferenceRewriteAction::LeaveAsIs
244
                        }
245
                        DeclarationKind::ValueLetting(_, _)
246
                        | DeclarationKind::TemporaryValueLetting(_) => {
247
                            ReferenceRewriteAction::ExpandValueLetting
248
                        }
249
                        DeclarationKind::Given(_)
250
                        | DeclarationKind::Quantified(_)
251
                        | DeclarationKind::QuantifiedExpr(_)
252
                        | DeclarationKind::DomainLetting(_)
253
                        | _ => ReferenceRewriteAction::LeaveAsIs,
254
                    }
255
                };
256

            
257
66824
                match action {
258
33412
                    ReferenceRewriteAction::LeaveAsIs => Some(subexpr),
259
                    ReferenceRewriteAction::ExpandValueLetting => reference.resolve_expression(),
260
33412
                    ReferenceRewriteAction::WrapInFromSolution => Some(Expression::FromSolution(
261
33412
                        Metadata::new(),
262
33412
                        Moo::new(Atom::Reference(reference.clone())),
263
33412
                    )),
264
                }
265
            }
266
275269
            _ => Some(subexpr),
267
342093
        });
268

            
269
54156
        if next == lifted {
270
36104
            return lifted;
271
18052
        }
272

            
273
18052
        lifted = next;
274
    }
275
36104
}
276

            
277
18054
fn combine_with_and_or(exprs: Vec<Expression>, is_or: bool) -> Expression {
278
18054
    match exprs.len() {
279
        0 => {
280
            if is_or {
281
                Expression::Or(Metadata::new(), Moo::new(into_matrix_expr![exprs]))
282
            } else {
283
                Expression::And(Metadata::new(), Moo::new(into_matrix_expr![exprs]))
284
            }
285
        }
286
2
        1 => match exprs.into_iter().next() {
287
2
            Some(expr) => expr,
288
            None => unreachable!("vector length already checked"),
289
        },
290
        _ => {
291
18052
            if is_or {
292
9025
                Expression::Or(Metadata::new(), Moo::new(into_matrix_expr![exprs]))
293
            } else {
294
9027
                Expression::And(Metadata::new(), Moo::new(into_matrix_expr![exprs]))
295
            }
296
        }
297
    }
298
18054
}