1
use std::sync::{Arc, RwLock};
2
use std::{fs, vec};
3

            
4
use conjure_cp_core::Model;
5
use conjure_cp_core::ast::assertions::debug_assert_model_well_formed;
6
use conjure_cp_core::ast::{DeclarationPtr, Expression, Metadata, Moo};
7
use conjure_cp_core::context::Context;
8
#[allow(unused)]
9
use uniplate::Uniplate;
10

            
11
use super::ParseContext;
12
use super::find::parse_find_statement;
13
use super::letting::parse_letting_statement;
14
use super::util::{TypecheckingContext, get_tree};
15
use crate::diagnostics::diagnostics_api::SymbolKind;
16
use crate::diagnostics::source_map::{HoverInfo, SourceMap, span_with_hover};
17
use crate::errors::{FatalParseError, ParseErrorCollection, RecoverableParseError};
18
use crate::expression::parse_expression;
19
use crate::field;
20
use crate::syntax_errors::detect_syntactic_errors;
21

            
22
/// Parse an Essence file into a Model using the tree-sitter parser.
23
1175
pub fn parse_essence_file_native(
24
1175
    path: &str,
25
1175
    context: Arc<RwLock<Context<'static>>>,
26
1175
) -> Result<Model, Box<ParseErrorCollection>> {
27
1175
    let source_code = fs::read_to_string(path)
28
1175
        .unwrap_or_else(|_| panic!("Failed to read the source code file {path}"));
29

            
30
1175
    let mut errors = vec![];
31
1175
    let model = parse_essence_with_context(&source_code, context, &mut errors);
32

            
33
1175
    match model {
34
866
        Ok(Some(m)) => {
35
866
            debug_assert_model_well_formed(&m, "tree-sitter");
36
866
            Ok(m)
37
        }
38
        Ok(None) => {
39
            // Recoverable errors were found, return them as a ParseErrorCollection
40
309
            Err(Box::new(ParseErrorCollection::multiple(
41
309
                errors,
42
309
                Some(source_code),
43
309
                Some(path.to_string()),
44
309
            )))
45
        }
46
        Err(fatal) => {
47
            // Fatal error - wrap in ParseErrorCollection::Fatal
48
            Err(Box::new(ParseErrorCollection::fatal(fatal)))
49
        }
50
    }
51
1175
}
52

            
53
2330
pub fn parse_essence_with_context(
54
2330
    src: &str,
55
2330
    context: Arc<RwLock<Context<'static>>>,
56
2330
    errors: &mut Vec<RecoverableParseError>,
57
2330
) -> Result<Option<Model>, FatalParseError> {
58
2330
    match parse_essence_with_context_and_map(src, context, errors)? {
59
899
        Some((model, _source_map)) => Ok(Some(model)),
60
1431
        None => Ok(None),
61
    }
62
2330
}
63

            
64
1671
pub fn parse_essence_with_context_and_map(
65
1671
    src: &str,
66
1671
    context: Arc<RwLock<Context<'static>>>,
67
1671
    errors: &mut Vec<RecoverableParseError>,
68
1671
) -> Result<Option<(Model, SourceMap)>, FatalParseError> {
69
1671
    let (tree, source_code) = match get_tree(src) {
70
1671
        Some(tree) => tree,
71
        None => {
72
            return Err(FatalParseError::TreeSitterError(
73
                "Failed to parse source code".to_string(),
74
            ));
75
        }
76
    };
77

            
78
1671
    if tree.root_node().has_error() {
79
690
        detect_syntactic_errors(src, &tree, errors);
80
690
        return Ok(None);
81
981
    }
82

            
83
981
    let mut model = Model::new(context);
84
981
    let mut source_map = SourceMap::default();
85
981
    let root_node = tree.root_node();
86

            
87
    // Create a ParseContext
88
981
    let mut ctx = ParseContext::new(
89
981
        &source_code,
90
981
        &root_node,
91
981
        Some(model.symbols_ptr_unchecked().clone()),
92
981
        errors,
93
981
        &mut source_map,
94
    );
95

            
96
981
    let mut cursor = root_node.walk();
97
5634
    for statement in root_node.children(&mut cursor) {
98
        /*
99
           since find and letting are unnamed children
100
           hover info is added here.
101
           other unnamed children will be skipped.
102
        */
103
5634
        if statement.kind() == "find" {
104
1157
            span_with_hover(
105
1157
                &statement,
106
1157
                ctx.source_code,
107
1157
                ctx.source_map,
108
1157
                HoverInfo {
109
1157
                    description: "Find keyword".to_string(),
110
1157
                    kind: Some(SymbolKind::Find),
111
1157
                    ty: None,
112
1157
                    decl_span: None,
113
1157
                },
114
1157
            );
115
4477
        } else if statement.kind() == "letting" {
116
706
            span_with_hover(
117
706
                &statement,
118
706
                ctx.source_code,
119
706
                ctx.source_map,
120
706
                HoverInfo {
121
706
                    description: "Letting keyword".to_string(),
122
706
                    kind: Some(SymbolKind::Letting),
123
706
                    ty: None,
124
706
                    decl_span: None,
125
706
                },
126
706
            );
127
3771
        }
128

            
129
5634
        if !statement.is_named() {
130
2817
            continue;
131
2817
        }
132

            
133
2817
        match statement.kind() {
134
2817
            "single_line_comment" => {}
135
2817
            "language_declaration" => {}
136
2817
            "find_statement" => {
137
1157
                let var_hashmap = parse_find_statement(&mut ctx, statement)?;
138
1293
                for (name, domain) in var_hashmap {
139
1293
                    model
140
1293
                        .symbols_mut()
141
1293
                        .insert(DeclarationPtr::new_find(name, domain));
142
1293
                }
143
            }
144
1660
            "bool_expr" | "atom" | "comparison_expr" => {
145
954
                ctx.typechecking_context = TypecheckingContext::Boolean;
146
954
                let Some(expr) = parse_expression(&mut ctx, statement)? else {
147
220
                    continue;
148
                };
149
734
                model.add_constraint(expr);
150
            }
151
706
            "language_label" => {}
152
706
            "letting_statement" => {
153
706
                let Some(letting_vars) = parse_letting_statement(&mut ctx, statement)? else {
154
                    continue;
155
                };
156
706
                model.symbols_mut().extend(letting_vars);
157
            }
158
            "dominance_relation" => {
159
                let inner = field!(statement, "expression");
160
                let Some(expr) = parse_expression(&mut ctx, inner)? else {
161
                    continue;
162
                };
163
                let dominance = Expression::DominanceRelation(Metadata::new(), Moo::new(expr));
164
                if model.dominance.is_some() {
165
                    ctx.record_error(RecoverableParseError::new(
166
                        "Duplicate dominance relation".to_string(),
167
                        None,
168
                    ));
169
                    continue;
170
                }
171
                model.dominance = Some(dominance);
172
            }
173
            _ => {
174
                return Err(FatalParseError::internal_error(
175
                    format!("Unexpected top-level statement: {}", statement.kind()),
176
                    Some(statement.range()),
177
                ));
178
            }
179
        }
180
    }
181

            
182
    // check for errors (keyword as identifier)
183
981
    keyword_as_identifier(&mut ctx);
184

            
185
    // Check if there were any recoverable errors
186
981
    if !errors.is_empty() {
187
363
        return Ok(None);
188
618
    }
189
    // otherwise return the model
190
618
    Ok(Some((model, source_map)))
191
1671
}
192

            
193
const KEYWORDS: [&str; 21] = [
194
    "forall", "exists", "such", "that", "letting", "find", "minimise", "maximise", "subject", "to",
195
    "where", "and", "or", "not", "if", "then", "else", "in", "sum", "product", "bool",
196
];
197

            
198
981
fn keyword_as_identifier(ctx: &mut ParseContext) {
199
981
    let mut stack = vec![*ctx.root];
200
56583
    while let Some(node) = stack.pop() {
201
55602
        if (node.kind() == "variable" || node.kind() == "identifier" || node.kind() == "parameter")
202
4072
            && let Ok(text) = node.utf8_text(ctx.source_code.as_bytes())
203
        {
204
4072
            let ident = text.trim();
205
4072
            if KEYWORDS.contains(&ident) {
206
88
                let start_point = node.start_position();
207
88
                let end_point = node.end_position();
208
88
                ctx.errors.push(RecoverableParseError::new(
209
88
                    format!("Keyword '{ident}' used as identifier"),
210
88
                    Some(tree_sitter::Range {
211
88
                        start_byte: node.start_byte(),
212
88
                        end_byte: node.end_byte(),
213
88
                        start_point,
214
88
                        end_point,
215
88
                    }),
216
88
                ));
217
3984
            }
218
51530
        }
219

            
220
        // push children onto stack
221
55602
        for i in 0..node.child_count() {
222
54621
            if let Some(child) = u32::try_from(i).ok().and_then(|i| node.child(i)) {
223
54621
                stack.push(child);
224
54621
            }
225
        }
226
    }
227
981
}
228

            
229
4
pub fn parse_essence(src: &str) -> Result<(Model, SourceMap), Box<ParseErrorCollection>> {
230
4
    let context = Arc::new(RwLock::new(Context::default()));
231
4
    let mut errors = vec![];
232
4
    match parse_essence_with_context_and_map(src, context, &mut errors) {
233
4
        Ok(Some((model, source_map))) => {
234
4
            debug_assert_model_well_formed(&model, "tree-sitter");
235
4
            Ok((model, source_map))
236
        }
237
        Ok(None) => {
238
            // Recoverable errors were found, return them as a ParseErrorCollection
239
            Err(Box::new(ParseErrorCollection::multiple(
240
                errors,
241
                Some(src.to_string()),
242
                None,
243
            )))
244
        }
245
        Err(fatal) => Err(Box::new(ParseErrorCollection::fatal(fatal))),
246
    }
247
4
}
248

            
249
mod test {
250
    #[allow(unused_imports)]
251
    use crate::parse_essence;
252
    #[allow(unused_imports)]
253
    use conjure_cp_core::ast::{Atom, Expression, Metadata, Moo, Name};
254
    #[allow(unused_imports)]
255
    use conjure_cp_core::{domain_int, matrix_expr, range};
256
    #[allow(unused_imports)]
257
    use std::ops::Deref;
258

            
259
    #[test]
260
2
    pub fn test_parse_xyz() {
261
2
        let src = "
262
2
        find x, y, z : int(1..4)
263
2
        such that x + y + z = 4
264
2
        such that x >= y
265
2
        ";
266

            
267
2
        let (model, _source_map) = parse_essence(src).unwrap();
268

            
269
2
        let st = model.symbols();
270
2
        let x = st.lookup(&Name::user("x")).unwrap();
271
2
        let y = st.lookup(&Name::user("y")).unwrap();
272
2
        let z = st.lookup(&Name::user("z")).unwrap();
273
2
        assert_eq!(x.domain(), Some(domain_int!(1..4)));
274
2
        assert_eq!(y.domain(), Some(domain_int!(1..4)));
275
2
        assert_eq!(z.domain(), Some(domain_int!(1..4)));
276

            
277
2
        let constraints = model.constraints();
278
2
        assert_eq!(constraints.len(), 2);
279

            
280
2
        let c1 = constraints[0].clone();
281
2
        let x_e = Expression::Atomic(Metadata::new(), Atom::new_ref(x));
282
2
        let y_e = Expression::Atomic(Metadata::new(), Atom::new_ref(y));
283
2
        let z_e = Expression::Atomic(Metadata::new(), Atom::new_ref(z));
284
2
        assert_eq!(
285
            c1,
286
2
            Expression::Eq(
287
2
                Metadata::new(),
288
2
                Moo::new(Expression::Sum(
289
2
                    Metadata::new(),
290
2
                    Moo::new(matrix_expr!(
291
2
                        Expression::Sum(
292
2
                            Metadata::new(),
293
2
                            Moo::new(matrix_expr!(x_e.clone(), y_e.clone()))
294
2
                        ),
295
2
                        z_e
296
2
                    ))
297
2
                )),
298
2
                Moo::new(Expression::Atomic(Metadata::new(), 4.into()))
299
2
            )
300
        );
301

            
302
2
        let c2 = constraints[1].clone();
303
2
        assert_eq!(
304
            c2,
305
2
            Expression::Geq(Metadata::new(), Moo::new(x_e), Moo::new(y_e))
306
        );
307
2
    }
308

            
309
    #[test]
310
2
    pub fn test_parse_letting_index() {
311
2
        let src = "
312
2
        letting a be [ [ 1,2,3 ; int(1,2,4) ], [ 1,3,2 ; int(1,2,4) ], [ 3,2,1 ; int(1,2,4) ] ; int(-2..0) ]
313
2
        find b: int(1..5)
314
2
        such that
315
2
        b < a[-2,2],
316
2
        allDiff(a[-2,..])
317
2
        ";
318

            
319
2
        let (model, _source_map) = parse_essence(src).unwrap();
320
2
        let st = model.symbols();
321
2
        let a_decl = st.lookup(&Name::user("a")).unwrap();
322
2
        let a = a_decl.as_value_letting().unwrap().deref().clone();
323
2
        assert_eq!(
324
            a,
325
2
            matrix_expr!(
326
2
                matrix_expr!(1.into(), 2.into(), 3.into() ; domain_int!(1, 2, 4)),
327
2
                matrix_expr!(1.into(), 3.into(), 2.into() ; domain_int!(1, 2, 4)),
328
2
                matrix_expr!(3.into(), 2.into(), 1.into() ; domain_int!(1, 2, 4));
329
2
                domain_int!(-2..0)
330
            )
331
        )
332
2
    }
333
}