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

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

            
10
use super::find::parse_find_statement;
11
use super::letting::parse_letting_statement;
12
use super::util::{get_tree, named_children};
13
use crate::errors::EssenceParseError;
14
use crate::expression::parse_expression;
15

            
16
/// Parse an Essence file into a Model using the tree-sitter parser.
17
269
pub fn parse_essence_file_native(
18
269
    path: &str,
19
269
    context: Arc<RwLock<Context<'static>>>,
20
269
) -> Result<Model, EssenceParseError> {
21
269
    let source_code = fs::read_to_string(path)
22
269
        .unwrap_or_else(|_| panic!("Failed to read the source code file {path}"));
23
269
    parse_essence_with_context(&source_code, context)
24
269
}
25

            
26
337
pub fn parse_essence_with_context(
27
337
    src: &str,
28
337
    context: Arc<RwLock<Context<'static>>>,
29
337
) -> Result<Model, EssenceParseError> {
30
337
    let (tree, source_code) = match get_tree(src) {
31
337
        Some(tree) => tree,
32
        None => {
33
            return Err(EssenceParseError::TreeSitterError(
34
                "Failed to parse source code".to_string(),
35
            ));
36
        }
37
    };
38

            
39
337
    let mut model = Model::new(context);
40
337
    let root_node = tree.root_node();
41
337
    let symbols_ptr = model.as_submodel().symbols_ptr_unchecked().clone();
42
960
    for statement in named_children(&root_node) {
43
960
        match statement.kind() {
44
960
            "single_line_comment" => {}
45
960
            "language_declaration" => {}
46
960
            "find_statement" => {
47
313
                let var_hashmap =
48
313
                    parse_find_statement(statement, &source_code, Some(symbols_ptr.clone()))?;
49
374
                for (name, domain) in var_hashmap {
50
374
                    model
51
374
                        .as_submodel_mut()
52
374
                        .symbols_mut()
53
374
                        .insert(DeclarationPtr::new_find(name, domain));
54
374
                }
55
            }
56
647
            "bool_expr" | "atom" | "comparison_expr" => {
57
371
                model.as_submodel_mut().add_constraint(parse_expression(
58
371
                    statement,
59
371
                    &source_code,
60
371
                    &statement,
61
371
                    Some(symbols_ptr.clone()),
62
35
                )?);
63
            }
64
276
            "language_label" => {}
65
276
            "letting_statement" => {
66
276
                let letting_vars =
67
276
                    parse_letting_statement(statement, &source_code, Some(symbols_ptr.clone()))?;
68
276
                model.as_submodel_mut().symbols_mut().extend(letting_vars);
69
            }
70
            "dominance_relation" => {
71
                let inner = statement
72
                    .child_by_field_name("expression")
73
                    .expect("Expected a sub-expression inside `dominanceRelation`");
74
                let expr =
75
                    parse_expression(inner, &source_code, &statement, Some(symbols_ptr.clone()))?;
76
                let dominance = Expression::DominanceRelation(Metadata::new(), Moo::new(expr));
77
                if model.dominance.is_some() {
78
                    return Err(EssenceParseError::syntax_error(
79
                        "Duplicate dominance relation".to_string(),
80
                        None,
81
                    ));
82
                }
83
                model.dominance = Some(dominance);
84
            }
85
            "ERROR" => {
86
                let raw_expr = &source_code[statement.start_byte()..statement.end_byte()];
87
                return Err(EssenceParseError::syntax_error(
88
                    format!("'{raw_expr}' is not a valid expression"),
89
                    Some(statement.range()),
90
                ));
91
            }
92
            _ => {
93
                let kind = statement.kind();
94
                return Err(EssenceParseError::syntax_error(
95
                    format!("Unrecognized top level statement kind: {kind}"),
96
                    Some(statement.range()),
97
                ));
98
            }
99
        }
100

            
101
        // check for errors (keyword as identifier)
102
925
        let result = keyword_as_identifier(root_node, &source_code);
103
925
        result?
104
    }
105
267
    Ok(model)
106
337
}
107

            
108
const KEYWORDS: [&str; 21] = [
109
    "forall", "exists", "such", "that", "letting", "find", "minimise", "maximise", "subject", "to",
110
    "where", "and", "or", "not", "if", "then", "else", "in", "sum", "product", "bool",
111
];
112

            
113
925
fn keyword_as_identifier(root: tree_sitter::Node, src: &str) -> Result<(), EssenceParseError> {
114
925
    let mut stack = vec![root];
115
71660
    while let Some(node) = stack.pop() {
116
70770
        if (node.kind() == "variable" || node.kind() == "identifier" || node.kind() == "parameter")
117
4552
            && let Ok(text) = node.utf8_text(src.as_bytes())
118
        {
119
4552
            let ident = text.trim();
120
4552
            if KEYWORDS.contains(&ident) {
121
35
                let start_point = node.start_position();
122
35
                let end_point = node.end_position();
123
35
                return Err(EssenceParseError::syntax_error(
124
35
                    format!("Keyword '{ident}' used as identifier"),
125
35
                    Some(tree_sitter::Range {
126
35
                        start_byte: node.start_byte(),
127
35
                        end_byte: node.end_byte(),
128
35
                        start_point,
129
35
                        end_point,
130
35
                    }),
131
35
                ));
132
4517
            }
133
66218
        }
134

            
135
        // push children onto stack
136
70735
        for i in 0..node.child_count() {
137
69888
            if let Some(child) = u32::try_from(i).ok().and_then(|i| node.child(i)) {
138
69888
                stack.push(child);
139
69888
            }
140
        }
141
    }
142
890
    Ok(())
143
925
}
144

            
145
2
pub fn parse_essence(src: &str) -> Result<Model, EssenceParseError> {
146
2
    let context = Arc::new(RwLock::new(Context::default()));
147
2
    parse_essence_with_context(src, context)
148
2
}
149

            
150
mod test {
151
    #[allow(unused_imports)]
152
    use crate::parse_essence;
153
    #[allow(unused_imports)]
154
    use conjure_cp_core::ast::{Atom, Expression, Metadata, Moo, Name};
155
    #[allow(unused_imports)]
156
    use conjure_cp_core::{domain_int, matrix_expr, range};
157
    #[allow(unused_imports)]
158
    use std::ops::Deref;
159

            
160
    #[test]
161
1
    pub fn test_parse_xyz() {
162
1
        let src = "
163
1
        find x, y, z : int(1..4)
164
1
        such that x + y + z = 4
165
1
        such that x >= y
166
1
        ";
167

            
168
1
        let model = parse_essence(src).unwrap();
169

            
170
1
        let st = model.as_submodel().symbols();
171
1
        let x = st.lookup(&Name::user("x")).unwrap();
172
1
        let y = st.lookup(&Name::user("y")).unwrap();
173
1
        let z = st.lookup(&Name::user("z")).unwrap();
174
1
        assert_eq!(x.domain(), Some(domain_int!(1..4)));
175
1
        assert_eq!(y.domain(), Some(domain_int!(1..4)));
176
1
        assert_eq!(z.domain(), Some(domain_int!(1..4)));
177

            
178
1
        let constraints = model.as_submodel().constraints();
179
1
        assert_eq!(constraints.len(), 2);
180

            
181
1
        let c1 = constraints[0].clone();
182
1
        let x_e = Expression::Atomic(Metadata::new(), Atom::new_ref(x));
183
1
        let y_e = Expression::Atomic(Metadata::new(), Atom::new_ref(y));
184
1
        let z_e = Expression::Atomic(Metadata::new(), Atom::new_ref(z));
185
1
        assert_eq!(
186
            c1,
187
1
            Expression::Eq(
188
1
                Metadata::new(),
189
1
                Moo::new(Expression::Sum(
190
1
                    Metadata::new(),
191
1
                    Moo::new(matrix_expr!(
192
1
                        Expression::Sum(
193
1
                            Metadata::new(),
194
1
                            Moo::new(matrix_expr!(x_e.clone(), y_e.clone()))
195
1
                        ),
196
1
                        z_e
197
1
                    ))
198
1
                )),
199
1
                Moo::new(Expression::Atomic(Metadata::new(), 4.into()))
200
1
            )
201
        );
202

            
203
1
        let c2 = constraints[1].clone();
204
1
        assert_eq!(
205
            c2,
206
1
            Expression::Geq(Metadata::new(), Moo::new(x_e), Moo::new(y_e))
207
        );
208
1
    }
209

            
210
    #[test]
211
1
    pub fn test_parse_letting_index() {
212
1
        let src = "
213
1
        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) ]
214
1
        find b: int(1..5)
215
1
        such that
216
1
        b < a[-2,2],
217
1
        allDiff(a[-2,..])
218
1
        ";
219

            
220
1
        let model = parse_essence(src).unwrap();
221
1
        let st = model.as_submodel().symbols();
222
1
        let a_decl = st.lookup(&Name::user("a")).unwrap();
223
1
        let a = a_decl.as_value_letting().unwrap().deref().clone();
224
1
        assert_eq!(
225
            a,
226
1
            matrix_expr!(
227
1
                matrix_expr!(1.into(), 2.into(), 3.into() ; domain_int!(1, 2, 4)),
228
1
                matrix_expr!(1.into(), 3.into(), 2.into() ; domain_int!(1, 2, 4)),
229
1
                matrix_expr!(3.into(), 2.into(), 1.into() ; domain_int!(1, 2, 4));
230
1
                domain_int!(-2..0)
231
            )
232
        )
233
1
    }
234
}