1
use crate::errors::{FatalParseError, RecoverableParseError};
2
use crate::expression::parse_expression;
3
use crate::field;
4
use crate::parser::ParseContext;
5
use crate::parser::domain::parse_domain;
6
use crate::util::{TypecheckingContext, named_children};
7
use conjure_cp_core::ast::{AbstractLiteral, DomainPtr, Expression};
8
use conjure_cp_core::{domain_int, range};
9
use tree_sitter::Node;
10

            
11
5018
pub fn parse_abstract(
12
5018
    ctx: &mut ParseContext,
13
5018
    node: &Node,
14
5018
) -> Result<Option<AbstractLiteral<Expression>>, FatalParseError> {
15
5018
    if typecheck_abstract_literal(ctx, node) {
16
26
        return Ok(None);
17
4992
    }
18

            
19
4992
    match node.kind() {
20
4992
        "record" => parse_record(ctx, node),
21
4966
        "tuple" => parse_tuple(ctx, node),
22
4862
        "matrix" => parse_matrix(ctx, node),
23
1482
        "set_literal" => parse_set_literal(ctx, node),
24
        _ => {
25
            ctx.record_error(RecoverableParseError::new(
26
                format!("Expected abstract literal, got: {}", node.kind()),
27
                Some(node.range()),
28
            ));
29
            Ok(None)
30
        }
31
    }
32
5018
}
33

            
34
5018
fn typecheck_abstract_literal(ctx: &mut ParseContext, node: &Node) -> bool {
35
5018
    let expected = match ctx.typechecking_context {
36
13
        TypecheckingContext::Boolean => "bool",
37
        TypecheckingContext::Arithmetic => "int",
38
832
        TypecheckingContext::Set => "set",
39
1751
        TypecheckingContext::SetOrMatrix => "set or matrix",
40
        TypecheckingContext::MSet => "mset",
41
208
        TypecheckingContext::Matrix => "matrix",
42
52
        TypecheckingContext::Tuple => "tuple",
43
26
        TypecheckingContext::Record => "record",
44
        TypecheckingContext::Partition => "partition",
45
        TypecheckingContext::Sequence => "sequence",
46
2136
        TypecheckingContext::Unknown => "unknown",
47
    };
48

            
49
5018
    let got = match node.kind() {
50
5018
        "set_literal" => "set",
51
3536
        "matrix" => "matrix",
52
156
        "tuple" => "tuple",
53
26
        "record" => "record",
54
        _ => {
55
            ctx.record_error(RecoverableParseError::new(
56
                format!("Expected abstract literal, got: {}", node.kind()),
57
                Some(node.range()),
58
            ));
59
            return true;
60
        }
61
    };
62

            
63
5018
    if expected != "unknown"
64
2882
        && !(ctx.typechecking_context == TypecheckingContext::SetOrMatrix
65
1751
            && matches!(got, "set" | "matrix"))
66
1144
        && expected != got
67
    {
68
26
        ctx.record_error(RecoverableParseError::new(
69
26
            format!(
70
                "Type error: {}\n\tExpected: {}\n\tGot: {}",
71
26
                ctx.source_code[node.start_byte()..node.end_byte()].trim(),
72
                expected,
73
                got
74
            ),
75
26
            Some(node.range()),
76
        ));
77
26
        return true;
78
4992
    }
79

            
80
4992
    false
81
5018
}
82

            
83
26
fn parse_record(
84
26
    ctx: &mut ParseContext,
85
26
    node: &Node,
86
26
) -> Result<Option<AbstractLiteral<Expression>>, FatalParseError> {
87
26
    let mut values = Vec::new();
88
52
    for child in node.children_by_field_name("name_value_pair", &mut node.walk()) {
89
52
        let Some(name_node) = field!(recover, ctx, child, "name") else {
90
            return Ok(None);
91
        };
92
52
        let name_str = &ctx.source_code[name_node.start_byte()..name_node.end_byte()];
93
52
        let name = conjure_cp_core::ast::Name::user(name_str);
94

            
95
52
        let Some(value_node) = field!(recover, ctx, child, "value") else {
96
            return Ok(None);
97
        };
98

            
99
        // Parse value with inner typechecking context
100
52
        let saved_ctx = ctx.typechecking_context;
101
52
        ctx.typechecking_context = ctx.inner_typechecking_context;
102
52
        ctx.inner_typechecking_context = TypecheckingContext::Unknown;
103

            
104
52
        let Some(value) = parse_expression(ctx, value_node)? else {
105
            return Ok(None);
106
        };
107

            
108
        // Reset contexts
109
52
        ctx.inner_typechecking_context = ctx.typechecking_context;
110
52
        ctx.typechecking_context = saved_ctx;
111

            
112
52
        values.push(conjure_cp_core::ast::records::FieldValue { name, value });
113
    }
114
26
    Ok(Some(AbstractLiteral::Record(values)))
115
26
}
116

            
117
104
fn parse_tuple(
118
104
    ctx: &mut ParseContext,
119
104
    node: &Node,
120
104
) -> Result<Option<AbstractLiteral<Expression>>, FatalParseError> {
121
    // Save the typechecking context
122
104
    let saved_ctx = ctx.typechecking_context;
123
104
    let saved_inner_ctx = ctx.inner_typechecking_context;
124

            
125
104
    let mut elements = Vec::new();
126
260
    for child in named_children(node) {
127
        // Parse elements with inner typechecking context
128
260
        ctx.typechecking_context = saved_inner_ctx;
129
260
        ctx.inner_typechecking_context = TypecheckingContext::Unknown;
130

            
131
260
        let Some(expr) = parse_expression(ctx, child)? else {
132
            ctx.typechecking_context = saved_ctx;
133
            ctx.inner_typechecking_context = saved_inner_ctx;
134
            return Ok(None);
135
        };
136
260
        elements.push(expr);
137
    }
138

            
139
104
    ctx.typechecking_context = saved_ctx;
140
104
    ctx.inner_typechecking_context = saved_inner_ctx;
141
104
    Ok(Some(AbstractLiteral::Tuple(elements)))
142
104
}
143

            
144
3380
fn parse_matrix(
145
3380
    ctx: &mut ParseContext,
146
3380
    node: &Node,
147
3380
) -> Result<Option<AbstractLiteral<Expression>>, FatalParseError> {
148
    // Save the typechecking contexts
149
3380
    let saved_ctx = ctx.typechecking_context;
150
3380
    let saved_inner_ctx = ctx.inner_typechecking_context;
151

            
152
3380
    let mut elements = vec![];
153
3380
    let mut domain: Option<DomainPtr> = None;
154
8967
    for child in named_children(node) {
155
8967
        if child.kind() == "arithmetic_expr"
156
8642
            || child.kind() == "bool_expr"
157
8632
            || child.kind() == "comparison_expr"
158
8290
            || child.kind() == "atom"
159
        {
160
            // Parse elements with inner typechecking context
161
8629
            ctx.typechecking_context = saved_inner_ctx;
162
8629
            ctx.inner_typechecking_context = TypecheckingContext::Unknown;
163

            
164
8629
            let Some(expr) = parse_expression(ctx, child)? else {
165
65
                ctx.typechecking_context = saved_ctx;
166
65
                ctx.inner_typechecking_context = saved_inner_ctx;
167
65
                return Ok(None);
168
            };
169
8564
            elements.push(expr);
170
        } else {
171
            // Parse domains with unknown typechecking context
172
338
            ctx.typechecking_context = TypecheckingContext::Unknown;
173

            
174
338
            let Some(parsed_domain) = parse_domain(ctx, child)? else {
175
13
                ctx.typechecking_context = saved_ctx;
176
13
                ctx.inner_typechecking_context = saved_inner_ctx;
177
13
                return Ok(None);
178
            };
179
325
            domain = Some(parsed_domain);
180
        }
181
    }
182
3302
    if domain.is_none() {
183
2977
        let count = elements.len() as i32;
184
2977
        domain = Some(domain_int!(1..count));
185
3159
    }
186

            
187
3302
    ctx.typechecking_context = saved_ctx;
188
3302
    ctx.inner_typechecking_context = saved_inner_ctx;
189
3302
    Ok(Some(AbstractLiteral::Matrix(elements, domain.unwrap())))
190
3380
}
191

            
192
1482
fn parse_set_literal(
193
1482
    ctx: &mut ParseContext,
194
1482
    node: &Node,
195
1482
) -> Result<Option<AbstractLiteral<Expression>>, FatalParseError> {
196
    // Save the typechecking contexts
197
1482
    let saved_ctx = ctx.typechecking_context;
198
1482
    let saved_inner_ctx = ctx.inner_typechecking_context;
199

            
200
1482
    let mut elements = Vec::new();
201
3848
    for child in named_children(node) {
202
        // Parse elements with inner typechecking context
203
3848
        ctx.typechecking_context = saved_inner_ctx;
204
3848
        ctx.inner_typechecking_context = TypecheckingContext::Unknown;
205

            
206
3848
        let Some(expr) = parse_expression(ctx, child)? else {
207
26
            ctx.typechecking_context = saved_ctx;
208
26
            ctx.inner_typechecking_context = saved_inner_ctx;
209
26
            return Ok(None);
210
        };
211
3822
        elements.push(expr);
212
    }
213

            
214
1456
    ctx.typechecking_context = saved_ctx;
215
1456
    ctx.inner_typechecking_context = saved_inner_ctx;
216
1456
    Ok(Some(AbstractLiteral::Set(elements)))
217
1482
}