Skip to main content

oak_go/builder/
mod.rs

1//! Go language builder
2
3use crate::{
4    ast::{self, Declaration, GoRoot},
5    language::GoLanguage,
6    lexer::{GoLexer, token_type::GoTokenType},
7    parser::{GoParser, element_type::GoElementType},
8};
9use oak_core::{
10    Builder, BuilderCache, GreenNode, Lexer, OakDiagnostics, OakError, ParseSession, Parser, SourceText, TextEdit, TokenType,
11    builder::BuildOutput,
12    source::Source,
13    tree::{RedNode, RedTree},
14};
15
16mod build_signatures;
17
18/// Go language builder that constructs typed AST from parse trees.
19pub struct GoBuilder<'config> {
20    /// The language configuration reference.
21    pub(crate) config: &'config GoLanguage,
22}
23
24impl<'config> Builder<GoLanguage> for GoBuilder<'config> {
25    fn build<'a, S: Source + ?Sized>(&self, text: &S, edits: &[TextEdit], cache: &'a mut impl BuilderCache<GoLanguage>) -> BuildOutput<GoLanguage> {
26        let parser = GoParser::new(self.config);
27        let cache_ptr = cache as *mut _;
28        let output = parser.parse(text, edits, unsafe { &mut *cache_ptr });
29        let result = output.result.map(|green| {
30            let source = SourceText::new(text.get_text_in((0..text.length()).into()).into_owned());
31            self.build_root(green, &source, unsafe { &mut *cache_ptr })
32        });
33        OakDiagnostics { result, diagnostics: output.diagnostics }
34    }
35}
36
37impl<'config> GoBuilder<'config> {
38    /// Creates a new Go builder with the given language configuration.
39    pub fn new(config: &'config GoLanguage) -> Self {
40        Self { config }
41    }
42
43    fn build_root<'a>(&self, green: &'a GreenNode<'a, GoLanguage>, source: &SourceText, cache: &mut impl BuilderCache<GoLanguage>) -> GoRoot {
44        if let Some(cached) = cache.get_typed_node::<std::sync::Arc<GoRoot>>(green) {
45            return (*cached).clone();
46        }
47
48        let red = RedNode::new(green, 0);
49        let mut package = None;
50        let mut imports = vec![];
51        let mut declarations = vec![];
52
53        for child in red.children() {
54            if let RedTree::Node(node) = child {
55                match node.green.kind {
56                    GoElementType::PackageClause => {
57                        package = self.extract_package(node, source);
58                    }
59                    GoElementType::ImportDeclaration => {
60                        imports.extend(self.extract_imports(node, source));
61                    }
62                    GoElementType::FunctionDeclaration => {
63                        if let Ok(func) = self.extract_function(node, source, cache) {
64                            declarations.push(Declaration::Function(func));
65                        }
66                    }
67                    GoElementType::VariableDeclaration => {
68                        if let Ok(vars) = self.extract_variables(node, source, cache) {
69                            declarations.extend(vars);
70                        }
71                    }
72                    GoElementType::ConstDeclaration => {
73                        if let Ok(consts) = self.extract_consts(node, source, cache) {
74                            declarations.extend(consts);
75                        }
76                    }
77                    GoElementType::TypeDeclaration => {
78                        if let Ok(types) = self.extract_types(node, source, cache) {
79                            declarations.extend(types);
80                        }
81                    }
82                    _ => {}
83                }
84            }
85        }
86
87        let root = GoRoot { package, imports, declarations };
88        // cache.set_typed_node(green, std::sync::Arc::new(root.clone()));
89        root
90    }
91
92    fn extract_package(&self, node: RedNode<GoLanguage>, source: &SourceText) -> Option<String> {
93        for child in node.children() {
94            if let RedTree::Leaf(leaf) = child {
95                if leaf.kind == GoTokenType::Identifier {
96                    return Some(source.get_text_in(leaf.span).trim().to_string());
97                }
98            }
99        }
100        None
101    }
102
103    fn extract_imports(&self, node: RedNode<GoLanguage>, source: &SourceText) -> Vec<ast::Import> {
104        let mut imports = vec![];
105        for child in node.children() {
106            if let RedTree::Node(n) = child {
107                match n.green.kind {
108                    GoElementType::ImportSpec => {
109                        let mut path = String::new();
110                        let mut alias = None;
111                        for spec_child in n.children() {
112                            match spec_child {
113                                RedTree::Leaf(leaf) if leaf.kind == GoTokenType::Identifier => {
114                                    alias = Some(source.get_text_in(leaf.span).trim().to_string());
115                                }
116                                RedTree::Leaf(leaf) if leaf.kind == GoTokenType::StringLiteral => {
117                                    path = source.get_text_in(leaf.span).trim_matches('"').to_string();
118                                }
119                                _ => {}
120                            }
121                        }
122                        imports.push(ast::Import { path, alias, span: n.span() });
123                    }
124                    _ => {
125                        // Handle import ( ... )
126                        imports.extend(self.extract_imports(n, source));
127                    }
128                }
129            }
130        }
131        imports
132    }
133
134    fn extract_block(&self, node: RedNode<GoLanguage>, source: &SourceText) -> Result<ast::Block, OakError> {
135        let mut statements = vec![];
136        let span = node.span();
137
138        for child in node.children() {
139            if let RedTree::Node(n) = child {
140                if let Some(stmt) = self.extract_statement(n, source)? {
141                    statements.push(stmt);
142                }
143            }
144        }
145
146        Ok(ast::Block { statements, span })
147    }
148
149    fn extract_statement(&self, node: RedNode<GoLanguage>, source: &SourceText) -> Result<Option<ast::Statement>, OakError> {
150        match node.green.kind {
151            GoElementType::ReturnStatement => {
152                let mut values = vec![];
153                for child in node.children() {
154                    if let RedTree::Node(n) = child {
155                        values.push(self.extract_expression(n, source)?);
156                    }
157                }
158                Ok(Some(ast::Statement::Return { values, span: node.span() }))
159            }
160            GoElementType::IfStatement => {
161                let mut condition = ast::Expression::Literal { value: "true".to_string(), span: node.span() };
162                let mut then_block = ast::Block { statements: vec![], span: node.span() };
163                let mut else_block = None;
164
165                for child in node.children() {
166                    if let RedTree::Node(n) = child {
167                        match n.green.kind {
168                            GoElementType::BinaryExpression | GoElementType::CallExpression | GoElementType::Identifier => {
169                                condition = self.extract_expression(n, source)?;
170                            }
171                            GoElementType::Block => {
172                                if then_block.statements.is_empty() {
173                                    then_block = self.extract_block(n, source)?;
174                                }
175                                else {
176                                    else_block = Some(self.extract_block(n, source)?);
177                                }
178                            }
179                            GoElementType::IfStatement => {
180                                // Handle else if
181                                let inner_if = self.extract_statement(n, source)?;
182                                if let Some(ast::Statement::If { condition, then_block, else_block: inner_else, span }) = inner_if {
183                                    else_block = Some(ast::Block { statements: vec![ast::Statement::If { condition, then_block, else_block: inner_else, span }], span });
184                                }
185                            }
186                            _ => {}
187                        }
188                    }
189                }
190                Ok(Some(ast::Statement::If { condition, then_block, else_block, span: node.span() }))
191            }
192            GoElementType::ForStatement => {
193                let mut init = None;
194                let mut condition = None;
195                let mut post = None;
196                let mut body = ast::Block { statements: vec![], span: node.span() };
197
198                for child in node.children() {
199                    if let RedTree::Node(n) = child {
200                        match n.green.kind {
201                            GoElementType::ShortVarDecl | GoElementType::AssignmentStatement | GoElementType::VariableDeclaration => {
202                                if init.is_none() {
203                                    init = self.extract_statement(n, source)?.map(Box::new);
204                                }
205                                else {
206                                    post = self.extract_statement(n, source)?.map(Box::new);
207                                }
208                            }
209                            GoElementType::BinaryExpression | GoElementType::CallExpression | GoElementType::Identifier => {
210                                condition = Some(self.extract_expression(n, source)?);
211                            }
212                            GoElementType::Block => {
213                                body = self.extract_block(n, source)?;
214                            }
215                            _ => {}
216                        }
217                    }
218                }
219                Ok(Some(ast::Statement::For { init, condition, post, body, span: node.span() }))
220            }
221            GoElementType::AssignmentStatement | GoElementType::ShortVarDecl | GoElementType::VariableDeclaration | GoElementType::VariableSpec => {
222                // Support multiple assignment
223                let mut targets = vec![];
224                let mut values = vec![];
225
226                for child in node.children() {
227                    match child {
228                        RedTree::Leaf(leaf) if leaf.kind == GoTokenType::Identifier => {
229                            targets.push(source.get_text_in(leaf.span).trim().to_string());
230                        }
231                        RedTree::Node(n) => {
232                            if n.green.kind == GoElementType::VariableSpec || n.green.kind == GoElementType::VariableDeclaration {
233                                // Recursive extraction
234                                if let Some(ast::Statement::Assignment { targets: t, values: v, .. }) = self.extract_statement(n, source)? {
235                                    targets.extend(t);
236                                    values.extend(v);
237                                }
238                            }
239                            else if n.green.kind != GoElementType::Identifier && !n.green.kind.is_keyword() {
240                                values.push(self.extract_expression(n, source)?);
241                            }
242                            else if n.green.kind == GoElementType::Identifier {
243                                targets.push(source.get_text_in(n.span()).trim().to_string());
244                            }
245                        }
246                        _ => {}
247                    }
248                }
249
250                if targets.is_empty() {
251                    return Ok(None);
252                }
253
254                // If no value (e.g., variable declaration), fill with default value
255                if values.is_empty() {
256                    for _ in &targets {
257                        values.push(ast::Expression::Literal { value: "0".to_string(), span: node.span() });
258                    }
259                }
260
261                Ok(Some(ast::Statement::Assignment { targets, values, span: node.span() }))
262            }
263            _ => {
264                // Could be an expression statement
265                if let Ok(expr) = self.extract_expression(node, source) { Ok(Some(ast::Statement::Expression(expr))) } else { Ok(None) }
266            }
267        }
268    }
269
270    fn extract_expression(&self, node: RedNode<GoLanguage>, source: &SourceText) -> Result<ast::Expression, OakError> {
271        self.extract_expression_internal(RedTree::Node(node), source)
272    }
273
274    fn extract_expression_internal(&self, tree: RedTree<GoLanguage>, source: &SourceText) -> Result<ast::Expression, OakError> {
275        match tree {
276            RedTree::Leaf(leaf) => {
277                if leaf.kind == GoTokenType::Identifier {
278                    let name = source.get_text_in(leaf.span).trim().to_string();
279                    if name.is_empty() {
280                        return Err(OakError::parse_error("Empty identifier leaf"));
281                    }
282                    Ok(ast::Expression::Identifier { name, span: leaf.span })
283                }
284                else if leaf.kind == GoTokenType::IntLiteral || leaf.kind == GoTokenType::StringLiteral || leaf.kind == GoTokenType::BoolLiteral {
285                    Ok(ast::Expression::Literal { value: source.get_text_in(leaf.span).to_string(), span: leaf.span })
286                }
287                else {
288                    Err(OakError::parse_error(format!("Unexpected leaf in expression: {:?}", leaf.kind)))
289                }
290            }
291            RedTree::Node(node) => match node.green.kind {
292                GoElementType::Identifier => {
293                    let mut name = String::new();
294                    for child in node.children() {
295                        match child {
296                            RedTree::Leaf(leaf) if leaf.kind == GoTokenType::Identifier => {
297                                name = source.get_text_in(leaf.span).trim().to_string();
298                                if !name.is_empty() {
299                                    break;
300                                }
301                            }
302                            RedTree::Node(n) => {
303                                println!("DEBUG: Identifier node has a Node child: kind={:?}, span={:?}", n.green.kind, n.span());
304                            }
305                            _ => {}
306                        }
307                    }
308                    if name.is_empty() {
309                        name = source.get_text_in(node.span()).trim().to_string();
310                    }
311                    if name.is_empty() {
312                        println!("DEBUG: Final empty identifier node details:");
313                        println!("  Node kind: {:?}", node.green.kind);
314                        println!("  Node span: {:?}", node.span());
315                        println!("  Children count: {}", node.children().count());
316                        for (i, child) in node.children().enumerate() {
317                            match child {
318                                RedTree::Node(n) => println!("    child {}: Node kind={:?}, span={:?}", i, n.green.kind, n.span()),
319                                RedTree::Leaf(l) => println!("    child {}: Leaf kind={:?}, span={:?}, text={:?}", i, l.kind, l.span, source.get_text_in(l.span)),
320                            }
321                        }
322                        return Err(OakError::parse_error(format!("Empty identifier at {:?}", node.span())));
323                    }
324                    Ok(ast::Expression::Identifier { name, span: node.span() })
325                }
326                GoElementType::IntLiteral | GoElementType::StringLiteral | GoElementType::BoolLiteral => Ok(ast::Expression::Literal { value: source.get_text_in(node.span()).trim().to_string(), span: node.span() }),
327                GoElementType::BinaryExpression => {
328                    let mut left = None;
329                    let mut op = String::new();
330                    let mut right = None;
331
332                    for child in node.children() {
333                        match child {
334                            RedTree::Node(n) => {
335                                if left.is_none() {
336                                    left = Some(Box::new(self.extract_expression(n, source)?));
337                                }
338                                else {
339                                    right = Some(Box::new(self.extract_expression(n, source)?));
340                                }
341                            }
342                            RedTree::Leaf(leaf) => {
343                                if leaf.kind == GoTokenType::Identifier {
344                                    if left.is_none() {
345                                        left = Some(Box::new(ast::Expression::Identifier { name: source.get_text_in(leaf.span).trim().to_string(), span: leaf.span }));
346                                    }
347                                    else {
348                                        right = Some(Box::new(ast::Expression::Identifier { name: source.get_text_in(leaf.span).trim().to_string(), span: leaf.span }));
349                                    }
350                                }
351                                else if TokenType::role(&leaf.kind) == oak_core::UniversalTokenRole::Operator {
352                                    op = source.get_text_in(leaf.span).to_string();
353                                }
354                            }
355                        }
356                    }
357
358                    if let (Some(left), Some(right)) = (left, right) { Ok(ast::Expression::Binary { left, op, right, span: node.span() }) } else { Ok(ast::Expression::Literal { value: source.get_text_in(node.span()).to_string(), span: node.span() }) }
359                }
360                GoElementType::CallExpression => {
361                    let mut func = None;
362                    let mut args = vec![];
363
364                    for child in node.children() {
365                        match child {
366                            RedTree::Node(n) => {
367                                if func.is_none() {
368                                    func = Some(Box::new(self.extract_expression(n, source)?));
369                                }
370                                else if n.green.kind == GoElementType::ExpressionList {
371                                    for list_child in n.children() {
372                                        if let RedTree::Node(ln) = list_child {
373                                            args.push(self.extract_expression(ln, source)?);
374                                        }
375                                        else if let RedTree::Leaf(leaf) = list_child {
376                                            if let Ok(expr) = self.extract_expression_internal(RedTree::Leaf(leaf), source) {
377                                                args.push(expr);
378                                            }
379                                        }
380                                    }
381                                }
382                                else {
383                                    args.push(self.extract_expression(n, source)?);
384                                }
385                            }
386                            RedTree::Leaf(leaf) if leaf.kind == GoTokenType::Identifier => {
387                                if func.is_none() {
388                                    func = Some(Box::new(ast::Expression::Identifier { name: source.get_text_in(leaf.span).trim().to_string(), span: leaf.span }));
389                                }
390                                else {
391                                    args.push(ast::Expression::Identifier { name: source.get_text_in(leaf.span).trim().to_string(), span: leaf.span });
392                                }
393                            }
394                            _ => {}
395                        }
396                    }
397
398                    if let Some(func) = func { Ok(ast::Expression::Call { func, args, span: node.span() }) } else { Ok(ast::Expression::Literal { value: source.get_text_in(node.span()).to_string(), span: node.span() }) }
399                }
400                _ => Ok(ast::Expression::Literal { value: source.get_text_in(node.span()).to_string(), span: node.span() }),
401            },
402        }
403    }
404
405    fn extract_variables(&self, node: RedNode<GoLanguage>, source: &SourceText, cache: &mut impl BuilderCache<GoLanguage>) -> Result<Vec<Declaration>, OakError> {
406        if let Some(cached) = cache.get_typed_node::<Vec<Declaration>>(&node.green) {
407            return Ok(cached);
408        }
409        let mut vars = vec![];
410        for child in node.children() {
411            if let RedTree::Node(n) = child {
412                match n.green.kind {
413                    GoElementType::VariableSpec => {
414                        let mut name = String::new();
415                        let mut var_type = None;
416                        let mut value = None;
417                        for spec_child in n.children() {
418                            match spec_child {
419                                RedTree::Leaf(leaf) if leaf.kind == GoTokenType::Identifier => {
420                                    if name.is_empty() {
421                                        name = source.get_text_in(leaf.span).trim().to_string();
422                                    }
423                                    else if var_type.is_none() {
424                                        var_type = Some(source.get_text_in(leaf.span).trim().to_string());
425                                    }
426                                }
427                                RedTree::Node(en) => {
428                                    value = Some(self.extract_expression(en, source)?);
429                                }
430                                _ => {}
431                            }
432                        }
433                        vars.push(Declaration::Variable(ast::Variable { name, var_type, value, span: n.span() }));
434                    }
435                    _ => {
436                        // Handle var ( ... )
437                        vars.extend(self.extract_variables(n, source, cache)?);
438                    }
439                }
440            }
441        }
442        cache.set_typed_node(&node.green, std::sync::Arc::new(vars.clone()));
443        Ok(vars)
444    }
445
446    fn extract_consts(&self, node: RedNode<GoLanguage>, source: &SourceText, cache: &mut impl BuilderCache<GoLanguage>) -> Result<Vec<Declaration>, OakError> {
447        if let Some(cached) = cache.get_typed_node::<Vec<Declaration>>(&node.green) {
448            return Ok(cached);
449        }
450        let mut consts = vec![];
451        for child in node.children() {
452            if let RedTree::Node(n) = child {
453                match n.green.kind {
454                    GoElementType::ConstSpec => {
455                        let mut name = String::new();
456                        let mut const_type = None;
457                        let mut value = ast::Expression::Literal { value: "0".to_string(), span: n.span() };
458                        for spec_child in n.children() {
459                            match spec_child {
460                                RedTree::Leaf(leaf) if leaf.kind == GoTokenType::Identifier => {
461                                    if name.is_empty() {
462                                        name = source.get_text_in(leaf.span).trim().to_string();
463                                    }
464                                    else if const_type.is_none() {
465                                        const_type = Some(source.get_text_in(leaf.span).trim().to_string());
466                                    }
467                                }
468                                RedTree::Node(en) => {
469                                    value = self.extract_expression(en, source)?;
470                                }
471                                _ => {}
472                            }
473                        }
474                        consts.push(Declaration::Const(ast::Const { name, const_type, value, span: n.span() }));
475                    }
476                    _ => {
477                        // Handle const ( ... )
478                        consts.extend(self.extract_consts(n, source, cache)?);
479                    }
480                }
481            }
482        }
483        cache.set_typed_node(&node.green, std::sync::Arc::new(consts.clone()));
484        Ok(consts)
485    }
486
487    fn extract_types(&self, node: RedNode<GoLanguage>, source: &SourceText, cache: &mut impl BuilderCache<GoLanguage>) -> Result<Vec<Declaration>, OakError> {
488        if let Some(cached) = cache.get_typed_node::<Vec<Declaration>>(&node.green) {
489            return Ok(cached);
490        }
491        let mut types = vec![];
492        for child in node.children() {
493            if let RedTree::Node(n) = child {
494                match n.green.kind {
495                    GoElementType::TypeSpec => {
496                        let mut name = String::new();
497                        let mut definition = String::new();
498                        for spec_child in n.children() {
499                            match spec_child {
500                                RedTree::Leaf(leaf) if leaf.kind == GoTokenType::Identifier => {
501                                    if name.is_empty() {
502                                        name = source.get_text_in(leaf.span).trim().to_string();
503                                    }
504                                    else {
505                                        definition = source.get_text_in(leaf.span).trim().to_string();
506                                    }
507                                }
508                                RedTree::Node(tn) => {
509                                    definition = source.get_text_in(tn.span()).to_string();
510                                }
511                                _ => {}
512                            }
513                        }
514                        types.push(Declaration::Type(ast::TypeDecl { name, definition, span: n.span() }));
515                    }
516                    _ => {
517                        // Handle type ( ... )
518                        types.extend(self.extract_types(n, source, cache)?);
519                    }
520                }
521            }
522        }
523        cache.set_typed_node(&node.green, std::sync::Arc::new(types.clone()));
524        Ok(types)
525    }
526}