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