1
use crate::expression::{parse_binary_expression, parse_expression};
2
use crate::parser::abstract_literal::parse_abstract;
3
use crate::parser::comprehension::parse_comprehension;
4
use crate::util::named_children;
5
use crate::{EssenceParseError, field, named_child};
6
use conjure_cp_core::ast::{Atom, Expression, Literal, Metadata, Moo, Name, SymbolTable};
7
use std::cell::RefCell;
8
use std::rc::Rc;
9
use tree_sitter::Node;
10
use ustr::Ustr;
11

            
12
732
pub fn parse_atom(
13
732
    node: &Node,
14
732
    source_code: &str,
15
732
    root: &Node,
16
732
    symbols_ptr: Option<Rc<RefCell<SymbolTable>>>,
17
732
) -> Result<Expression, EssenceParseError> {
18
732
    match node.kind() {
19
732
        "atom" => parse_atom(&named_child!(node), source_code, root, symbols_ptr),
20
366
        "metavar" => {
21
318
            let ident = field!(node, "identifier");
22
318
            let name_str = &source_code[ident.start_byte()..ident.end_byte()];
23
318
            Ok(Expression::Metavar(Metadata::new(), Ustr::from(name_str)))
24
        }
25
48
        "identifier" => parse_variable(node, source_code, symbols_ptr)
26
            .map(|var| Expression::Atomic(Metadata::new(), var)),
27
48
        "from_solution" => {
28
            if root.kind() != "dominance_relation" {
29
                return Err(EssenceParseError::syntax_error(
30
                    "fromSolution only allowed inside dominance relations".to_string(),
31
                    Some(node.range()),
32
                ));
33
            }
34

            
35
            let inner = parse_variable(&field!(node, "variable"), source_code, symbols_ptr)?;
36
            Ok(Expression::FromSolution(Metadata::new(), Moo::new(inner)))
37
        }
38
48
        "constant" => {
39
48
            let lit = parse_constant(node, source_code)?;
40
48
            Ok(Expression::Atomic(Metadata::new(), Atom::Literal(lit)))
41
        }
42
        "matrix" | "record" | "tuple" | "set_literal" => {
43
            parse_abstract(node, source_code, symbols_ptr)
44
                .map(|l| Expression::AbstractLiteral(Metadata::new(), l))
45
        }
46
        "flatten" => parse_flatten(node, source_code, root, symbols_ptr),
47
        "index_or_slice" => parse_index_or_slice(node, source_code, root, symbols_ptr),
48
        // for now, assume is binary since powerset isn't implemented
49
        // TODO: add powerset support under "set_operation"
50
        "set_operation" => parse_binary_expression(node, source_code, root, symbols_ptr),
51
        "comprehension" => parse_comprehension(node, source_code, root, symbols_ptr),
52
        _ => Err(EssenceParseError::syntax_error(
53
            format!("Expected atom, got: {}", node.kind()),
54
            Some(node.range()),
55
        )),
56
    }
57
732
}
58

            
59
fn parse_flatten(
60
    node: &Node,
61
    source_code: &str,
62
    root: &Node,
63
    symbols_ptr: Option<Rc<RefCell<SymbolTable>>>,
64
) -> Result<Expression, EssenceParseError> {
65
    let expr_node = field!(node, "expression");
66
    let expr = parse_atom(&expr_node, source_code, root, symbols_ptr)?;
67

            
68
    if node.child_by_field_name("depth").is_some() {
69
        let depth_node = field!(node, "depth");
70
        let depth = parse_int(&depth_node, source_code)?;
71
        let depth_expression =
72
            Expression::Atomic(Metadata::new(), Atom::Literal(Literal::Int(depth)));
73
        Ok(Expression::Flatten(
74
            Metadata::new(),
75
            Some(Moo::new(depth_expression)),
76
            Moo::new(expr),
77
        ))
78
    } else {
79
        Ok(Expression::Flatten(Metadata::new(), None, Moo::new(expr)))
80
    }
81
}
82

            
83
fn parse_index_or_slice(
84
    node: &Node,
85
    source_code: &str,
86
    root: &Node,
87
    symbols_ptr: Option<Rc<RefCell<SymbolTable>>>,
88
) -> Result<Expression, EssenceParseError> {
89
    let collection = parse_atom(
90
        &field!(node, "collection"),
91
        source_code,
92
        root,
93
        symbols_ptr.clone(),
94
    )?;
95
    let mut indices = Vec::new();
96
    for idx_node in named_children(&field!(node, "indices")) {
97
        indices.push(parse_index(&idx_node, source_code, symbols_ptr.clone())?);
98
    }
99

            
100
    let has_null_idx = indices.iter().any(|idx| idx.is_none());
101
    // TODO: We could check whether the slice/index is safe here
102
    if has_null_idx {
103
        // It's a slice
104
        Ok(Expression::UnsafeSlice(
105
            Metadata::new(),
106
            Moo::new(collection),
107
            indices,
108
        ))
109
    } else {
110
        // It's an index
111
        let idx_exprs: Vec<Expression> = indices.into_iter().map(|idx| idx.unwrap()).collect();
112
        Ok(Expression::UnsafeIndex(
113
            Metadata::new(),
114
            Moo::new(collection),
115
            idx_exprs,
116
        ))
117
    }
118
}
119

            
120
fn parse_index(
121
    node: &Node,
122
    source_code: &str,
123
    symbols_ptr: Option<Rc<RefCell<SymbolTable>>>,
124
) -> Result<Option<Expression>, EssenceParseError> {
125
    match node.kind() {
126
        "arithmetic_expr" => Ok(Some(parse_expression(
127
            *node,
128
            source_code,
129
            node,
130
            symbols_ptr,
131
        )?)),
132
        "null_index" => Ok(None),
133
        _ => Err(EssenceParseError::syntax_error(
134
            format!("Expected an index, got: '{}'", node.kind()),
135
            Some(node.range()),
136
        )),
137
    }
138
}
139

            
140
fn parse_variable(
141
    node: &Node,
142
    source_code: &str,
143
    symbols_ptr: Option<Rc<RefCell<SymbolTable>>>,
144
) -> Result<Atom, EssenceParseError> {
145
    let raw_name = &source_code[node.start_byte()..node.end_byte()];
146
    let name = Name::user(raw_name.trim());
147
    if let Some(symbols) = symbols_ptr {
148
        if let Some(decl) = symbols.borrow().lookup(&name) {
149
            Ok(Atom::Reference(conjure_cp_core::ast::Reference::new(decl)))
150
        } else {
151
            Err(EssenceParseError::syntax_error(
152
                format!("Undefined variable: '{}'", raw_name),
153
                Some(node.range()),
154
            ))
155
        }
156
    } else {
157
        Err(EssenceParseError::syntax_error(
158
            format!(
159
                "Found variable '{raw_name}'; Did you mean to pass a meta-variable '&{raw_name}'?\n\
160
            A symbol table is needed to resolve variable names, but none exists in this context."
161
            ),
162
            Some(node.range()),
163
        ))
164
    }
165
}
166

            
167
48
fn parse_constant(node: &Node, source_code: &str) -> Result<Literal, EssenceParseError> {
168
48
    let inner = named_child!(node);
169
48
    let raw_value = &source_code[inner.start_byte()..inner.end_byte()];
170
48
    match inner.kind() {
171
48
        "integer" => {
172
31
            let value = parse_int(&inner, source_code)?;
173
31
            Ok(Literal::Int(value))
174
        }
175
17
        "TRUE" => Ok(Literal::Bool(true)),
176
6
        "FALSE" => Ok(Literal::Bool(false)),
177
        _ => Err(EssenceParseError::syntax_error(
178
            format!(
179
                "'{}' (kind: '{}') is not a valid constant",
180
                raw_value,
181
                inner.kind()
182
            ),
183
            Some(inner.range()),
184
        )),
185
    }
186
48
}
187

            
188
31
fn parse_int(node: &Node, source_code: &str) -> Result<i32, EssenceParseError> {
189
31
    let raw_value = &source_code[node.start_byte()..node.end_byte()];
190
31
    raw_value.parse::<i32>().map_err(|_e| {
191
        if raw_value.is_empty() {
192
            EssenceParseError::syntax_error(
193
                "Expected an integer here".to_string(),
194
                Some(node.range()),
195
            )
196
        } else {
197
            EssenceParseError::syntax_error(
198
                format!("'{raw_value}' is not a valid integer"),
199
                Some(node.range()),
200
            )
201
        }
202
    })
203
31
}