1
use super::atom::parse_int;
2
use super::util::named_children;
3
use crate::diagnostics::diagnostics_api::SymbolKind;
4
use crate::diagnostics::source_map::{HoverInfo, span_with_hover};
5
use crate::errors::FatalParseError;
6
use crate::expression::parse_expression;
7
use crate::parser::ParseContext;
8
use crate::{RecoverableParseError, child};
9
use conjure_cp_core::ast::{
10
    DeclarationPtr, Domain, DomainPtr, Field, IntVal, Moo, Name, Range, Reference, SetAttr,
11
};
12
use tree_sitter::Node;
13

            
14
use crate::field;
15

            
16
/// Parse an Essence variable domain into its Conjure AST representation.
17
98100
pub fn parse_domain(
18
98100
    ctx: &mut ParseContext,
19
98100
    domain: Node,
20
98100
) -> Result<Option<DomainPtr>, FatalParseError> {
21
98100
    match domain.kind() {
22
98100
        "domain" => {
23
48792
            let inner = match domain.child(0) {
24
48792
                Some(node) => node,
25
                None => {
26
                    ctx.record_error(RecoverableParseError::new(
27
                        format!("{} in expression of kind '{}'", "domain", domain.kind()),
28
                        Some(domain.range()),
29
                    ));
30
                    return Ok(None);
31
                }
32
            };
33
48792
            parse_domain(ctx, inner)
34
        }
35
49308
        "bool_domain" => {
36
21043
            ctx.add_span_and_doc_hover(&domain, "L_bool", SymbolKind::Domain, None, None);
37
21043
            Ok(Some(Domain::bool()))
38
        }
39
28265
        "int_domain" => parse_int_domain(ctx, domain),
40
3499
        "identifier" => {
41
780
            let Some(decl) = get_declaration_ptr_from_identifier(ctx, domain)? else {
42
52
                return Ok(None);
43
            };
44
728
            let Some(dom) = Domain::reference(decl) else {
45
13
                ctx.record_error(crate::errors::RecoverableParseError::new(
46
13
                    format!(
47
                        "The identifier '{}' is not a valid domain",
48
13
                        &ctx.source_code[domain.start_byte()..domain.end_byte()]
49
                    ),
50
13
                    Some(domain.range()),
51
                ));
52
13
                return Ok(None);
53
            };
54
715
            let name = &ctx.source_code[domain.start_byte()..domain.end_byte()];
55

            
56
            // Not form docs, because we need context specific hover info
57
715
            span_with_hover(
58
715
                &domain,
59
715
                ctx.source_code,
60
715
                ctx.source_map,
61
715
                HoverInfo {
62
715
                    description: format!("Domain reference: {name}"),
63
715
                    doc_key: None,
64
715
                    kind: Some(SymbolKind::Variable),
65
715
                    ty: None,
66
715
                    decl_span: None, // could link to the declaration span if we wanted
67
715
                },
68
            );
69
715
            Ok(Some(dom))
70
        }
71
2719
        "tuple_domain" => parse_tuple_domain(ctx, domain),
72
2524
        "matrix_domain" => parse_matrix_domain(ctx, domain),
73
234
        "record_domain" => parse_record_domain(ctx, domain),
74
195
        "set_domain" => parse_set_domain(ctx, domain),
75
        _ => {
76
13
            ctx.record_error(RecoverableParseError::new(
77
13
                format!("{} is not a supported domain type", domain.kind()),
78
13
                Some(domain.range()),
79
            ));
80
13
            Ok(None)
81
        }
82
    }
83
98100
}
84

            
85
2505
fn get_declaration_ptr_from_identifier(
86
2505
    ctx: &mut ParseContext,
87
2505
    identifier: Node,
88
2505
) -> Result<Option<DeclarationPtr>, FatalParseError> {
89
2505
    let name = Name::user(&ctx.source_code[identifier.start_byte()..identifier.end_byte()]);
90
2505
    let decl = ctx.symbols.as_ref().unwrap().read().lookup(&name);
91

            
92
2505
    if decl.is_none() {
93
78
        ctx.record_error(crate::errors::RecoverableParseError::new(
94
78
            format!("The identifier '{}' is not defined", name),
95
78
            Some(identifier.range()),
96
        ));
97
78
        return Ok(None);
98
2427
    }
99
2427
    match decl {
100
2427
        Some(decl) => Ok(Some(decl)),
101
        None => {
102
            ctx.record_error(crate::errors::RecoverableParseError::new(
103
                format!("The identifier '{}' is not defined", name),
104
                Some(identifier.range()),
105
            ));
106
            Ok(None)
107
        }
108
    }
109
2505
}
110

            
111
/// Parse an integer domain. Can be a single integer or a range.
112
24766
fn parse_int_domain(
113
24766
    ctx: &mut ParseContext,
114
24766
    int_domain: Node,
115
24766
) -> Result<Option<DomainPtr>, FatalParseError> {
116
24766
    let int_keyword_node = child!(int_domain, 0, "int");
117
24766
    if int_domain.child_count() == 1 {
118
        // for domains of just 'int' with no range
119
328
        ctx.add_span_and_doc_hover(&int_keyword_node, "L_int", SymbolKind::Domain, None, None);
120
328
        return Ok(Some(Domain::int(vec![Range::Bounded(i32::MIN, i32::MAX)])));
121
24438
    }
122

            
123
24438
    let Some(range_list) = field!(recover, ctx, int_domain, "ranges") else {
124
        return Ok(None);
125
    };
126
24438
    let mut ranges_unresolved: Vec<Range<IntVal>> = Vec::new();
127
24438
    let mut all_resolved = true;
128

            
129
25952
    for domain_component in named_children(&range_list) {
130
25952
        match domain_component.kind() {
131
25952
            "atom" | "arithmetic_expr" => {
132
2173
                let Some(int_val) = parse_int_val(ctx, domain_component)? else {
133
                    return Ok(None);
134
                };
135

            
136
2173
                if !matches!(int_val, IntVal::Const(_)) {
137
68
                    all_resolved = false;
138
2157
                }
139
2173
                ranges_unresolved.push(Range::Single(int_val));
140
            }
141
23779
            "int_range" => {
142
23779
                let lower_bound = match domain_component.child_by_field_name("lower") {
143
23766
                    Some(node) => {
144
23766
                        match parse_int_val(ctx, node)? {
145
23766
                            Some(val) => Some(val),
146
                            None => return Ok(None), // semantic error occurred
147
                        }
148
                    }
149
13
                    None => None,
150
                };
151
23779
                let upper_bound = match domain_component.child_by_field_name("upper") {
152
23753
                    Some(node) => {
153
23753
                        match parse_int_val(ctx, node)? {
154
23727
                            Some(val) => Some(val),
155
26
                            None => return Ok(None), // semantic error occurred
156
                        }
157
                    }
158
26
                    None => None,
159
                };
160

            
161
23753
                match (lower_bound, upper_bound) {
162
23714
                    (Some(lower), Some(upper)) => {
163
                        // Check if both bounds are constants and validate lower <= upper
164
23714
                        if let (IntVal::Const(l), IntVal::Const(u)) = (&lower, &upper) {
165
21635
                            if l > u {
166
39
                                ctx.record_error(crate::errors::RecoverableParseError::new(
167
39
                                    format!(
168
39
                                        "Invalid integer range: lower bound {} is greater than upper bound {}",
169
39
                                        l, u
170
39
                                    ),
171
39
                                    Some(domain_component.range()),
172
39
                                ));
173
21596
                            }
174
2079
                        } else {
175
2079
                            all_resolved = false;
176
2079
                        }
177
23714
                        ranges_unresolved.push(Range::Bounded(lower, upper));
178
                    }
179
26
                    (Some(lower), None) => {
180
26
                        if !matches!(lower, IntVal::Const(_)) {
181
                            all_resolved = false;
182
26
                        }
183
26
                        ranges_unresolved.push(Range::UnboundedR(lower));
184
                    }
185
13
                    (None, Some(upper)) => {
186
13
                        if !matches!(upper, IntVal::Const(_)) {
187
                            all_resolved = false;
188
13
                        }
189
13
                        ranges_unresolved.push(Range::UnboundedL(upper));
190
                    }
191
                    _ => {
192
                        ctx.record_error(RecoverableParseError::new(
193
                            "Invalid int range: must have at least a lower or upper bound"
194
                                .to_string(),
195
                            Some(domain_component.range()),
196
                        ));
197
                        return Ok(None);
198
                    }
199
                }
200
            }
201
            _ => {
202
                ctx.record_error(RecoverableParseError::new(
203
                    format!(
204
                        "Unexpected int domain component: {}",
205
                        domain_component.kind()
206
                    ),
207
                    Some(domain_component.range()),
208
                ));
209
                return Ok(None);
210
            }
211
        }
212
    }
213

            
214
    // If all values are resolved constants, convert IntVals to raw integers
215
24412
    if all_resolved {
216
22317
        let ranges: Vec<Range<i32>> = ranges_unresolved
217
22317
            .into_iter()
218
23779
            .map(|r| r.resolve())
219
22317
            .collect::<Result<_, _>>()
220
22317
            .map_err(|e| {
221
                FatalParseError::internal_error(
222
                    format!("could not resolve range: {e}"),
223
                    Some(int_domain.range()),
224
                )
225
            })?;
226

            
227
22317
        ctx.add_span_and_doc_hover(&int_keyword_node, "L_int", SymbolKind::Domain, None, None);
228
22317
        Ok(Some(Domain::int(ranges)))
229
    } else {
230
        // Otherwise, keep as an expression-based domain
231

            
232
        // Adding int keyword to the source map with hover info from documentation
233
2095
        ctx.add_span_and_doc_hover(&int_keyword_node, "L_int", SymbolKind::Domain, None, None);
234
2095
        Ok(Some(Domain::int(ranges_unresolved)))
235
    }
236
24766
}
237

            
238
// Helper function to parse a node into an IntVal
239
// Handles constants, references, and arbitrary expressions
240
49692
fn parse_int_val(ctx: &mut ParseContext, node: Node) -> Result<Option<IntVal>, FatalParseError> {
241
    // For atoms, try to parse as a constant integer first
242
49692
    if node.kind() == "atom" {
243
48918
        let text = &ctx.source_code[node.start_byte()..node.end_byte()];
244
48918
        if let Ok(integer) = text.parse::<i32>() {
245
47193
            return Ok(Some(IntVal::new_const(integer)));
246
1725
        }
247
        // Otherwise, check if it's an identifier reference
248
1725
        let Some(decl) = get_declaration_ptr_from_identifier(ctx, node)? else {
249
            // If identifier isn't defined, it's a semantic error
250
26
            return Ok(None);
251
        };
252
1699
        return Ok(Some(IntVal::Reference(Reference::new(decl))));
253
774
    }
254

            
255
    // For anything else, parse as an expression
256
774
    let Some(expr) = parse_expression(ctx, node)? else {
257
        return Ok(None);
258
    };
259
774
    Ok(Some(IntVal::Expr(Moo::new(expr))))
260
49692
}
261

            
262
195
fn parse_tuple_domain(
263
195
    ctx: &mut ParseContext,
264
195
    tuple_domain: Node,
265
195
) -> Result<Option<DomainPtr>, FatalParseError> {
266
195
    let mut domains: Vec<DomainPtr> = Vec::new();
267
390
    for domain in named_children(&tuple_domain) {
268
390
        let Some(parsed_domain) = parse_domain(ctx, domain)? else {
269
13
            return Ok(None);
270
        };
271
377
        domains.push(parsed_domain);
272
    }
273

            
274
    // extract the first child node which should be the 'tuple' keyword for hover info
275
182
    if let Some(first) = tuple_domain.child(0)
276
182
        && first.kind() == "tuple"
277
156
    {
278
        // Adding tuple to the source map with hover info from documentation
279
156
        ctx.add_span_and_doc_hover(&first, "L_tuple", SymbolKind::Domain, None, None);
280
156
    }
281

            
282
182
    Ok(Some(Domain::tuple(domains)))
283
195
}
284

            
285
2290
fn parse_matrix_domain(
286
2290
    ctx: &mut ParseContext,
287
2290
    matrix_domain: Node,
288
2290
) -> Result<Option<DomainPtr>, FatalParseError> {
289
2290
    let mut domains: Vec<DomainPtr> = Vec::new();
290
2290
    let Some(index_domain_list) = field!(recover, ctx, matrix_domain, "index_domain_list") else {
291
        return Ok(None);
292
    };
293
2841
    for domain in named_children(&index_domain_list) {
294
2841
        let Some(parsed_domain) = parse_domain(ctx, domain)? else {
295
            return Ok(None);
296
        };
297
2841
        domains.push(parsed_domain);
298
    }
299
2290
    let Some(value_domain_node) = field!(recover, ctx, matrix_domain, "value_domain") else {
300
        return Ok(None);
301
    };
302
2290
    let Some(value_domain) = parse_domain(ctx, value_domain_node)? else {
303
        return Ok(None);
304
    };
305

            
306
    // Adding matrix to the source map with hover info from documentation
307
2290
    let matrix_keyword_node = child!(matrix_domain, 0, "matrix");
308
2290
    ctx.add_span_and_doc_hover(
309
2290
        &matrix_keyword_node,
310
2290
        "matrix",
311
2290
        SymbolKind::Domain,
312
2290
        None,
313
2290
        None,
314
    );
315
2290
    Ok(Some(Domain::matrix(value_domain, domains)))
316
2290
}
317

            
318
39
fn parse_record_domain(
319
39
    ctx: &mut ParseContext,
320
39
    record_domain: Node,
321
39
) -> Result<Option<DomainPtr>, FatalParseError> {
322
39
    let mut record_entries: Vec<Field<DomainPtr>> = Vec::new();
323
78
    for record_entry in named_children(&record_domain) {
324
78
        let Some(name_node) = field!(recover, ctx, record_entry, "name") else {
325
            return Ok(None);
326
        };
327
78
        let name = Name::user(&ctx.source_code[name_node.start_byte()..name_node.end_byte()]);
328
78
        let Some(domain_node) = field!(recover, ctx, record_entry, "domain") else {
329
            return Ok(None);
330
        };
331
78
        let Some(value) = parse_domain(ctx, domain_node)? else {
332
            return Ok(None);
333
        };
334
78
        record_entries.push(Field { name, value });
335
    }
336

            
337
    // Adding record keyword to the source map with hover info from documentation
338
39
    let record_keyword_node = child!(record_domain, 0, "record");
339
39
    ctx.add_span_and_doc_hover(
340
39
        &record_keyword_node,
341
39
        "L_record",
342
39
        SymbolKind::Domain,
343
39
        None,
344
39
        None,
345
    );
346
39
    Ok(Some(Domain::record(record_entries)))
347
39
}
348

            
349
182
pub fn parse_set_domain(
350
182
    ctx: &mut ParseContext,
351
182
    set_domain: Node,
352
182
) -> Result<Option<DomainPtr>, FatalParseError> {
353
182
    let mut set_attribute: Option<SetAttr> = None;
354
182
    let mut value_domain: Option<DomainPtr> = None;
355

            
356
260
    for child in named_children(&set_domain) {
357
260
        match child.kind() {
358
260
            "set_attributes" => {
359
                // Check if we have both minSize and maxSize (minMax case)
360
78
                let min_value_node = child.child_by_field_name("min_value");
361
78
                let max_value_node = child.child_by_field_name("max_value");
362
78
                let size_value_node = child.child_by_field_name("size_value");
363

            
364
78
                if let (Some(min_node), Some(max_node)) = (min_value_node, max_value_node) {
365
                    // MinMax case
366
                    let Some(min_val) = parse_int(ctx, &min_node) else {
367
                        return Ok(None);
368
                    };
369
                    let Some(max_val) = parse_int(ctx, &max_node) else {
370
                        return Ok(None);
371
                    };
372

            
373
                    set_attribute = Some(SetAttr::new_min_max_size(min_val, max_val));
374
78
                } else if let Some(size_node) = size_value_node {
375
                    // Size case
376
                    let Some(size_val) = parse_int(ctx, &size_node) else {
377
                        return Ok(None);
378
                    };
379
                    set_attribute = Some(SetAttr::new_size(size_val));
380
78
                } else if let Some(min_node) = min_value_node {
381
                    // MinSize only case
382
78
                    let Some(min_val) = parse_int(ctx, &min_node) else {
383
                        return Ok(None);
384
                    };
385
78
                    set_attribute = Some(SetAttr::new_min_size(min_val));
386
                } else if let Some(max_node) = max_value_node {
387
                    // MaxSize only case
388
                    let Some(max_val) = parse_int(ctx, &max_node) else {
389
                        return Ok(None);
390
                    };
391
                    set_attribute = Some(SetAttr::new_max_size(max_val));
392
                }
393
            }
394
182
            "domain" => {
395
182
                let Some(parsed_domain) = parse_domain(ctx, child)? else {
396
                    return Ok(None);
397
                };
398
182
                value_domain = Some(parsed_domain);
399
            }
400
            _ => {
401
                ctx.record_error(RecoverableParseError::new(
402
                    format!("Unrecognized set domain child kind: {}", child.kind()),
403
                    Some(child.range()),
404
                ));
405
                return Ok(None);
406
            }
407
        }
408
    }
409

            
410
182
    if let Some(domain) = value_domain {
411
        // Adding set to the source map with hover info from documentation
412
182
        let set_keyword_node = child!(set_domain, 0, "set");
413
        // No documentation available for set domain, using fallback description
414
182
        ctx.add_span_and_doc_hover(&set_keyword_node, "set", SymbolKind::Domain, None, None);
415
182
        Ok(Some(Domain::set(set_attribute.unwrap_or_default(), domain)))
416
    } else {
417
        ctx.record_error(RecoverableParseError::new(
418
            "Set domain must have a value domain".to_string(),
419
            Some(set_domain.range()),
420
        ));
421
        Ok(None)
422
    }
423
182
}