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
5209
pub fn parse_abstract(
12
5209
    ctx: &mut ParseContext,
13
5209
    node: &Node,
14
5209
) -> Result<Option<AbstractLiteral<Expression>>, FatalParseError> {
15
5209
    if typecheck_abstract_literal(ctx, node) {
16
26
        return Ok(None);
17
5183
    }
18

            
19
5183
    match node.kind() {
20
5183
        "record" => parse_record(ctx, node),
21
5157
        "tuple" => parse_tuple(ctx, node),
22
5053
        "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
5209
}
33

            
34
5209
fn typecheck_abstract_literal(ctx: &mut ParseContext, node: &Node) -> bool {
35
5209
    let expected = match ctx.typechecking_context {
36
13
        TypecheckingContext::Boolean => "bool",
37
        TypecheckingContext::Arithmetic => "int",
38
832
        TypecheckingContext::Set => "set",
39
1812
        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
2266
        TypecheckingContext::Unknown => "unknown",
47
    };
48

            
49
5209
    let got = match node.kind() {
50
5209
        "set_literal" => "set",
51
3727
        "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
5209
    if expected != "unknown"
64
2943
        && !(ctx.typechecking_context == TypecheckingContext::SetOrMatrix
65
1812
            && 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
5183
    }
79

            
80
5183
    false
81
5209
}
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
26
    let mut had_error = false;
89
52
    for child in node.children_by_field_name("name_value_pair", &mut node.walk()) {
90
52
        let Some(name_node) = field!(recover, ctx, child, "name") else {
91
            return Ok(None);
92
        };
93
52
        let name_str = &ctx.source_code[name_node.start_byte()..name_node.end_byte()];
94
52
        let name = conjure_cp_core::ast::Name::user(name_str);
95

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

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

            
105
52
        let Some(value) = parse_expression(ctx, value_node)? else {
106
            had_error = true;
107
            ctx.typechecking_context = saved_ctx;
108
            ctx.inner_typechecking_context = TypecheckingContext::Unknown;
109
            continue; // continue parsing other elements
110
        };
111

            
112
        // Reset contexts
113
52
        ctx.inner_typechecking_context = ctx.typechecking_context;
114
52
        ctx.typechecking_context = saved_ctx;
115

            
116
52
        values.push(conjure_cp_core::ast::records::Field { name, value });
117
    }
118

            
119
26
    if had_error {
120
        Ok(None)
121
    } else {
122
26
        Ok(Some(AbstractLiteral::Record(values)))
123
    }
124
26
}
125

            
126
104
fn parse_tuple(
127
104
    ctx: &mut ParseContext,
128
104
    node: &Node,
129
104
) -> Result<Option<AbstractLiteral<Expression>>, FatalParseError> {
130
    // Save the typechecking context
131
104
    let saved_ctx = ctx.typechecking_context;
132
104
    let saved_inner_ctx = ctx.inner_typechecking_context;
133

            
134
104
    let mut elements = Vec::new();
135
104
    let mut had_error = false;
136
260
    for child in named_children(node) {
137
        // Parse elements with inner typechecking context
138
260
        ctx.typechecking_context = saved_inner_ctx;
139
260
        ctx.inner_typechecking_context = TypecheckingContext::Unknown;
140

            
141
260
        let Some(expr) = parse_expression(ctx, child)? else {
142
            had_error = true;
143
            continue;
144
        };
145
260
        elements.push(expr);
146
    }
147

            
148
104
    ctx.typechecking_context = saved_ctx;
149
104
    ctx.inner_typechecking_context = saved_inner_ctx;
150
104
    if had_error {
151
        Ok(None)
152
    } else {
153
104
        Ok(Some(AbstractLiteral::Tuple(elements)))
154
    }
155
104
}
156

            
157
3571
fn parse_matrix(
158
3571
    ctx: &mut ParseContext,
159
3571
    node: &Node,
160
3571
) -> Result<Option<AbstractLiteral<Expression>>, FatalParseError> {
161
    // Save the typechecking contexts
162
3571
    let saved_ctx = ctx.typechecking_context;
163
3571
    let saved_inner_ctx = ctx.inner_typechecking_context;
164

            
165
3571
    let mut elements = vec![];
166
3571
    let mut domain: Option<DomainPtr> = None;
167
3571
    let mut had_error = false;
168
10040
    for child in named_children(node) {
169
10040
        if child.kind() == "arithmetic_expr"
170
9702
            || child.kind() == "bool_expr"
171
9682
            || child.kind() == "comparison_expr"
172
9232
            || child.kind() == "atom"
173
        {
174
            // Parse elements with inner typechecking context
175
9524
            ctx.typechecking_context = saved_inner_ctx;
176
9524
            ctx.inner_typechecking_context = TypecheckingContext::Unknown;
177

            
178
9524
            let Some(expr) = parse_expression(ctx, child)? else {
179
117
                had_error = true;
180
117
                continue;
181
            };
182
9407
            elements.push(expr);
183
        } else {
184
            // Parse domains with unknown typechecking context
185
516
            ctx.typechecking_context = TypecheckingContext::Unknown;
186

            
187
516
            let Some(parsed_domain) = parse_domain(ctx, child)? else {
188
13
                had_error = true;
189
13
                continue;
190
            };
191
503
            domain = Some(parsed_domain);
192
        }
193
    }
194
3571
    if domain.is_none() {
195
3068
        let count = elements.len() as i32;
196
3068
        domain = Some(domain_int!(1..count));
197
3428
    }
198

            
199
3571
    ctx.typechecking_context = saved_ctx;
200
3571
    ctx.inner_typechecking_context = saved_inner_ctx;
201
3571
    if had_error {
202
91
        Ok(None)
203
    } else {
204
3480
        Ok(Some(AbstractLiteral::Matrix(elements, domain.unwrap())))
205
    }
206
3571
}
207

            
208
1482
fn parse_set_literal(
209
1482
    ctx: &mut ParseContext,
210
1482
    node: &Node,
211
1482
) -> Result<Option<AbstractLiteral<Expression>>, FatalParseError> {
212
    // Save the typechecking contexts
213
1482
    let saved_ctx = ctx.typechecking_context;
214
1482
    let saved_inner_ctx = ctx.inner_typechecking_context;
215

            
216
1482
    let mut elements = Vec::new();
217
1482
    let mut had_error = false;
218
3848
    for child in named_children(node) {
219
        // Parse elements with inner typechecking context
220
3848
        ctx.typechecking_context = saved_inner_ctx;
221
3848
        ctx.inner_typechecking_context = TypecheckingContext::Unknown;
222

            
223
3848
        let Some(expr) = parse_expression(ctx, child)? else {
224
26
            had_error = true;
225
26
            continue;
226
        };
227
3822
        elements.push(expr);
228
    }
229

            
230
1482
    ctx.typechecking_context = saved_ctx;
231
1482
    ctx.inner_typechecking_context = saved_inner_ctx;
232
1482
    if had_error {
233
26
        Ok(None)
234
    } else {
235
1456
        Ok(Some(AbstractLiteral::Set(elements)))
236
    }
237
1482
}