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, FieldEntry, 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
96768
pub fn parse_domain(
18
96768
    ctx: &mut ParseContext,
19
96768
    domain: Node,
20
96768
) -> Result<Option<DomainPtr>, FatalParseError> {
21
96768
    match domain.kind() {
22
96768
        "domain" => {
23
48215
            let inner = match domain.child(0) {
24
48215
                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
48215
            parse_domain(ctx, inner)
34
        }
35
48553
        "bool_domain" => {
36
20980
            ctx.add_span_and_doc_hover(&domain, "L_bool", SymbolKind::Domain, None, None);
37
20980
            Ok(Some(Domain::bool()))
38
        }
39
27573
        "int_domain" => parse_int_domain(ctx, domain),
40
3433
        "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
2653
        "tuple_domain" => parse_tuple_domain(ctx, domain),
72
2458
        "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
96768
}
84

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

            
92
2091
    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
2013
    }
99
2013
    match decl {
100
2013
        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
2091
}
110

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

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

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

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

            
161
23229
                match (lower_bound, upper_bound) {
162
23190
                    (Some(lower), Some(upper)) => {
163
                        // Check if both bounds are constants and validate lower <= upper
164
23190
                        if let (IntVal::Const(l), IntVal::Const(u)) = (&lower, &upper) {
165
21481
                            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
21442
                            }
174
1709
                        } else {
175
1709
                            all_resolved = false;
176
1709
                        }
177
23190
                        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
23846
    if all_resolved {
216
22155
        let ranges: Vec<Range<i32>> = ranges_unresolved
217
22155
            .into_iter()
218
22155
            .map(|r| match r {
219
2097
                Range::Single(IntVal::Const(v)) => Range::Single(v),
220
21481
                Range::Bounded(IntVal::Const(l), IntVal::Const(u)) => Range::Bounded(l, u),
221
26
                Range::UnboundedR(IntVal::Const(l)) => Range::UnboundedR(l),
222
13
                Range::UnboundedL(IntVal::Const(u)) => Range::UnboundedL(u),
223
                Range::Unbounded => Range::Unbounded,
224
                _ => unreachable!("all_resolved should be true only if all are Const"),
225
23617
            })
226
22155
            .collect();
227

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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