oak_mojo/builder/
mod.rs

1use crate::{
2    MojoLanguage,
3    ast::{MojoExpression, MojoLiteral, MojoStatement},
4    parser::MojoElementType,
5};
6use oak_core::{GreenNode, GreenTree, OakError, Source, source::SourceText};
7
8/// Mojo AST builder.
9pub struct MojoBuilder<'a> {
10    source: &'a SourceText,
11}
12
13impl<'a> MojoBuilder<'a> {
14    fn build_param_list(&self, node: &GreenNode<MojoLanguage>, offset: usize) -> Result<Vec<(String, Option<String>)>, OakError> {
15        let mut params = Vec::new();
16        let mut current_offset = offset;
17        let mut current_name = String::new();
18
19        for child in node.children() {
20            if let GreenTree::Node(child_node) = child {
21                match child_node.kind {
22                    MojoElementType::Identifier => {
23                        let text = self.source.get_text_in((current_offset..current_offset + child.len() as usize).into()).to_string();
24                        if current_name.is_empty() {
25                            current_name = text;
26                        }
27                        else {
28                            params.push((current_name.clone(), Some(text)));
29                            current_name.clear();
30                        }
31                    }
32                    MojoElementType::Comma => {
33                        if !current_name.is_empty() {
34                            params.push((current_name.clone(), None));
35                            current_name.clear();
36                        }
37                    }
38                    _ => {}
39                }
40            }
41            current_offset += child.len() as usize;
42        }
43
44        if !current_name.is_empty() {
45            params.push((current_name, None));
46        }
47
48        Ok(params)
49    }
50
51    /// Creates a new builder
52    pub fn new(source: &'a SourceText) -> Self {
53        Self { source }
54    }
55
56    /// Builds AST from GreenNode
57    pub fn build_root(&self, green: &GreenNode<MojoLanguage>) -> Result<Vec<MojoStatement>, OakError> {
58        let mut statements = Vec::new();
59        let mut offset = 0;
60
61        for child in green.children() {
62            if let GreenTree::Node(node) = child {
63                if !node.kind.is_trivia() {
64                    if let Some(stmt) = self.build_statement(node, offset)? {
65                        statements.push(stmt);
66                    }
67                }
68            }
69            offset += child.len() as usize;
70        }
71
72        Ok(statements)
73    }
74
75    fn build_statement(&self, node: &GreenNode<MojoLanguage>, offset: usize) -> Result<Option<MojoStatement>, OakError> {
76        match node.kind {
77            MojoElementType::FunctionDef => self.build_function_def(node, offset).map(Some),
78            MojoElementType::VariableDecl => self.build_variable_decl(node, offset).map(Some),
79            MojoElementType::IfStatement => self.build_if_stmt(node, offset).map(Some),
80            MojoElementType::WhileStatement => self.build_while_stmt(node, offset).map(Some),
81            MojoElementType::ReturnStatement => self.build_return_stmt(node, offset).map(Some),
82            MojoElementType::ExpressionStatement => self.build_expression_stmt(node, offset).map(Some),
83            _ => Ok(None),
84        }
85    }
86
87    fn build_function_def(&self, node: &GreenNode<MojoLanguage>, offset: usize) -> Result<MojoStatement, OakError> {
88        let mut name = String::new();
89        let mut params = Vec::new();
90        let mut return_type = None;
91        let mut body = Vec::new();
92        let mut current_offset = offset;
93
94        for child in node.children() {
95            if let GreenTree::Node(child_node) = child {
96                match child_node.kind {
97                    MojoElementType::Identifier => {
98                        let text = self.source.get_text_in((current_offset..current_offset + child.len() as usize).into()).to_string();
99                        if name.is_empty() {
100                            name = text;
101                        }
102                        else {
103                            return_type = Some(text);
104                        }
105                    }
106                    MojoElementType::ParamList => {
107                        params = self.build_param_list(child_node, current_offset)?;
108                    }
109                    MojoElementType::Block => {
110                        body = self.build_root(child_node)?;
111                    }
112                    _ => {}
113                }
114            }
115            current_offset += child.len() as usize;
116        }
117
118        Ok(MojoStatement::Function { name, params, return_type, body })
119    }
120
121    fn build_variable_decl(&self, node: &GreenNode<MojoLanguage>, offset: usize) -> Result<MojoStatement, OakError> {
122        let mut name = String::new();
123        let mut ty = None;
124        let mut value = None;
125        let mut is_let = false;
126        let mut current_offset = offset;
127
128        for child in node.children() {
129            if let GreenTree::Node(child_node) = child {
130                match child_node.kind {
131                    MojoElementType::Var => is_let = false,
132                    MojoElementType::Let => is_let = true,
133                    MojoElementType::Identifier => {
134                        name = self.source.get_text_in((current_offset..current_offset + child.len() as usize).into()).to_string();
135                    }
136                    MojoElementType::BinaryExpr | MojoElementType::LiteralExpr | MojoElementType::IdentifierExpr | MojoElementType::Grouping => {
137                        value = Some(self.build_expression(child_node, current_offset)?);
138                    }
139                    _ => {}
140                }
141            }
142            current_offset += child.len() as usize;
143        }
144
145        Ok(MojoStatement::Variable { name, ty, value, is_let })
146    }
147
148    fn build_if_stmt(&self, node: &GreenNode<MojoLanguage>, offset: usize) -> Result<MojoStatement, OakError> {
149        let mut condition = MojoExpression::Literal(MojoLiteral::None);
150        let mut then_body = Vec::new();
151        let mut else_body = None;
152        let mut current_offset = offset;
153
154        for child in node.children() {
155            if let GreenTree::Node(child_node) = child {
156                match child_node.kind {
157                    MojoElementType::BinaryExpr | MojoElementType::LiteralExpr | MojoElementType::IdentifierExpr | MojoElementType::Grouping => {
158                        condition = self.build_expression(child_node, current_offset)?;
159                    }
160                    MojoElementType::Block => {
161                        if then_body.is_empty() {
162                            then_body = self.build_root(child_node)?;
163                        }
164                        else {
165                            else_body = Some(self.build_root(child_node)?);
166                        }
167                    }
168                    MojoElementType::IfStatement => {
169                        // Elif case, but Mojo usually handles this via nested Ifs in Block
170                        // or we could transform it here.
171                    }
172                    _ => {}
173                }
174            }
175            current_offset += child.len() as usize;
176        }
177
178        Ok(MojoStatement::If { condition, then_body, else_body })
179    }
180
181    fn build_while_stmt(&self, node: &GreenNode<MojoLanguage>, offset: usize) -> Result<MojoStatement, OakError> {
182        let mut condition = MojoExpression::Literal(MojoLiteral::None);
183        let mut body = Vec::new();
184        let mut current_offset = offset;
185
186        for child in node.children() {
187            if let GreenTree::Node(child_node) = child {
188                match child_node.kind {
189                    MojoElementType::BinaryExpr | MojoElementType::LiteralExpr | MojoElementType::IdentifierExpr | MojoElementType::Grouping => {
190                        condition = self.build_expression(child_node, current_offset)?;
191                    }
192                    MojoElementType::Block => {
193                        body = self.build_root(child_node)?;
194                    }
195                    _ => {}
196                }
197            }
198            current_offset += child.len() as usize;
199        }
200
201        Ok(MojoStatement::While { condition, body })
202    }
203
204    fn build_return_stmt(&self, node: &GreenNode<MojoLanguage>, offset: usize) -> Result<MojoStatement, OakError> {
205        let mut value = None;
206        let mut current_offset = offset;
207
208        for child in node.children() {
209            if let GreenTree::Node(child_node) = child {
210                if !child_node.kind.is_trivia() && child_node.kind != MojoElementType::Return {
211                    value = Some(self.build_expression(child_node, current_offset)?);
212                }
213            }
214            current_offset += child.len() as usize;
215        }
216
217        Ok(MojoStatement::Return(value))
218    }
219
220    fn build_expression_stmt(&self, node: &GreenNode<MojoLanguage>, offset: usize) -> Result<MojoStatement, OakError> {
221        let mut expr = MojoExpression::Literal(MojoLiteral::None);
222        let mut current_offset = offset;
223
224        for child in node.children() {
225            if let GreenTree::Node(child_node) = child {
226                if !child_node.kind.is_trivia() {
227                    expr = self.build_expression(child_node, current_offset)?;
228                }
229            }
230            current_offset += child.len() as usize;
231        }
232
233        Ok(MojoStatement::Expression(expr))
234    }
235
236    fn build_expression(&self, node: &GreenNode<MojoLanguage>, offset: usize) -> Result<MojoExpression, OakError> {
237        match node.kind {
238            MojoElementType::LiteralExpr => {
239                let text = self.source.get_text_in((offset..offset + node.text_len() as usize).into());
240                if text.contains('.') {
241                    Ok(MojoExpression::Literal(MojoLiteral::Float(text.parse().unwrap_or(0.0))))
242                }
243                else if text.starts_with('"') || text.starts_with('\'') {
244                    Ok(MojoExpression::Literal(MojoLiteral::String(text[1..text.len() - 1].to_string())))
245                }
246                else if text == "True" {
247                    Ok(MojoExpression::Literal(MojoLiteral::Bool(true)))
248                }
249                else if text == "False" {
250                    Ok(MojoExpression::Literal(MojoLiteral::Bool(false)))
251                }
252                else if text == "None" {
253                    Ok(MojoExpression::Literal(MojoLiteral::None))
254                }
255                else {
256                    Ok(MojoExpression::Literal(MojoLiteral::Int(text.parse().unwrap_or(0))))
257                }
258            }
259            MojoElementType::IdentifierExpr => {
260                let text = self.source.get_text_in((offset..offset + node.text_len() as usize).into());
261                Ok(MojoExpression::Identifier(text.to_string()))
262            }
263            MojoElementType::BinaryExpr => {
264                let mut left = None;
265                let mut op = String::new();
266                let mut right = None;
267                let mut current_offset = offset;
268
269                for child in node.children() {
270                    if let GreenTree::Node(child_node) = child {
271                        if !child_node.kind.is_trivia() {
272                            if left.is_none() {
273                                left = Some(Box::new(self.build_expression(child_node, current_offset)?));
274                            }
275                            else if right.is_none()
276                                && (child_node.kind == MojoElementType::BinaryExpr || child_node.kind == MojoElementType::LiteralExpr || child_node.kind == MojoElementType::IdentifierExpr || child_node.kind == MojoElementType::Grouping)
277                            {
278                                right = Some(Box::new(self.build_expression(child_node, current_offset)?));
279                            }
280                            else {
281                                // Assume it's the operator
282                                op = self.source.get_text_in((current_offset..current_offset + child.len() as usize).into()).to_string();
283                            }
284                        }
285                    }
286                    current_offset += child.len() as usize;
287                }
288
289                Ok(MojoExpression::Binary { left: left.unwrap_or_else(|| Box::new(MojoExpression::Literal(MojoLiteral::None))), op, right: right.unwrap_or_else(|| Box::new(MojoExpression::Literal(MojoLiteral::None))) })
290            }
291            MojoElementType::Grouping => {
292                let mut expr = MojoExpression::Literal(MojoLiteral::None);
293                let mut current_offset = offset;
294                for child in node.children() {
295                    if let GreenTree::Node(child_node) = child {
296                        if !child_node.kind.is_trivia() && child_node.kind != MojoElementType::LeftParen && child_node.kind != MojoElementType::RightParen {
297                            expr = self.build_expression(child_node, current_offset)?;
298                        }
299                    }
300                    current_offset += child.len() as usize;
301                }
302                Ok(expr)
303            }
304            _ => Ok(MojoExpression::Literal(MojoLiteral::None)),
305        }
306    }
307}
oak_mojo/builder/mod.rs

oak_mojo/builder/
mod.rs
