Skip to main content

oak_swift/builder/
mod.rs

1use crate::{
2    ast::{Expression, Literal, Parameter, Program, Statement, SwiftRoot, Type},
3    language::SwiftLanguage,
4    lexer::token_type::SwiftTokenType,
5    parser::{SwiftParser, element_type::SwiftElementType},
6};
7use oak_core::{Builder, BuilderCache, GreenNode, GreenTree, OakDiagnostics, OakError, Parser, SourceText, TextEdit, builder::BuildOutput, source::Source};
8
9/// AST builder for Swift language that constructs typed AST nodes from green tree.
10#[derive(Clone)]
11pub struct SwiftBuilder<'config> {
12    config: &'config SwiftLanguage,
13}
14
15impl<'config> SwiftBuilder<'config> {
16    /// Creates a new SwiftBuilder with the given language configuration.
17    pub fn new(config: &'config SwiftLanguage) -> Self {
18        Self { config }
19    }
20}
21
22impl<'config> Builder<SwiftLanguage> for SwiftBuilder<'config> {
23    fn build<'a, S: Source + ?Sized>(&self, source: &S, edits: &[TextEdit], _cache: &'a mut impl BuilderCache<SwiftLanguage>) -> BuildOutput<SwiftLanguage> {
24        let parser = SwiftParser::new(self.config);
25        let mut parse_cache = oak_core::parser::session::ParseSession::<SwiftLanguage>::default();
26        let parse_result = parser.parse(source, edits, &mut parse_cache);
27
28        match parse_result.result {
29            Ok(green_tree) => {
30                let source_text = SourceText::new(source.get_text_in((0..source.length()).into()).into_owned());
31                match self.build_root(green_tree, &source_text) {
32                    Ok(ast_root) => OakDiagnostics { result: Ok(ast_root), diagnostics: parse_result.diagnostics },
33                    Err(build_error) => {
34                        let mut diagnostics = parse_result.diagnostics;
35                        diagnostics.push(build_error.clone());
36                        OakDiagnostics { result: Err(build_error), diagnostics }
37                    }
38                }
39            }
40            Err(e) => OakDiagnostics { result: Err(e), diagnostics: parse_result.diagnostics },
41        }
42    }
43}
44
45impl<'config> SwiftBuilder<'config> {
46    /// Builds a SwiftRoot AST node from a green tree and source text.
47    pub fn build_root(&self, green_tree: &GreenNode<SwiftLanguage>, source: &SourceText) -> Result<SwiftRoot, OakError> {
48        let mut statements = Vec::new();
49        let mut current_offset = 0;
50
51        for child in green_tree.children() {
52            let child_len = child.len() as usize;
53            match child {
54                GreenTree::Node(node) => {
55                    if !node.kind.is_trivia() {
56                        if let Some(stmt) = self.build_statement(node, current_offset, source)? {
57                            statements.push(stmt);
58                        }
59                    }
60                }
61                _ => {}
62            }
63            current_offset += child_len;
64        }
65
66        Ok(SwiftRoot { program: Program { statements }, span: (0..green_tree.text_len() as usize).into() })
67    }
68
69    fn build_statement(&self, node: &GreenNode<SwiftLanguage>, offset: usize, source: &SourceText) -> Result<Option<Statement>, OakError> {
70        // eprintln!("Building statement: {:?} at {}", node.kind, offset);
71        match node.kind {
72            SwiftElementType::FunctionDeclaration => {
73                let mut name = String::new();
74                let mut parameters = Vec::new();
75                let mut body = Vec::new();
76                let mut return_type = None;
77                let mut current_offset = offset;
78                let mut found_arrow = false;
79
80                for child in node.children() {
81                    let child_len = child.len() as usize;
82                    match child {
83                        GreenTree::Leaf(leaf) => {
84                            if leaf.kind == SwiftTokenType::Identifier {
85                                if name.is_empty() {
86                                    name = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
87                                    eprintln!("Found identifier: '{}' at {}", name, current_offset);
88                                }
89                                else if found_arrow && return_type.is_none() {
90                                    let t = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
91                                    return_type = Some(Type { name: t });
92                                }
93                            }
94                            else if leaf.kind == SwiftTokenType::Arrow {
95                                found_arrow = true;
96                            }
97                        }
98                        GreenTree::Node(n) if n.kind == SwiftElementType::ParameterList => {
99                            parameters = self.build_parameters(n, current_offset, source)?;
100                        }
101                        GreenTree::Node(n) if n.kind == SwiftElementType::Block => {
102                            body = self.build_block(n, current_offset, source)?;
103                        }
104                        _ => {}
105                    }
106                    current_offset += child_len;
107                }
108                Ok(Some(Statement::FunctionDef { name, parameters, return_type, body }))
109            }
110            SwiftElementType::VariableDeclaration => {
111                let mut name = String::new();
112                let mut value = None;
113                let mut is_mutable = true;
114                let mut type_annotation = None;
115                let mut current_offset = offset;
116                let mut found_colon = false;
117
118                for child in node.children() {
119                    let child_len = child.len() as usize;
120                    match child {
121                        GreenTree::Leaf(leaf) => match leaf.kind {
122                            SwiftTokenType::Let => is_mutable = false,
123                            SwiftTokenType::Var => is_mutable = true,
124                            SwiftTokenType::Identifier => {
125                                if name.is_empty() {
126                                    name = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
127                                }
128                                else if found_colon && type_annotation.is_none() {
129                                    let t = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
130                                    type_annotation = Some(Type { name: t });
131                                }
132                            }
133                            SwiftTokenType::Colon => found_colon = true,
134                            _ => {}
135                        },
136                        GreenTree::Node(n) => {
137                            if !n.kind.is_trivia() && n.kind != SwiftElementType::Identifier && n.kind != SwiftElementType::Colon {
138                                value = Some(self.build_expression(n, current_offset, source)?);
139                            }
140                        }
141                    }
142                    current_offset += child_len;
143                }
144                Ok(Some(Statement::VariableDecl { is_mutable, name, type_annotation, value }))
145            }
146            SwiftElementType::ReturnStatement => {
147                let mut value = None;
148                let mut current_offset = offset;
149                for child in node.children() {
150                    let child_len = child.len() as usize;
151                    if let GreenTree::Node(n) = child {
152                        if !n.kind.is_trivia() {
153                            value = Some(self.build_expression(n, current_offset, source)?);
154                        }
155                    }
156                    current_offset += child_len;
157                }
158                Ok(Some(Statement::Return(value)))
159            }
160            SwiftElementType::IfStatement => {
161                let mut test = None;
162                let mut body = Vec::new();
163                let mut orelse = None;
164                let mut current_offset = offset;
165
166                for child in node.children() {
167                    let child_len = child.len() as usize;
168                    match child {
169                        GreenTree::Node(n) if n.kind == SwiftElementType::Block => {
170                            if body.is_empty() {
171                                body = self.build_block(n, current_offset, source)?;
172                            }
173                            else {
174                                orelse = Some(self.build_block(n, current_offset, source)?);
175                            }
176                        }
177                        GreenTree::Node(n) if n.kind == SwiftElementType::IfStatement => {
178                            orelse = Some(vec![self.build_statement(n, current_offset, source)?.unwrap()]);
179                        }
180                        GreenTree::Node(n) if !n.kind.is_trivia() => {
181                            test = Some(self.build_expression(n, current_offset, source)?);
182                        }
183                        _ => {}
184                    }
185                    current_offset += child_len;
186                }
187                Ok(Some(Statement::If { test: test.unwrap(), body, orelse }))
188            }
189            SwiftElementType::WhileStatement => {
190                let mut test = None;
191                let mut body = Vec::new();
192                let mut current_offset = offset;
193
194                for child in node.children() {
195                    let child_len = child.len() as usize;
196                    match child {
197                        GreenTree::Node(n) if n.kind == SwiftElementType::Block => {
198                            body = self.build_block(n, current_offset, source)?;
199                        }
200                        GreenTree::Node(n) if !n.kind.is_trivia() => {
201                            test = Some(self.build_expression(n, current_offset, source)?);
202                        }
203                        _ => {}
204                    }
205                    current_offset += child_len;
206                }
207                Ok(Some(Statement::While { test: test.unwrap(), body }))
208            }
209            SwiftElementType::ForStatement => {
210                let mut variable = String::new();
211                let mut iterable = None;
212                let mut body = Vec::new();
213                let mut current_offset = offset;
214
215                for child in node.children() {
216                    let child_len = child.len() as usize;
217                    match child {
218                        GreenTree::Leaf(leaf) if leaf.kind == SwiftTokenType::Identifier => {
219                            variable = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
220                        }
221                        GreenTree::Node(n) if n.kind == SwiftElementType::Block => {
222                            body = self.build_block(n, current_offset, source)?;
223                        }
224                        GreenTree::Node(n) if !n.kind.is_trivia() => {
225                            iterable = Some(self.build_expression(n, current_offset, source)?);
226                        }
227                        _ => {}
228                    }
229                    current_offset += child_len;
230                }
231                Ok(Some(Statement::For { variable, iterable: iterable.unwrap(), body }))
232            }
233            SwiftElementType::ExpressionStatement => {
234                let mut expr = None;
235                let mut current_offset = offset;
236                for child in node.children() {
237                    let child_len = child.len() as usize;
238                    if let GreenTree::Node(n) = child {
239                        if !n.kind.is_trivia() {
240                            expr = Some(self.build_expression(n, current_offset, source)?);
241                        }
242                    }
243                    current_offset += child_len;
244                }
245                Ok(expr.map(Statement::Expression))
246            }
247            SwiftElementType::Block => Ok(Some(Statement::Block(self.build_block(node, offset, source)?))),
248            _ => Ok(None),
249        }
250    }
251
252    fn build_block(&self, node: &GreenNode<SwiftLanguage>, offset: usize, source: &SourceText) -> Result<Vec<Statement>, OakError> {
253        let mut statements = Vec::new();
254        let mut current_offset = offset;
255
256        for child in node.children() {
257            let child_len = child.len() as usize;
258            if let GreenTree::Node(n) = child {
259                if !n.kind.is_trivia() {
260                    if let Some(stmt) = self.build_statement(n, current_offset, source)? {
261                        statements.push(stmt);
262                    }
263                }
264            }
265            current_offset += child_len;
266        }
267        Ok(statements)
268    }
269
270    fn build_parameters(&self, node: &GreenNode<SwiftLanguage>, offset: usize, source: &SourceText) -> Result<Vec<Parameter>, OakError> {
271        let mut parameters = Vec::new();
272        let mut current_offset = offset;
273
274        for child in node.children() {
275            let child_len = child.len() as usize;
276            if let GreenTree::Node(n) = child {
277                if n.kind == SwiftElementType::Parameter {
278                    parameters.push(self.build_parameter(n, current_offset, source)?);
279                }
280            }
281            current_offset += child_len;
282        }
283        Ok(parameters)
284    }
285
286    fn build_parameter(&self, node: &GreenNode<SwiftLanguage>, offset: usize, source: &SourceText) -> Result<Parameter, OakError> {
287        let mut name = String::new();
288        let mut type_name = String::new();
289        let mut current_offset = offset;
290
291        for child in node.children() {
292            let child_len = child.len() as usize;
293            match child {
294                GreenTree::Leaf(leaf) if leaf.kind == SwiftTokenType::Identifier => {
295                    if name.is_empty() {
296                        name = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
297                        eprintln!("Found parameter name: '{}' at {}", name, current_offset);
298                    }
299                    else {
300                        type_name = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
301                        eprintln!("Found parameter type: '{}' at {}", type_name, current_offset);
302                    }
303                }
304                _ => {}
305            }
306            current_offset += child_len;
307        }
308
309        Ok(Parameter { name, type_annotation: Type { name: type_name } })
310    }
311
312    fn build_expression(&self, node: &GreenNode<SwiftLanguage>, offset: usize, source: &SourceText) -> Result<Expression, OakError> {
313        match node.kind {
314            SwiftElementType::BinaryExpression => {
315                let mut left = None;
316                let mut operator = String::new();
317                let mut right = None;
318                let mut current_offset = offset;
319
320                for child in node.children() {
321                    let child_len = child.len() as usize;
322                    match child {
323                        GreenTree::Node(n) => {
324                            if left.is_none() {
325                                left = Some(Box::new(self.build_expression(n, current_offset, source)?));
326                            }
327                            else {
328                                right = Some(Box::new(self.build_expression(n, current_offset, source)?));
329                            }
330                        }
331                        GreenTree::Leaf(leaf) if !leaf.kind.is_trivia() => {
332                            operator = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
333                        }
334                        _ => {}
335                    }
336                    current_offset += child_len;
337                }
338                Ok(Expression::Binary { left: left.unwrap(), operator, right: right.unwrap() })
339            }
340            SwiftElementType::UnaryExpression => {
341                let mut operator = String::new();
342                let mut operand = None;
343                let mut current_offset = offset;
344
345                for child in node.children() {
346                    let child_len = child.len() as usize;
347                    match child {
348                        GreenTree::Leaf(leaf) if !leaf.kind.is_trivia() => {
349                            operator = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
350                        }
351                        GreenTree::Node(n) => {
352                            operand = Some(Box::new(self.build_expression(n, current_offset, source)?));
353                        }
354                        _ => {}
355                    }
356                    current_offset += child_len;
357                }
358                Ok(Expression::Unary { operator, operand: operand.unwrap() })
359            }
360            SwiftElementType::IdentifierExpression => {
361                let text = source.get_text_in((offset..offset + node.text_len() as usize).into()).trim().to_string();
362                Ok(Expression::Identifier(text))
363            }
364            SwiftElementType::LiteralExpression => {
365                let text = source.get_text_in((offset..offset + node.text_len() as usize).into()).trim().to_string();
366                // Simple literal detection
367                if text == "true" {
368                    Ok(Expression::Literal(Literal::Boolean(true)))
369                }
370                else if text == "false" {
371                    Ok(Expression::Literal(Literal::Boolean(false)))
372                }
373                else if text == "nil" {
374                    Ok(Expression::Literal(Literal::Nil))
375                }
376                else if text.starts_with('"') {
377                    Ok(Expression::Literal(Literal::String(text[1..text.len() - 1].to_string())))
378                }
379                else {
380                    Ok(Expression::Literal(Literal::Number(text)))
381                }
382            }
383            SwiftElementType::CallExpression => {
384                let mut callee = None;
385                let mut arguments = Vec::new();
386                let mut current_offset = offset;
387
388                for child in node.children() {
389                    let child_len = child.len() as usize;
390                    match child {
391                        GreenTree::Node(n) if callee.is_none() => {
392                            callee = Some(Box::new(self.build_expression(n, current_offset, source)?));
393                        }
394                        GreenTree::Node(n) => {
395                            arguments.push(self.build_expression(n, current_offset, source)?);
396                        }
397                        _ => {}
398                    }
399                    current_offset += child_len;
400                }
401                Ok(Expression::Call { callee: callee.unwrap(), arguments })
402            }
403            SwiftElementType::MemberExpression => {
404                let mut object = None;
405                let mut member = String::new();
406                let mut current_offset = offset;
407
408                for child in node.children() {
409                    let child_len = child.len() as usize;
410                    match child {
411                        GreenTree::Node(n) if !n.kind.is_trivia() => {
412                            object = Some(Box::new(self.build_expression(n, current_offset, source)?));
413                        }
414                        GreenTree::Leaf(leaf) if leaf.kind == SwiftTokenType::Identifier => {
415                            member = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
416                        }
417                        _ => {}
418                    }
419                    current_offset += child_len;
420                }
421                Ok(Expression::Member { object: object.unwrap(), member })
422            }
423            _ => Err(OakError::custom_error(format!("Unsupported expression kind: {:?}", node.kind))),
424        }
425    }
426}