Skip to main content

oak_haskell/builder/
mod.rs

1use crate::{
2    ast::*,
3    language::HaskellLanguage,
4    lexer::token_type::HaskellTokenType,
5    parser::{HaskellParser, element_type::HaskellElementType},
6};
7use oak_core::{Builder, BuilderCache, GreenNode, OakDiagnostics, OakError, RedNode, RedTree, Source, SourceText, TextEdit};
8
9/// AST builder for the Haskell language.
10#[derive(Clone)]
11pub struct HaskellBuilder<'config> {
12    /// Language configuration.
13    config: &'config HaskellLanguage,
14}
15
16impl<'config> HaskellBuilder<'config> {
17    /// Creates a new Haskell builder with the given configuration.
18    pub fn new(config: &'config HaskellLanguage) -> Self {
19        Self { config }
20    }
21}
22
23impl<'config> Builder<HaskellLanguage> for HaskellBuilder<'config> {
24    fn build<'a, S: Source + ?Sized>(&self, source: &S, edits: &[TextEdit], _cache: &'a mut impl BuilderCache<HaskellLanguage>) -> OakDiagnostics<HaskellRoot> {
25        let parser = HaskellParser::new(self.config);
26        let lexer = crate::lexer::HaskellLexer::new(&self.config);
27
28        let mut cache = oak_core::parser::session::ParseSession::<HaskellLanguage>::default();
29        let parse_result = oak_core::parser::parse(&parser, &lexer, source, edits, &mut cache);
30
31        match parse_result.result {
32            Ok(green_tree) => {
33                let source_text = SourceText::new(source.get_text_in((0..source.length()).into()).into_owned());
34                match self.build_root(green_tree.clone(), &source_text) {
35                    Ok(ast_root) => OakDiagnostics { result: Ok(ast_root), diagnostics: parse_result.diagnostics },
36                    Err(build_error) => {
37                        let mut diagnostics = parse_result.diagnostics;
38                        diagnostics.push(build_error.clone());
39                        OakDiagnostics { result: Err(build_error), diagnostics }
40                    }
41                }
42            }
43            Err(parse_error) => OakDiagnostics { result: Err(parse_error), diagnostics: parse_result.diagnostics },
44        }
45    }
46}
47
48impl<'config> HaskellBuilder<'config> {
49    /// Builds the root node.
50    pub(crate) fn build_root(&self, green_tree: GreenNode<HaskellLanguage>, source: &SourceText) -> Result<HaskellRoot, OakError> {
51        let red_root = RedNode::new(&green_tree, 0);
52        let mut module_name = None;
53        let mut items = Vec::new();
54
55        for child in red_root.children() {
56            if let RedTree::Node(node) = child {
57                match node.green.kind {
58                    HaskellElementType::ModuleDeclaration => {
59                        module_name = self.extract_module_name(&node, source);
60                    }
61                    _ => {
62                        if let Some(item) = self.build_item(&node, source) {
63                            items.push(item);
64                        }
65                    }
66                }
67            }
68        }
69
70        Ok(HaskellRoot { module_name, items })
71    }
72
73    fn build_item(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<Item> {
74        match node.green.kind {
75            HaskellElementType::Function => self.build_function(node, source).map(Item::Function),
76            HaskellElementType::DataDeclaration => self.build_data(node, source).map(Item::DataDeclaration),
77            HaskellElementType::TypeAliasDeclaration => self.build_type_alias(node, source).map(Item::TypeAlias),
78            HaskellElementType::ImportDeclaration => self.build_import(node, source).map(Item::Import),
79            _ => None,
80        }
81    }
82
83    fn build_function(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<Function> {
84        let mut name = None;
85        let mut type_signature = None;
86        let mut equations = Vec::new();
87
88        for child in node.children() {
89            if let RedTree::Node(n) = child {
90                match n.green.kind {
91                    HaskellElementType::Identifier | HaskellElementType::Constructor => {
92                        if name.is_none() {
93                            name = Some(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() });
94                        }
95                    }
96                    HaskellElementType::TypeSignature => {
97                        type_signature = self.build_type_signature(&n, source);
98                    }
99                    HaskellElementType::Equation => {
100                        if let Some(equation) = self.build_equation(&n, source) {
101                            equations.push(equation);
102                        }
103                    }
104                    _ => {}
105                }
106            }
107        }
108
109        name.map(|name| Function { name, type_signature, equations, span: node.span() })
110    }
111
112    fn build_type_signature(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<Type> {
113        for child in node.children() {
114            if let RedTree::Node(n) = child {
115                if n.green.kind == HaskellElementType::Type {
116                    return self.build_type(&n, source);
117                }
118            }
119        }
120        None
121    }
122
123    fn build_type(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<Type> {
124        let mut types = Vec::new();
125        let mut is_function = false;
126
127        for child in node.children() {
128            match child {
129                RedTree::Node(n) => match n.green.kind {
130                    HaskellElementType::Identifier => {
131                        types.push(Type::Variable(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() }));
132                    }
133                    HaskellElementType::Constructor => {
134                        types.push(Type::Constructor(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() }, vec![]));
135                    }
136                    HaskellElementType::Type => {
137                        if let Some(t) = self.build_type(&n, source) {
138                            types.push(t);
139                        }
140                    }
141                    _ => {}
142                },
143                RedTree::Leaf(l) => {
144                    if l.kind == HaskellTokenType::Arrow {
145                        is_function = true;
146                    }
147                }
148            }
149        }
150
151        if is_function && types.len() >= 2 {
152            let mut result = types.pop().unwrap();
153            while let Some(t) = types.pop() {
154                result = Type::Function(Box::new(t), Box::new(result));
155            }
156            Some(result)
157        }
158        else if let Some(first) = types.into_iter().next() {
159            Some(first)
160        }
161        else {
162            None
163        }
164    }
165
166    fn build_data(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<DataDeclaration> {
167        let mut name = None;
168        let mut type_params = Vec::new();
169        let mut constructors = Vec::new();
170
171        let mut found_assign = false;
172
173        for child in node.children() {
174            if let RedTree::Node(n) = child {
175                match n.green.kind {
176                    HaskellElementType::Constructor if !found_assign => {
177                        name = Some(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() });
178                    }
179                    HaskellElementType::Identifier if !found_assign => {
180                        type_params.push(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() });
181                    }
182                    HaskellElementType::Constructor if found_assign => {
183                        if let Some(cons) = self.build_constructor_def(&n, source) {
184                            constructors.push(cons);
185                        }
186                    }
187                    _ => {}
188                }
189            }
190            else if let RedTree::Leaf(l) = child {
191                if l.kind == HaskellTokenType::Assign {
192                    found_assign = true;
193                }
194            }
195        }
196
197        name.map(|name| DataDeclaration { name, type_params, constructors })
198    }
199
200    fn build_constructor_def(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<ConstructorDef> {
201        let mut name = None;
202        let mut fields = Vec::new();
203
204        for child in node.children() {
205            if let RedTree::Node(n) = child {
206                match n.green.kind {
207                    HaskellElementType::Constructor if name.is_none() => {
208                        name = Some(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() });
209                    }
210                    HaskellElementType::Type => {
211                        if let Some(t) = self.build_type(&n, source) {
212                            fields.push(t);
213                        }
214                    }
215                    _ => {}
216                }
217            }
218        }
219
220        name.map(|name| ConstructorDef { name, fields })
221    }
222
223    fn build_type_alias(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<TypeAlias> {
224        let mut name = None;
225        let mut type_params = Vec::new();
226        let mut target = None;
227
228        let mut found_assign = false;
229
230        for child in node.children() {
231            if let RedTree::Node(n) = child {
232                match n.green.kind {
233                    HaskellElementType::Identifier if !found_assign => {
234                        if name.is_none() {
235                            name = Some(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() });
236                        }
237                        else {
238                            type_params.push(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() });
239                        }
240                    }
241                    HaskellElementType::Type if found_assign => {
242                        target = self.build_type(&n, source);
243                    }
244                    _ => {}
245                }
246            }
247            else if let RedTree::Leaf(l) = child {
248                if l.kind == HaskellTokenType::Assign {
249                    found_assign = true;
250                }
251            }
252        }
253
254        if let (Some(name), Some(target)) = (name, target) { Some(TypeAlias { name, type_params, target }) } else { None }
255    }
256
257    fn build_import(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<Import> {
258        let mut module = None;
259        let mut qualified = false;
260        let mut as_name = None;
261
262        let mut found_as = false;
263
264        for child in node.children() {
265            if let RedTree::Node(n) = child {
266                match n.green.kind {
267                    HaskellElementType::Constructor => {
268                        if found_as {
269                            as_name = Some(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() });
270                        }
271                        else {
272                            module = Some(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() });
273                        }
274                    }
275                    _ => {}
276                }
277            }
278            else if let RedTree::Leaf(l) = child {
279                match l.kind {
280                    HaskellTokenType::Qualified => qualified = true,
281                    HaskellTokenType::As => found_as = true,
282                    _ => {}
283                }
284            }
285        }
286
287        module.map(|module| Import { module, qualified, as_name })
288    }
289
290    fn build_equation(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<Equation> {
291        let mut patterns = Vec::new();
292        let mut body = None;
293
294        for child in node.children() {
295            if let RedTree::Node(n) = child {
296                match n.green.kind {
297                    HaskellElementType::Pattern => {
298                        if let Some(pattern) = self.build_pattern(&n, source) {
299                            patterns.push(pattern);
300                        }
301                    }
302                    _ if n.green.kind as u16 >= HaskellElementType::LiteralExpression as u16 && n.green.kind as u16 <= HaskellElementType::CaseExpression as u16 => {
303                        body = self.build_expression(&n, source);
304                    }
305                    _ => {}
306                }
307            }
308        }
309
310        body.map(|body| Equation { patterns, body })
311    }
312
313    fn build_expression(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<Expression> {
314        match node.green.kind {
315            HaskellElementType::IdentifierExpression => {
316                for child in node.children() {
317                    if let RedTree::Node(n) = child {
318                        if n.green.kind == HaskellElementType::Identifier {
319                            return Some(Expression::Variable(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() }));
320                        }
321                    }
322                }
323                None
324            }
325            HaskellElementType::LiteralExpression => {
326                for child in node.children() {
327                    if let RedTree::Node(n) = child {
328                        if let Some(lit) = self.build_literal(&n, source) {
329                            return Some(Expression::Literal(lit));
330                        }
331                    }
332                }
333                None
334            }
335            HaskellElementType::ApplicationExpression => {
336                let mut exprs = Vec::new();
337                for child in node.children() {
338                    if let RedTree::Node(n) = child {
339                        if let Some(expr) = self.build_expression(&n, source) {
340                            exprs.push(expr);
341                        }
342                    }
343                }
344                if exprs.len() >= 2 {
345                    let mut iter = exprs.into_iter();
346                    let mut result = iter.next().unwrap();
347                    for next in iter {
348                        result = Expression::Application(Box::new(result), Box::new(next));
349                    }
350                    Some(result)
351                }
352                else {
353                    exprs.into_iter().next()
354                }
355            }
356            HaskellElementType::LambdaExpression => {
357                let mut patterns = Vec::new();
358                let mut body = None;
359                for child in node.children() {
360                    if let RedTree::Node(n) = child {
361                        match n.green.kind {
362                            HaskellElementType::Pattern => {
363                                if let Some(p) = self.build_pattern(&n, source) {
364                                    patterns.push(p);
365                                }
366                            }
367                            _ if n.green.kind as u16 >= HaskellElementType::LiteralExpression as u16 && n.green.kind as u16 <= HaskellElementType::CaseExpression as u16 => {
368                                body = self.build_expression(&n, source);
369                            }
370                            _ => {}
371                        }
372                    }
373                }
374                body.map(|body| Expression::Lambda(patterns, Box::new(body)))
375            }
376            HaskellElementType::LetExpression => {
377                let mut items = Vec::new();
378                let mut body = None;
379                for child in node.children() {
380                    if let RedTree::Node(n) = child {
381                        match n.green.kind {
382                            _ if n.green.kind as u16 >= HaskellElementType::Function as u16 && n.green.kind as u16 <= HaskellElementType::TypeAliasDeclaration as u16 => {
383                                if let Some(item) = self.build_item(&n, source) {
384                                    items.push(item);
385                                }
386                            }
387                            _ if n.green.kind as u16 >= HaskellElementType::LiteralExpression as u16 && n.green.kind as u16 <= HaskellElementType::CaseExpression as u16 => {
388                                body = self.build_expression(&n, source);
389                            }
390                            _ => {}
391                        }
392                    }
393                }
394                body.map(|body| Expression::Let(items, Box::new(body)))
395            }
396            HaskellElementType::CaseExpression => {
397                let mut target = None;
398                let mut arms = Vec::new();
399                for child in node.children() {
400                    if let RedTree::Node(n) = child {
401                        match n.green.kind {
402                            HaskellElementType::CaseArm => {
403                                if let Some(arm) = self.build_case_arm(&n, source) {
404                                    arms.push(arm);
405                                }
406                            }
407                            _ if n.green.kind as u16 >= HaskellElementType::LiteralExpression as u16 && n.green.kind as u16 <= HaskellElementType::CaseExpression as u16 => {
408                                if target.is_none() {
409                                    target = self.build_expression(&n, source);
410                                }
411                            }
412                            _ => {}
413                        }
414                    }
415                }
416                target.map(|target| Expression::Case(Box::new(target), arms))
417            }
418            _ => None,
419        }
420    }
421
422    fn build_case_arm(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<CaseArm> {
423        let mut pattern = None;
424        let mut body = None;
425
426        for child in node.children() {
427            if let RedTree::Node(n) = child {
428                match n.green.kind {
429                    HaskellElementType::Pattern => {
430                        pattern = self.build_pattern(&n, source);
431                    }
432                    _ if n.green.kind as u16 >= HaskellElementType::LiteralExpression as u16 && n.green.kind as u16 <= HaskellElementType::CaseExpression as u16 => {
433                        body = self.build_expression(&n, source);
434                    }
435                    _ => {}
436                }
437            }
438        }
439
440        if let (Some(pattern), Some(body)) = (pattern, body) { Some(CaseArm { pattern, body }) } else { None }
441    }
442
443    fn build_literal(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<Literal> {
444        for child in node.children() {
445            if let RedTree::Leaf(l) = child {
446                let text = source.get_text_in(l.span());
447                match l.kind {
448                    HaskellTokenType::Integer => {
449                        if let Ok(val) = text.parse::<i64>() {
450                            return Some(Literal::Integer(val));
451                        }
452                    }
453                    HaskellTokenType::Float => {
454                        if let Ok(val) = text.parse::<f64>() {
455                            return Some(Literal::Float(val));
456                        }
457                    }
458                    HaskellTokenType::StringLiteral => {
459                        return Some(Literal::String(text[1..text.len() - 1].to_string()));
460                    }
461                    HaskellTokenType::CharLiteral => {
462                        if let Some(c) = text.chars().nth(1) {
463                            return Some(Literal::Char(c));
464                        }
465                    }
466                    _ => {}
467                }
468            }
469        }
470        None
471    }
472
473    fn build_pattern(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<Pattern> {
474        for child in node.children() {
475            match child {
476                RedTree::Leaf(leaf) => {
477                    if leaf.kind == HaskellTokenType::Underscore {
478                        return Some(Pattern::Wildcard);
479                    }
480                }
481                RedTree::Node(n) => match n.green.kind {
482                    HaskellElementType::Identifier => {
483                        return Some(Pattern::Variable(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() }));
484                    }
485                    HaskellElementType::Constructor => {
486                        let mut name = None;
487                        let mut args = Vec::new();
488                        for c in n.children() {
489                            if let RedTree::Node(cn) = c {
490                                if cn.green.kind == HaskellElementType::Constructor && name.is_none() {
491                                    name = Some(Identifier { name: source.get_text_in(cn.span()).to_string(), span: cn.span() });
492                                }
493                                else if cn.green.kind == HaskellElementType::Pattern {
494                                    if let Some(p) = self.build_pattern(&cn, source) {
495                                        args.push(p);
496                                    }
497                                }
498                            }
499                        }
500                        if let Some(name) = name {
501                            return Some(Pattern::Constructor(name, args));
502                        }
503                    }
504                    _ if n.green.kind == HaskellElementType::LiteralExpression => {
505                        if let Some(expr) = self.build_expression(&n, source) {
506                            if let Expression::Literal(lit) = expr {
507                                return Some(Pattern::Literal(lit));
508                            }
509                        }
510                    }
511                    _ => {}
512                },
513            }
514        }
515        None
516    }
517
518    fn extract_module_name(&self, node: &RedNode<HaskellLanguage>, source: &SourceText) -> Option<Identifier> {
519        for child in node.children() {
520            if let RedTree::Node(n) = child {
521                if n.green.kind == HaskellElementType::Constructor {
522                    return Some(Identifier { name: source.get_text_in(n.span()).to_string(), span: n.span() });
523                }
524            }
525        }
526        None
527    }
528}