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, IntVal, Moo, Name, Range, RecordEntry, 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
95559
pub fn parse_domain(
18
95559
    ctx: &mut ParseContext,
19
95559
    domain: Node,
20
95559
) -> Result<Option<DomainPtr>, FatalParseError> {
21
95559
    match domain.kind() {
22
95559
        "domain" => {
23
47604
            let inner = match domain.child(0) {
24
47604
                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
47604
            parse_domain(ctx, inner)
34
        }
35
47955
        "bool_domain" => {
36
20941
            ctx.add_span_and_doc_hover(&domain, "L_bool", SymbolKind::Domain, None, None);
37
20941
            Ok(Some(Domain::bool()))
38
        }
39
27014
        "int_domain" => parse_int_domain(ctx, domain),
40
3381
        "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
                    kind: Some(SymbolKind::Variable),
64
715
                    ty: None,
65
715
                    decl_span: None, // could link to the declaration span if we wanted
66
715
                },
67
            );
68
715
            Ok(Some(dom))
69
        }
70
2601
        "tuple_domain" => parse_tuple_domain(ctx, domain),
71
2406
        "matrix_domain" => parse_matrix_domain(ctx, domain),
72
208
        "record_domain" => parse_record_domain(ctx, domain),
73
182
        "set_domain" => parse_set_domain(ctx, domain),
74
        _ => {
75
13
            ctx.record_error(RecoverableParseError::new(
76
13
                format!("{} is not a supported domain type", domain.kind()),
77
13
                Some(domain.range()),
78
            ));
79
13
            Ok(None)
80
        }
81
    }
82
95559
}
83

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

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

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

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

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

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

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

            
213
    // If all values are resolved constants, convert IntVals to raw integers
214
23339
    if all_resolved {
215
21648
        let ranges: Vec<Range<i32>> = ranges_unresolved
216
21648
            .into_iter()
217
21648
            .map(|r| match r {
218
2097
                Range::Single(IntVal::Const(v)) => Range::Single(v),
219
20974
                Range::Bounded(IntVal::Const(l), IntVal::Const(u)) => Range::Bounded(l, u),
220
26
                Range::UnboundedR(IntVal::Const(l)) => Range::UnboundedR(l),
221
13
                Range::UnboundedL(IntVal::Const(u)) => Range::UnboundedL(u),
222
                Range::Unbounded => Range::Unbounded,
223
                _ => unreachable!("all_resolved should be true only if all are Const"),
224
23110
            })
225
21648
            .collect();
226

            
227
21648
        ctx.add_span_and_doc_hover(&int_keyword_node, "L_int", SymbolKind::Domain, None, None);
228
21648
        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
1691
        ctx.add_span_and_doc_hover(&int_keyword_node, "L_int", SymbolKind::Domain, None, None);
234
1691
        Ok(Some(Domain::int(ranges_unresolved)))
235
    }
236
23633
}
237

            
238
// Helper function to parse a node into an IntVal
239
// Handles constants, references, and arbitrary expressions
240
47588
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
47588
    if node.kind() == "atom" {
243
46824
        let text = &ctx.source_code[node.start_byte()..node.end_byte()];
244
46824
        if let Ok(integer) = text.parse::<i32>() {
245
45513
            return Ok(Some(IntVal::Const(integer)));
246
1311
        }
247
        // Otherwise, check if it's an identifier reference
248
1311
        let Some(decl) = get_declaration_ptr_from_identifier(ctx, node)? else {
249
            // If identifier isn't defined, its a semantic error
250
26
            return Ok(None);
251
        };
252
1285
        return Ok(Some(IntVal::Reference(Reference::new(decl))));
253
764
    }
254

            
255
    // For anything else, parse as an expression
256
764
    let Some(expr) = parse_expression(ctx, node)? else {
257
        return Ok(None);
258
    };
259
764
    Ok(Some(IntVal::Expr(Moo::new(expr))))
260
47588
}
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
2198
fn parse_matrix_domain(
286
2198
    ctx: &mut ParseContext,
287
2198
    matrix_domain: Node,
288
2198
) -> Result<Option<DomainPtr>, FatalParseError> {
289
2198
    let mut domains: Vec<DomainPtr> = Vec::new();
290
2198
    let Some(index_domain_list) = field!(recover, ctx, matrix_domain, "index_domain_list") else {
291
        return Ok(None);
292
    };
293
2675
    for domain in named_children(&index_domain_list) {
294
2675
        let Some(parsed_domain) = parse_domain(ctx, domain)? else {
295
            return Ok(None);
296
        };
297
2675
        domains.push(parsed_domain);
298
    }
299
2198
    let Some(value_domain_node) = field!(recover, ctx, matrix_domain, "value_domain") else {
300
        return Ok(None);
301
    };
302
2198
    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
2198
    let matrix_keyword_node = child!(matrix_domain, 0, "matrix");
308
2198
    ctx.add_span_and_doc_hover(
309
2198
        &matrix_keyword_node,
310
2198
        "matrix",
311
2198
        SymbolKind::Domain,
312
2198
        None,
313
2198
        None,
314
    );
315
2198
    Ok(Some(Domain::matrix(value_domain, domains)))
316
2198
}
317

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

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

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

            
356
247
    for child in named_children(&set_domain) {
357
247
        match child.kind() {
358
247
            "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
169
            "domain" => {
395
169
                let Some(parsed_domain) = parse_domain(ctx, child)? else {
396
                    return Ok(None);
397
                };
398
169
                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
169
    if let Some(domain) = value_domain {
411
        // Adding set to the source map with hover info from documentation
412
169
        let set_keyword_node = child!(set_domain, 0, "set");
413
        // No documentation available for set domain, using fallback description
414
169
        ctx.add_span_and_doc_hover(&set_keyword_node, "set", SymbolKind::Domain, None, None);
415
169
        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
169
}