Skip to main content

conjure_cp_essence_parser/parser/
expression.rs

1use crate::diagnostics::diagnostics_api::SymbolKind;
2use crate::errors::{FatalParseError, RecoverableParseError};
3use crate::parser::ParseContext;
4use crate::parser::atom::{parse_atom, parse_table_operand};
5use crate::parser::comprehension::parse_quantifier_or_aggregate_expr;
6use crate::parser::domain::parse_domain;
7use crate::parser::global_constraints::{
8    ALL_DIFFERENT, ALL_DIFFERENT_EXCEPT, AT_LEAST, AT_MOST, GLOBAL_CARDINALITY,
9    is_all_different_except_operator, is_all_different_operator, is_at_least_operator,
10    is_at_most_operator, is_global_cardinality_operator,
11};
12use crate::util::TypecheckingContext;
13use crate::{child, field, named_child};
14use conjure_cp_core::ast::ac_operators::ACOperatorKind;
15use conjure_cp_core::ast::{
16    Atom, Expression, GroundDomain, Literal, Metadata, Moo, ReturnType, Typeable,
17};
18use conjure_cp_core::{domain_int, matrix_expr, range};
19use tree_sitter::Node;
20
21pub fn parse_expression(
22    ctx: &mut ParseContext,
23    node: Node,
24) -> Result<Option<Expression>, FatalParseError> {
25    match node.kind() {
26        "atom"
27        | "primary_atom"
28        | "constant"
29        | "identifier"
30        | "metavar"
31        | "matrix"
32        | "record"
33        | "variant"
34        | "tuple"
35        | "set_literal"
36        | "mset_literal"
37        | "sequence_literal"
38        | "function_literal"
39        | "relation_literal"
40        | "partition_literal"
41        | "comprehension"
42        | "index_or_slice"
43        | "flatten"
44        | "element_id"
45        | "table"
46        | "negative_table"
47        | "apply_expr"
48        | "image_expr"
49        | "image_set_expr"
50        | "pre_image_expr"
51        | "inverse_expr"
52        | "restrict_expr"
53        | "defined_expr"
54        | "range_expr"
55        | "to_set_expr"
56        | "to_mset_expr"
57        | "to_relation_expr"
58        | "attribute_as_constraint_expr"
59        | "together_expr"
60        | "apart_expr"
61        | "party_expr"
62        | "parts_expr"
63        | "participants_expr" => parse_atom(ctx, &node),
64        "bool_expr" => {
65            if ctx.typechecking_context == TypecheckingContext::Arithmetic {
66                ctx.record_error(RecoverableParseError::new(
67                    format!(
68                        "Type error: {}\n\tExpected: int\n\tGot: boolean expression",
69                        &ctx.source_code[node.start_byte()..node.end_byte()]
70                    ),
71                    Some(node.range()),
72                ));
73                return Ok(None);
74            }
75            parse_boolean_expression(ctx, &node)
76        }
77        "arithmetic_expr" => {
78            if ctx.typechecking_context == TypecheckingContext::Boolean {
79                ctx.record_error(RecoverableParseError::new(
80                    format!(
81                        "Type error: {}\n\tExpected: bool\n\tGot: int",
82                        &ctx.source_code[node.start_byte()..node.end_byte()]
83                    ),
84                    Some(node.range()),
85                ));
86                return Ok(None);
87            }
88            parse_arithmetic_expression(ctx, &node)
89        }
90        "comparison_expr" => {
91            if ctx.typechecking_context == TypecheckingContext::Arithmetic {
92                ctx.record_error(RecoverableParseError::new(
93                    format!(
94                        "Type error: {}\n\tExpected: int\n\tGot: comparison expression",
95                        &ctx.source_code[node.start_byte()..node.end_byte()]
96                    ),
97                    Some(node.range()),
98                ));
99                return Ok(None);
100            }
101            parse_comparison_expression(ctx, &node)
102        }
103        "all_diff_comparison" => {
104            if ctx.typechecking_context == TypecheckingContext::Arithmetic {
105                ctx.record_error(RecoverableParseError::new(
106                    format!("Type error: {}\n\tExpected: arithmetic expression\n\tGot: comparison expression", &ctx.source_code[node.start_byte()..node.end_byte()]),
107                    Some(node.range()),
108                ));
109                return Ok(None);
110            }
111            ctx.typechecking_context = TypecheckingContext::Matrix;
112            parse_all_diff_comparison(ctx, &node)
113        }
114        "all_different_except_comparison" => {
115            if ctx.typechecking_context == TypecheckingContext::Arithmetic {
116                ctx.record_error(RecoverableParseError::new(
117                    format!("Type error: {}\n\tExpected: arithmetic expression\n\tGot: comparison expression", &ctx.source_code[node.start_byte()..node.end_byte()]),
118                    Some(node.range()),
119                ));
120                return Ok(None);
121            }
122            ctx.typechecking_context = TypecheckingContext::Matrix;
123            parse_all_different_except_comparison(ctx, &node)
124        }
125        "global_cardinality_comparison" => {
126            if ctx.typechecking_context == TypecheckingContext::Arithmetic {
127                ctx.record_error(RecoverableParseError::new(
128                    format!("Type error: {}\n\tExpected: arithmetic expression\n\tGot: comparison expression", &ctx.source_code[node.start_byte()..node.end_byte()]),
129                    Some(node.range()),
130                ));
131                return Ok(None);
132            }
133            parse_global_cardinality_comparison(ctx, &node)
134        }
135        "annotation_expr" | "type_annotation" | "domain_annotation" => {
136            parse_annotation_expression(ctx, &node)
137        }
138        _ => {
139            ctx.record_error(RecoverableParseError::new(
140                format!("Unexpected expression type: '{}'", node.kind()),
141                Some(node.range()),
142            ));
143            Ok(None)
144        }
145    }
146}
147
148fn parse_arithmetic_expression(
149    ctx: &mut ParseContext,
150    node: &Node,
151) -> Result<Option<Expression>, FatalParseError> {
152    ctx.typechecking_context = TypecheckingContext::Arithmetic;
153    ctx.inner_typechecking_context = TypecheckingContext::Unknown;
154    let Some(inner) = named_child!(recover, ctx, node) else {
155        return Ok(None);
156    };
157    match inner.kind() {
158        "atom" => parse_atom(ctx, &inner),
159        "negative_expr" | "abs_value" | "sub_arith_expr" | "factorial_expr" => {
160            parse_unary_expression(ctx, &inner)
161        }
162        "toInt_expr" => {
163            // add special handling for toInt, as it is arithmetic but takes a non-arithmetic operand
164            ctx.typechecking_context = TypecheckingContext::Unknown;
165            parse_unary_expression(ctx, &inner)
166        }
167        "catch_undef_expr" => parse_catch_undef_expression(ctx, &inner),
168        "exponent" | "product_expr" | "sum_expr" => parse_binary_expression(ctx, &inner),
169        "list_combining_expr_arith" => {
170            // list-combining arithmetic operators accept either set or matrix operands
171            ctx.typechecking_context = TypecheckingContext::SetOrMatrix;
172
173            // set inner context to arithmetic to ensure elements of list are arithmetic expressions
174            ctx.inner_typechecking_context = TypecheckingContext::Arithmetic;
175            parse_list_combining_expression(ctx, &inner)
176        }
177        "aggregate_expr" => {
178            ctx.inner_typechecking_context = TypecheckingContext::Arithmetic;
179            parse_quantifier_or_aggregate_expr(ctx, &inner)
180        }
181        _ => {
182            ctx.record_error(RecoverableParseError::new(
183                format!("Expected arithmetic expression, found: {}", inner.kind()),
184                Some(inner.range()),
185            ));
186            Ok(None)
187        }
188    }
189}
190
191fn parse_comparison_expression(
192    ctx: &mut ParseContext,
193    node: &Node,
194) -> Result<Option<Expression>, FatalParseError> {
195    let Some(inner) = named_child!(recover, ctx, node) else {
196        return Ok(None);
197    };
198    match inner.kind() {
199        "arithmetic_comparison" => {
200            // <, <=, >, >= work on any orderable type (everything except unnamed types, which
201            // don't support ordering); typechecking of operands is handled within
202            // parse_binary_expression, same as equality.
203            ctx.typechecking_context = TypecheckingContext::Unknown;
204            parse_binary_expression(ctx, &inner)
205        }
206        "lex_comparison" => {
207            // <lex, <=lex, >lex, >=lex are specialised to matrices; other types use the plain
208            // ordering operators above instead.
209            ctx.typechecking_context = TypecheckingContext::Matrix;
210            parse_binary_expression(ctx, &inner)
211        }
212        "equality_comparison" => {
213            // Equality works on any type, typechecking of operands will be handled within parse_binary_expression
214            ctx.typechecking_context = TypecheckingContext::Unknown;
215            parse_binary_expression(ctx, &inner)
216        }
217        "set_comparison" => {
218            // Set comparisons require set operands (except 'in', which is hadled later)
219            ctx.typechecking_context = TypecheckingContext::Set;
220            parse_binary_expression(ctx, &inner)
221        }
222        "sequence_comparison" => {
223            // Sequence comparisons (substring/subsequence) require sequence operands
224            ctx.typechecking_context = TypecheckingContext::Sequence;
225            parse_binary_expression(ctx, &inner)
226        }
227        "all_diff_comparison" => {
228            ctx.typechecking_context = TypecheckingContext::Matrix;
229            parse_all_diff_comparison(ctx, &inner)
230        }
231        "all_different_except_comparison" => {
232            ctx.typechecking_context = TypecheckingContext::Matrix;
233            parse_all_different_except_comparison(ctx, &inner)
234        }
235        "global_cardinality_comparison" => parse_global_cardinality_comparison(ctx, &inner),
236        _ => {
237            ctx.record_error(RecoverableParseError::new(
238                format!("Expected comparison expression, found '{}'", inner.kind()),
239                Some(inner.range()),
240            ));
241            Ok(None)
242        }
243    }
244}
245
246fn parse_boolean_expression(
247    ctx: &mut ParseContext,
248    node: &Node,
249) -> Result<Option<Expression>, FatalParseError> {
250    ctx.typechecking_context = TypecheckingContext::Boolean;
251    ctx.inner_typechecking_context = TypecheckingContext::Unknown;
252    let Some(inner) = named_child!(recover, ctx, node) else {
253        return Ok(None);
254    };
255    match inner.kind() {
256        "atom" | "table" | "negative_table" => parse_atom(ctx, &inner),
257        "active_expr" => parse_active_expression(ctx, &inner),
258        "not_expr" | "sub_bool_expr" => parse_unary_expression(ctx, &inner),
259        "and_expr" | "or_expr" | "implication" | "iff_expr" => parse_binary_expression(ctx, &inner),
260        "list_combining_expr_bool" => {
261            // list-combining boolean operators accept either set or matrix operands
262            ctx.typechecking_context = TypecheckingContext::SetOrMatrix;
263
264            // set inner context to boolean to ensure elements of list are boolean expressions
265            ctx.inner_typechecking_context = TypecheckingContext::Boolean;
266            parse_list_combining_expression(ctx, &inner)
267        }
268        "quantifier_expr" => parse_quantifier_or_aggregate_expr(ctx, &inner),
269        _ => {
270            ctx.record_error(RecoverableParseError::new(
271                format!("Expected boolean expression, found '{}'", inner.kind()),
272                Some(inner.range()),
273            ));
274            Ok(None)
275        }
276    }
277}
278
279fn parse_active_expression(
280    ctx: &mut ParseContext,
281    node: &Node,
282) -> Result<Option<Expression>, FatalParseError> {
283    let Some(variant_node) = field!(recover, ctx, node, "variant") else {
284        return Ok(None);
285    };
286    let saved_context = ctx.typechecking_context;
287    ctx.typechecking_context = TypecheckingContext::Unknown;
288    let variant = parse_expression(ctx, variant_node)?;
289    ctx.typechecking_context = saved_context;
290    let Some(variant) = variant else {
291        return Ok(None);
292    };
293    let Some(name_node) = field!(recover, ctx, node, "name") else {
294        return Ok(None);
295    };
296    let name = conjure_cp_core::ast::Name::user(
297        &ctx.source_code[name_node.start_byte()..name_node.end_byte()],
298    );
299    let ReturnType::Variant(entries) = variant.return_type() else {
300        ctx.record_error(RecoverableParseError::new(
301            format!("active expects a variant, got {variant}"),
302            Some(node.range()),
303        ));
304        return Ok(None);
305    };
306    if !entries.iter().any(|entry| entry.name == name) {
307        ctx.record_error(RecoverableParseError::new(
308            format!("unknown variant field `{name}`"),
309            Some(name_node.range()),
310        ));
311        return Ok(None);
312    }
313    Ok(Some(Expression::Active(
314        Metadata::new(),
315        Moo::new(variant),
316        name,
317    )))
318}
319
320fn parse_list_combining_expression(
321    ctx: &mut ParseContext,
322    node: &Node,
323) -> Result<Option<Expression>, FatalParseError> {
324    let Some(operator_node) = field!(recover, ctx, node, "operator") else {
325        return Ok(None);
326    };
327    let operator_str = &ctx.source_code[operator_node.start_byte()..operator_node.end_byte()];
328
329    let Some(arg_node) = field!(recover, ctx, node, "arg") else {
330        return Ok(None);
331    };
332    // While parsing inner, the typechecking context is SetOrMatrix
333    // The inner context is either Boolean or Arithmetic so the elements of the set/matrix are typechecked correctly.
334    let Some(inner) = parse_expression(ctx, arg_node)? else {
335        return Ok(None);
336    };
337
338    let skip_operator = list_combining_skip_operator(operator_str);
339    let inner = set_comprehension_skip_operator(inner, skip_operator);
340
341    let expr = match operator_str {
342        "and" => Ok(Some(Expression::And(Metadata::new(), Moo::new(inner)))),
343        "or" => Ok(Some(Expression::Or(Metadata::new(), Moo::new(inner)))),
344        "sum" => Ok(Some(Expression::Sum(Metadata::new(), Moo::new(inner)))),
345        "product" => Ok(Some(Expression::Product(Metadata::new(), Moo::new(inner)))),
346        "min" => Ok(Some(Expression::Min(Metadata::new(), Moo::new(inner)))),
347        "max" => Ok(Some(Expression::Max(Metadata::new(), Moo::new(inner)))),
348        _ => {
349            ctx.record_error(RecoverableParseError::new(
350                format!("Invalid operator: '{operator_str}'"),
351                Some(operator_node.range()),
352            ));
353            Ok(None)
354        }
355    };
356
357    if expr.is_ok() {
358        ctx.add_span_and_doc_hover(
359            &operator_node,
360            operator_str,
361            SymbolKind::Function,
362            None,
363            None,
364        );
365    }
366
367    expr
368}
369
370fn parse_all_diff_comparison(
371    ctx: &mut ParseContext,
372    node: &Node,
373) -> Result<Option<Expression>, FatalParseError> {
374    let Some(arg_node) = field!(recover, ctx, node, "arg") else {
375        return Ok(None);
376    };
377    let Some(inner) = parse_expression(ctx, arg_node)? else {
378        return Ok(None);
379    };
380
381    let operator_node = node
382        .child_by_field_name("operator")
383        .unwrap_or_else(|| node.child(0).unwrap());
384    let operator_str =
385        ctx.source_code[operator_node.start_byte()..operator_node.end_byte()].to_string();
386    if !is_all_different_operator(&operator_str) {
387        ctx.record_error(RecoverableParseError::new(
388            format!("Invalid operator: '{operator_str}'"),
389            Some(operator_node.range()),
390        ));
391        return Ok(None);
392    }
393
394    ctx.add_span_and_doc_hover(
395        &operator_node,
396        ALL_DIFFERENT,
397        SymbolKind::Function,
398        None,
399        None,
400    );
401    Ok(Some(Expression::AllDiff(Metadata::new(), Moo::new(inner))))
402}
403
404fn parse_all_different_except_comparison(
405    ctx: &mut ParseContext,
406    node: &Node,
407) -> Result<Option<Expression>, FatalParseError> {
408    let operator_node = node
409        .child_by_field_name("operator")
410        .unwrap_or_else(|| node.child(0).unwrap());
411    let operator_str =
412        ctx.source_code[operator_node.start_byte()..operator_node.end_byte()].to_string();
413    if !is_all_different_except_operator(&operator_str) {
414        ctx.record_error(RecoverableParseError::new(
415            format!("Invalid operator: '{operator_str}'"),
416            Some(operator_node.range()),
417        ));
418        return Ok(None);
419    }
420
421    let Some(matrix_node) = field!(recover, ctx, node, "matrix") else {
422        return Ok(None);
423    };
424    let Some(except_node) = field!(recover, ctx, node, "except") else {
425        return Ok(None);
426    };
427
428    let saved_context = ctx.typechecking_context;
429    ctx.typechecking_context = TypecheckingContext::Unknown;
430    let Some(matrix) = parse_table_operand(ctx, &matrix_node)? else {
431        ctx.typechecking_context = saved_context;
432        return Ok(None);
433    };
434    let Some(except) = parse_table_operand(ctx, &except_node)? else {
435        ctx.typechecking_context = saved_context;
436        return Ok(None);
437    };
438    ctx.typechecking_context = saved_context;
439
440    ctx.add_span_and_doc_hover(
441        &operator_node,
442        ALL_DIFFERENT_EXCEPT,
443        SymbolKind::Function,
444        None,
445        None,
446    );
447    Ok(Some(Expression::AllDifferentExcept(
448        Metadata::new(),
449        Moo::new(matrix),
450        Moo::new(except),
451    )))
452}
453
454fn parse_global_cardinality_comparison(
455    ctx: &mut ParseContext,
456    node: &Node,
457) -> Result<Option<Expression>, FatalParseError> {
458    let Some(operator_node) = field!(recover, ctx, node, "operator") else {
459        return Ok(None);
460    };
461    let operator_str =
462        ctx.source_code[operator_node.start_byte()..operator_node.end_byte()].to_string();
463
464    let Some(variables_node) = field!(recover, ctx, node, "variables") else {
465        return Ok(None);
466    };
467    let Some(arg2_node) = field!(recover, ctx, node, "arg2") else {
468        return Ok(None);
469    };
470    let Some(arg3_node) = field!(recover, ctx, node, "arg3") else {
471        return Ok(None);
472    };
473
474    let saved_context = ctx.typechecking_context;
475    ctx.typechecking_context = TypecheckingContext::Unknown;
476    let Some(variables) = parse_expression(ctx, variables_node)? else {
477        ctx.typechecking_context = saved_context;
478        return Ok(None);
479    };
480    let Some(arg2) = parse_expression(ctx, arg2_node)? else {
481        ctx.typechecking_context = saved_context;
482        return Ok(None);
483    };
484    let Some(arg3) = parse_expression(ctx, arg3_node)? else {
485        ctx.typechecking_context = saved_context;
486        return Ok(None);
487    };
488    ctx.typechecking_context = saved_context;
489
490    ctx.add_span_and_doc_hover(
491        &operator_node,
492        match operator_str.as_str() {
493            op if is_at_least_operator(op) => AT_LEAST,
494            op if is_at_most_operator(op) => AT_MOST,
495            _ => GLOBAL_CARDINALITY,
496        },
497        SymbolKind::Function,
498        None,
499        None,
500    );
501
502    match operator_str.as_str() {
503        op if is_at_least_operator(op) => Ok(Some(Expression::AtLeast(
504            Metadata::new(),
505            Moo::new(variables),
506            Moo::new(arg2),
507            Moo::new(arg3),
508        ))),
509        op if is_at_most_operator(op) => Ok(Some(Expression::AtMost(
510            Metadata::new(),
511            Moo::new(variables),
512            Moo::new(arg2),
513            Moo::new(arg3),
514        ))),
515        op if is_global_cardinality_operator(op) => Ok(Some(Expression::Gcc(
516            Metadata::new(),
517            Moo::new(variables),
518            Moo::new(arg2),
519            Moo::new(arg3),
520        ))),
521        _ => {
522            ctx.record_error(RecoverableParseError::new(
523                format!("Invalid operator: '{operator_str}'"),
524                Some(operator_node.range()),
525            ));
526            Ok(None)
527        }
528    }
529}
530
531fn parse_unary_expression(
532    ctx: &mut ParseContext,
533    node: &Node,
534) -> Result<Option<Expression>, FatalParseError> {
535    let saved_context = ctx.typechecking_context;
536    if node.kind() == "abs_value" {
537        // Bars are overloaded for numeric absolute value and collection cardinality.
538        ctx.typechecking_context = TypecheckingContext::Unknown;
539    }
540    let Some(expr_node) = field!(recover, ctx, node, "expression") else {
541        ctx.typechecking_context = saved_context;
542        return Ok(None);
543    };
544    let Some(inner) = parse_expression(ctx, expr_node)? else {
545        ctx.typechecking_context = saved_context;
546        return Ok(None);
547    };
548    ctx.typechecking_context = saved_context;
549
550    match node.kind() {
551        "negative_expr" => {
552            if let Expression::Atomic(_, Atom::Literal(Literal::Int(value))) = inner {
553                Ok(Some(Expression::Atomic(
554                    Metadata::new(),
555                    Atom::Literal(Literal::Int(-value)),
556                )))
557            } else {
558                Ok(Some(Expression::Neg(Metadata::new(), Moo::new(inner))))
559            }
560        }
561        "abs_value" => {
562            let constructor = match inner.return_type() {
563                ReturnType::Matrix(_)
564                | ReturnType::Set(_)
565                | ReturnType::MSet(_)
566                | ReturnType::Sequence(_)
567                | ReturnType::Function(_, _)
568                | ReturnType::Relation(_)
569                | ReturnType::Partition(_)
570                | ReturnType::Permutation(_) => Expression::Card,
571                _ => Expression::Abs,
572            };
573            Ok(Some(constructor(Metadata::new(), Moo::new(inner))))
574        }
575        "not_expr" => Ok(Some(Expression::Not(Metadata::new(), Moo::new(inner)))),
576        "toInt_expr" => {
577            let to_int_keyword_node = child!(node, 0, "toInt");
578            ctx.add_span_and_doc_hover(
579                &to_int_keyword_node,
580                "toInt",
581                SymbolKind::Function,
582                None,
583                None,
584            );
585            Ok(Some(Expression::ToInt(Metadata::new(), Moo::new(inner))))
586        }
587        "factorial_expr" => {
588            // looking for the operator node (either '!' at the end or 'factorial' at the start) to add hover info
589            if let Some(op_node) = (0..node.child_count())
590                .filter_map(|i| node.child(i))
591                .find(|c| matches!(c.kind(), "!" | "factorial"))
592            {
593                ctx.add_span_and_doc_hover(
594                    &op_node,
595                    "post_factorial",
596                    SymbolKind::Function,
597                    None,
598                    None,
599                );
600            }
601
602            Ok(Some(Expression::Factorial(
603                Metadata::new(),
604                Moo::new(inner),
605            )))
606        }
607        "sub_bool_expr" | "sub_arith_expr" => Ok(Some(inner)),
608        _ => {
609            ctx.record_error(RecoverableParseError::new(
610                format!("Unrecognised unary operation: '{}'", node.kind()),
611                Some(node.range()),
612            ));
613            Ok(None)
614        }
615    }
616}
617
618pub fn parse_binary_expression(
619    ctx: &mut ParseContext,
620    node: &Node,
621) -> Result<Option<Expression>, FatalParseError> {
622    let Some(op_node) = field!(recover, ctx, node, "operator") else {
623        return Ok(None);
624    };
625    let op_str = &ctx.source_code[op_node.start_byte()..op_node.end_byte()];
626
627    let saved_ctx = ctx.typechecking_context;
628
629    // Special handling for 'in' operator, as the left operand doesn't have to be a set
630    if op_str == "in" {
631        ctx.typechecking_context = TypecheckingContext::Unknown
632    }
633
634    // Minus spells both arithmetic subtraction and set difference, so neither operand can be held
635    // to the arithmetic context until the operand types say which one this is.
636    if op_str == "-" {
637        ctx.typechecking_context = TypecheckingContext::Unknown;
638    }
639
640    // parse left operand
641    let Some(left_node) = field!(recover, ctx, node, "left") else {
642        return Ok(None);
643    };
644    let Some(left) = parse_expression(ctx, left_node)? else {
645        return Ok(None);
646    };
647
648    // reset context, if needed
649    ctx.typechecking_context = saved_ctx;
650
651    // Equality/inequality/ordering: enforce right operand to match left operand type when
652    // inferable. Ordering (<, <=, >, >=) works on any orderable type now, the same as equality,
653    // so it needs the same operand-matching check.
654    if matches!(op_str, "=" | "!=" | "<" | "<=" | ">" | ">=") {
655        ctx.typechecking_context = inferred_context_from_expression(&left);
656    }
657
658    // 'in's right operand is the container, which need not be a set (mset/sequence-via-toSet/
659    // relation/etc all support membership too); the grammar's "set_comparison" node kind sets
660    // `saved_ctx` to `Set` unconditionally, which only the left operand's override above escapes.
661    if op_str == "in" {
662        ctx.typechecking_context = TypecheckingContext::Unknown;
663    }
664
665    if op_str == "-" {
666        ctx.typechecking_context = TypecheckingContext::Unknown;
667    }
668
669    // parse right operand
670    let Some(right_node) = field!(recover, ctx, node, "right") else {
671        return Ok(None);
672    };
673    let Some(right) = parse_expression(ctx, right_node)? else {
674        return Ok(None);
675    };
676
677    // restore original contexts for parent expression parsing
678    ctx.typechecking_context = saved_ctx;
679
680    let mut doc_name = "";
681    let expr = match op_str {
682        // NB: We are deliberately setting the index domain to 1.., not 1..2.
683        // Semantically, this means "a list that can grow/shrink arbitrarily".
684        // This is expected by rules which will modify the terms of the sum expression
685        // (e.g. by partially evaluating them).
686        "+" => {
687            doc_name = "L_Plus";
688            Ok(Some(Expression::Sum(
689                Metadata::new(),
690                Moo::new(matrix_expr![left, right; domain_int!(1..)]),
691            )))
692        }
693        "-" => {
694            doc_name = "L_Minus";
695            if is_set_valued(&left) || is_set_valued(&right) {
696                Ok(Some(Expression::Difference(
697                    Metadata::new(),
698                    Moo::new(left),
699                    Moo::new(right),
700                )))
701            } else {
702                // Subtraction, so the operands do have to be integers. The arithmetic context
703                // could not be imposed while parsing them, since it would have rejected the set
704                // operands that make this a difference instead.
705                if !require_int_operand(ctx, &left_node, &left)
706                    || !require_int_operand(ctx, &right_node, &right)
707                {
708                    return Ok(None);
709                }
710                Ok(Some(Expression::Minus(
711                    Metadata::new(),
712                    Moo::new(left),
713                    Moo::new(right),
714                )))
715            }
716        }
717        "*" => {
718            doc_name = "L_Times";
719            Ok(Some(Expression::Product(
720                Metadata::new(),
721                Moo::new(matrix_expr![left, right; domain_int!(1..)]),
722            )))
723        }
724        "/\\" => {
725            doc_name = "and";
726            Ok(Some(Expression::And(
727                Metadata::new(),
728                Moo::new(matrix_expr![left, right; domain_int!(1..)]),
729            )))
730        }
731        "\\/" => {
732            // No documentation for or in Bits yet
733            doc_name = "or";
734            Ok(Some(Expression::Or(
735                Metadata::new(),
736                Moo::new(matrix_expr![left, right; domain_int!(1..)]),
737            )))
738        }
739        "**" => {
740            doc_name = "L_Pow";
741            Ok(Some(Expression::UnsafePow(
742                Metadata::new(),
743                Moo::new(left),
744                Moo::new(right),
745            )))
746        }
747        "/" => {
748            //TODO: add checks for if division is safe or not
749            doc_name = "L_Div";
750            Ok(Some(Expression::UnsafeDiv(
751                Metadata::new(),
752                Moo::new(left),
753                Moo::new(right),
754            )))
755        }
756        "%" => {
757            //TODO: add checks for if mod is safe or not
758            doc_name = "L_Mod";
759            Ok(Some(Expression::UnsafeMod(
760                Metadata::new(),
761                Moo::new(left),
762                Moo::new(right),
763            )))
764        }
765
766        "=" => {
767            doc_name = "L_Eq"; //no docs yet
768            Ok(Some(Expression::Eq(
769                Metadata::new(),
770                Moo::new(left),
771                Moo::new(right),
772            )))
773        }
774        "!=" => {
775            doc_name = "L_Neq"; //no docs yet
776            Ok(Some(Expression::Neq(
777                Metadata::new(),
778                Moo::new(left),
779                Moo::new(right),
780            )))
781        }
782        "<=" => {
783            doc_name = "L_Leq"; //no docs yet
784            Ok(Some(Expression::Leq(
785                Metadata::new(),
786                Moo::new(left),
787                Moo::new(right),
788            )))
789        }
790        ">=" => {
791            doc_name = "L_Geq"; //no docs yet
792            Ok(Some(Expression::Geq(
793                Metadata::new(),
794                Moo::new(left),
795                Moo::new(right),
796            )))
797        }
798        "<" => {
799            doc_name = "L_Lt"; //no docs yet
800            Ok(Some(Expression::Lt(
801                Metadata::new(),
802                Moo::new(left),
803                Moo::new(right),
804            )))
805        }
806        ">" => {
807            doc_name = "L_Gt"; //no docs yet
808            Ok(Some(Expression::Gt(
809                Metadata::new(),
810                Moo::new(left),
811                Moo::new(right),
812            )))
813        }
814
815        "->" => {
816            doc_name = "L_Imply"; //no docs yet
817            Ok(Some(Expression::Imply(
818                Metadata::new(),
819                Moo::new(left),
820                Moo::new(right),
821            )))
822        }
823        "<->" => {
824            doc_name = "L_Iff"; //no docs yet
825            Ok(Some(Expression::Iff(
826                Metadata::new(),
827                Moo::new(left),
828                Moo::new(right),
829            )))
830        }
831        "<lex" => {
832            doc_name = "L_LexLt"; //no docs yet
833            Ok(Some(Expression::LexLt(
834                Metadata::new(),
835                Moo::new(left),
836                Moo::new(right),
837            )))
838        }
839        ">lex" => {
840            doc_name = "L_LexGt"; //no docs yet
841            Ok(Some(Expression::LexGt(
842                Metadata::new(),
843                Moo::new(left),
844                Moo::new(right),
845            )))
846        }
847        "<=lex" => {
848            doc_name = "L_LexLeq"; //no docs yet
849            Ok(Some(Expression::LexLeq(
850                Metadata::new(),
851                Moo::new(left),
852                Moo::new(right),
853            )))
854        }
855        ">=lex" => {
856            doc_name = "L_LexGeq"; //no docs yet
857            Ok(Some(Expression::LexGeq(
858                Metadata::new(),
859                Moo::new(left),
860                Moo::new(right),
861            )))
862        }
863        "in" => {
864            doc_name = "L_in";
865            Ok(Some(Expression::In(
866                Metadata::new(),
867                Moo::new(left),
868                Moo::new(right),
869            )))
870        }
871        "subset" => {
872            doc_name = "L_subset";
873            Ok(Some(Expression::Subset(
874                Metadata::new(),
875                Moo::new(left),
876                Moo::new(right),
877            )))
878        }
879        "subsetEq" => {
880            doc_name = "L_subsetEq";
881            Ok(Some(Expression::SubsetEq(
882                Metadata::new(),
883                Moo::new(left),
884                Moo::new(right),
885            )))
886        }
887        "supset" => {
888            doc_name = "L_supset";
889            Ok(Some(Expression::Supset(
890                Metadata::new(),
891                Moo::new(left),
892                Moo::new(right),
893            )))
894        }
895        "supsetEq" => {
896            doc_name = "L_supsetEq";
897            Ok(Some(Expression::SupsetEq(
898                Metadata::new(),
899                Moo::new(left),
900                Moo::new(right),
901            )))
902        }
903        "substring" => {
904            doc_name = "L_substring";
905            Ok(Some(Expression::Substring(
906                Metadata::new(),
907                Moo::new(left),
908                Moo::new(right),
909            )))
910        }
911        "subsequence" => {
912            doc_name = "L_subsequence";
913            Ok(Some(Expression::Subsequence(
914                Metadata::new(),
915                Moo::new(left),
916                Moo::new(right),
917            )))
918        }
919        "union" => {
920            doc_name = "L_union";
921            Ok(Some(Expression::Union(
922                Metadata::new(),
923                Moo::new(left),
924                Moo::new(right),
925            )))
926        }
927        "intersect" => {
928            doc_name = "L_intersect";
929            Ok(Some(Expression::Intersect(
930                Metadata::new(),
931                Moo::new(left),
932                Moo::new(right),
933            )))
934        }
935        _ => {
936            ctx.record_error(RecoverableParseError::new(
937                format!("Invalid operator: '{op_str}'"),
938                Some(op_node.range()),
939            ));
940            Ok(None)
941        }
942    };
943
944    if expr.is_ok() {
945        ctx.add_span_and_doc_hover(&op_node, doc_name, SymbolKind::Function, None, None);
946    }
947
948    expr
949}
950
951pub fn parse_annotation_expression(
952    ctx: &mut ParseContext,
953    node: &Node,
954) -> Result<Option<Expression>, FatalParseError> {
955    let annotation_node = if node.kind() == "annotation_expr" {
956        let Some(inner) = named_child!(recover, ctx, node) else {
957            return Ok(None);
958        };
959        inner
960    } else {
961        *node
962    };
963
964    let Some(left_node) = field!(recover, ctx, &annotation_node, "left") else {
965        return Ok(None);
966    };
967    let left_node = if left_node.kind() == "annotation_subject" {
968        let Some(inner) = named_child!(recover, ctx, &left_node) else {
969            return Ok(None);
970        };
971        inner
972    } else {
973        left_node
974    };
975    let left_node = match left_node.kind() {
976        "sub_arith_expr" | "sub_bool_expr" | "sub_atom_expr" => {
977            let Some(inner) = field!(recover, ctx, &left_node, "expression") else {
978                return Ok(None);
979            };
980            inner
981        }
982        _ => left_node,
983    };
984    let saved_context = ctx.typechecking_context;
985    ctx.typechecking_context = TypecheckingContext::Unknown;
986    let left = match left_node.kind() {
987        "constant" | "identifier" | "metavar" => parse_atom(ctx, &left_node)?,
988        _ => parse_expression(ctx, left_node)?,
989    };
990    let Some(left) = left else {
991        ctx.typechecking_context = saved_context;
992        return Ok(None);
993    };
994    ctx.typechecking_context = saved_context;
995
996    let Some(operator_node) = field!(recover, ctx, &annotation_node, "operator") else {
997        return Ok(None);
998    };
999    let operator_str = &ctx.source_code[operator_node.start_byte()..operator_node.end_byte()];
1000
1001    match annotation_node.kind() {
1002        "type_annotation" => {
1003            let Some(type_node) = field!(recover, ctx, &annotation_node, "type") else {
1004                return Ok(None);
1005            };
1006            let Some(domain) = parse_domain(ctx, type_node)? else {
1007                return Ok(None);
1008            };
1009            ctx.add_span_and_doc_hover(
1010                &operator_node,
1011                operator_str,
1012                SymbolKind::Function,
1013                None,
1014                None,
1015            );
1016            Ok(Some(Expression::TypeAnnotation(
1017                Metadata::new(),
1018                Moo::new(left),
1019                domain,
1020            )))
1021        }
1022        "domain_annotation" => {
1023            let Some(domain_node) = field!(recover, ctx, &annotation_node, "domain") else {
1024                return Ok(None);
1025            };
1026            let Some(domain) = parse_domain(ctx, domain_node)? else {
1027                return Ok(None);
1028            };
1029            ctx.add_span_and_doc_hover(
1030                &operator_node,
1031                operator_str,
1032                SymbolKind::Function,
1033                None,
1034                None,
1035            );
1036            Ok(Some(Expression::DomainAnnotation(
1037                Metadata::new(),
1038                Moo::new(left),
1039                domain,
1040            )))
1041        }
1042        _ => {
1043            ctx.record_error(RecoverableParseError::new(
1044                format!(
1045                    "Unrecognised annotation expression: '{}'",
1046                    annotation_node.kind()
1047                ),
1048                Some(annotation_node.range()),
1049            ));
1050            Ok(None)
1051        }
1052    }
1053}
1054
1055fn inferred_context_from_expression(expr: &Expression) -> TypecheckingContext {
1056    // TODO: typechecking for index/slice expressions
1057    if matches!(
1058        expr,
1059        Expression::UnsafeIndex(_, _, _) | Expression::UnsafeSlice(_, _, _)
1060    ) {
1061        return TypecheckingContext::Unknown;
1062    }
1063
1064    let Some(domain) = expr.domain_of() else {
1065        return TypecheckingContext::Unknown;
1066    };
1067    let Ok(ground) = domain.resolve() else {
1068        return TypecheckingContext::Unknown;
1069    };
1070
1071    match ground.as_ref() {
1072        GroundDomain::Empty(_) => TypecheckingContext::Unknown,
1073        GroundDomain::Bool => TypecheckingContext::Boolean,
1074        GroundDomain::Int(_) => TypecheckingContext::Arithmetic,
1075        GroundDomain::Tuple(_) => TypecheckingContext::Tuple,
1076        GroundDomain::Record(_) => TypecheckingContext::Record,
1077        GroundDomain::Variant(_) => TypecheckingContext::Unknown,
1078        GroundDomain::Matrix(_, _) => TypecheckingContext::Matrix,
1079        GroundDomain::Sequence(_, _) => TypecheckingContext::Sequence,
1080        GroundDomain::Set(_, _) => TypecheckingContext::Set,
1081        GroundDomain::MSet(_, _) => TypecheckingContext::MSet,
1082        GroundDomain::Function(_, _, _) => TypecheckingContext::Function,
1083        GroundDomain::Relation(_, _) => TypecheckingContext::Relation,
1084        GroundDomain::Partition(_, _) => TypecheckingContext::Partition,
1085        GroundDomain::Permutation(_, _) => TypecheckingContext::Permutation,
1086    }
1087}
1088
1089fn list_combining_skip_operator(operator_str: &str) -> Option<ACOperatorKind> {
1090    match operator_str {
1091        "and" => Some(ACOperatorKind::And),
1092        "or" => Some(ACOperatorKind::Or),
1093        "sum" => Some(ACOperatorKind::Sum),
1094        "product" => Some(ACOperatorKind::Product),
1095        // min/max are not true AC operators (no universal identity element), but still need
1096        // their own skip-operator tag so a symbolic guard inside `min([... | ...])` lowers
1097        // correctly (previously mapped to Sum, which silently substituted 0 for guarded-out
1098        // elements -- wrong for anything but a min/max whose correct answer happens to be <=/>= 0).
1099        "min" => Some(ACOperatorKind::Min),
1100        "max" => Some(ACOperatorKind::Max),
1101        _ => None,
1102    }
1103}
1104
1105fn set_comprehension_skip_operator(
1106    inner: Expression,
1107    skip_operator: Option<ACOperatorKind>,
1108) -> Expression {
1109    let Expression::Comprehension(meta, comprehension) = inner else {
1110        return inner;
1111    };
1112    if let Some(skip_operator) = skip_operator {
1113        let mut comprehension = Moo::unwrap_or_clone(comprehension);
1114        comprehension.skip_operator = Some(skip_operator);
1115        Expression::Comprehension(meta, Moo::new(comprehension))
1116    } else {
1117        Expression::Comprehension(meta, comprehension)
1118    }
1119}
1120
1121/// Parses `catchUndef(expression, default)`.
1122fn parse_catch_undef_expression(
1123    ctx: &mut ParseContext,
1124    node: &Node,
1125) -> Result<Option<Expression>, FatalParseError> {
1126    ctx.typechecking_context = TypecheckingContext::Arithmetic;
1127
1128    let Some(expression_node) = field!(recover, ctx, node, "expression") else {
1129        return Ok(None);
1130    };
1131    let Some(default_node) = field!(recover, ctx, node, "default") else {
1132        return Ok(None);
1133    };
1134
1135    let Some(expression) = parse_expression(ctx, expression_node)? else {
1136        return Ok(None);
1137    };
1138    ctx.typechecking_context = TypecheckingContext::Arithmetic;
1139    let Some(default) = parse_expression(ctx, default_node)? else {
1140        return Ok(None);
1141    };
1142
1143    Ok(Some(Expression::CatchUndef(
1144        Metadata::new(),
1145        Moo::new(expression),
1146        Moo::new(default),
1147    )))
1148}
1149
1150/// Whether `expr` is known to be set-valued, used to read `-` as set difference.
1151fn is_set_valued(expr: &Expression) -> bool {
1152    matches!(expr.try_return_type(), Some(ReturnType::Set(_)))
1153}
1154
1155/// Reports a type error when an arithmetic operand is not an integer.
1156///
1157/// Returns whether the operand is acceptable.
1158fn require_int_operand(ctx: &mut ParseContext, node: &Node, expr: &Expression) -> bool {
1159    // An operand whose type cannot be worked out yet gets the benefit of the doubt.
1160    let Some(actual) = expr.try_return_type() else {
1161        return true;
1162    };
1163    if matches!(actual, ReturnType::Int | ReturnType::Unknown) {
1164        return true;
1165    }
1166
1167    ctx.record_error(RecoverableParseError::new(
1168        format!(
1169            "Type error: {}\n\tExpected: int\n\tGot: {}",
1170            ctx.source_code[node.start_byte()..node.end_byte()].trim(),
1171            type_name(&actual)
1172        ),
1173        Some(node.range()),
1174    ));
1175    false
1176}
1177
1178/// The bare name of a type, as type errors spell it: `partition`, not `partition of int`.
1179fn type_name(return_type: &ReturnType) -> &'static str {
1180    match return_type {
1181        ReturnType::Unknown => "unknown",
1182        ReturnType::Bool => "bool",
1183        ReturnType::Int => "int",
1184        ReturnType::Tuple(_) => "tuple",
1185        ReturnType::Record(_) => "record",
1186        ReturnType::Variant(_) => "variant",
1187        ReturnType::Matrix(_) => "matrix",
1188        ReturnType::Sequence(_) => "sequence",
1189        ReturnType::Set(_) => "set",
1190        ReturnType::MSet(_) => "mset",
1191        ReturnType::Function(_, _) => "function",
1192        ReturnType::Relation(_) => "relation",
1193        ReturnType::Partition(_) => "partition",
1194        ReturnType::Permutation(_) => "permutation",
1195    }
1196}