gbf_core/decompiler/ast/visitors/
emitter.rs

1#![deny(missing_docs)]
2
3use super::{
4    AstVisitor,
5    emit_context::{EmitContext, IndentStyle},
6};
7use crate::decompiler::ast::{AstKind, AstVisitable};
8use crate::decompiler::ast::{
9    array_access::ArrayAccessNode, array_kind::ArrayKind, array_node::ArrayNode,
10    control_flow::ControlFlowType, expr::ExprKind, phi::PhiNode, phi_array::PhiArrayNode,
11};
12use crate::decompiler::ast::{assignment::AssignmentNode, statement::StatementKind};
13use crate::decompiler::ast::{
14    bin_op::{BinOpType, BinaryOperationNode},
15    func_call::FunctionCallNode,
16};
17use crate::decompiler::ast::{block::BlockNode, ptr::P};
18use crate::decompiler::ast::{control_flow::ControlFlowNode, unary_op::UnaryOperationNode};
19use crate::decompiler::ast::{function::FunctionNode, literal::LiteralNode};
20use crate::decompiler::ast::{member_access::MemberAccessNode, ret::ReturnNode};
21use crate::{decompiler::ast::identifier::IdentifierNode, utils::escape_string};
22
23/// An emitter for the AST.
24///
25/// This emitter now builds up and returns a `String` for every AST node rather
26/// than writing to a shared output buffer.
27pub struct Gs2Emitter {
28    /// The context of the emitter.
29    context: EmitContext,
30}
31
32impl Gs2Emitter {
33    /// Creates a new `Gs2Emitter` with the given `context`.
34    pub fn new(context: EmitContext) -> Self {
35        Self { context }
36    }
37
38    /// Merges multiple comment vectors into a single one.
39    ///
40    /// This helper centralizes comment merging so that any future changes (such
41    /// as deduplication or filtering) can be applied in one place.
42    fn merge_comments(&self, comments: Vec<Vec<String>>) -> Vec<String> {
43        comments.into_iter().flatten().collect()
44    }
45
46    /// Returns a string containing spaces corresponding to the current indentation level.
47    fn emit_indent(&self) -> String {
48        " ".repeat(self.context.indent)
49    }
50}
51
52/// The output of the emitter.
53pub struct AstOutput {
54    /// The emitted node.
55    pub node: String,
56    /// The comments associated with the node.
57    pub comments: Vec<String>,
58}
59
60impl AstVisitor for Gs2Emitter {
61    type Output = AstOutput;
62
63    /// Visits an AST node.
64    fn visit_node(&mut self, node: &AstKind) -> AstOutput {
65        match node {
66            AstKind::Expression(expr) => expr.accept(self),
67            AstKind::Statement(stmt) => stmt.accept(self),
68            AstKind::Function(func) => func.accept(self),
69            AstKind::Block(block) => block.accept(self),
70            AstKind::ControlFlow(control_flow) => control_flow.accept(self),
71        }
72    }
73
74    /// Visits a statement node.
75    fn visit_statement(&mut self, node: &StatementKind) -> AstOutput {
76        self.context = self.context.with_expr_root(true);
77        let stmt_str = match node {
78            StatementKind::Assignment(assignment) => assignment.accept(self),
79            StatementKind::Return(ret) => ret.accept(self),
80            StatementKind::VirtualBranch(vbranch) => vbranch.accept(self),
81        };
82        AstOutput {
83            node: format!("{};", stmt_str.node),
84            comments: stmt_str.comments,
85        }
86    }
87
88    /// Visits an assignment node.
89    fn visit_assignment(&mut self, stmt_node: &P<AssignmentNode>) -> AstOutput {
90        let base_comments = stmt_node.metadata().comments().clone();
91        // Step 1: Visit and emit the LHS.
92        let lhs_str = stmt_node.lhs.accept(self);
93
94        // Step 2: Check for binary operations that use the LHS.
95        if let ExprKind::BinOp(bin_op_node) = stmt_node.rhs.clone() {
96            let lhs_in_rhs = bin_op_node.lhs == stmt_node.lhs.clone();
97            if lhs_in_rhs {
98                match bin_op_node.op_type {
99                    BinOpType::Add => {
100                        if let ExprKind::Literal(lit) = bin_op_node.rhs.clone() {
101                            if let LiteralNode::Number(num) = lit.as_ref() {
102                                if *num == 1 {
103                                    return AstOutput {
104                                        node: format!("{}++", lhs_str.node),
105                                        comments: self.merge_comments(vec![
106                                            base_comments.clone(),
107                                            lhs_str.comments.clone(),
108                                        ]),
109                                    };
110                                } else {
111                                    let rhs_str = bin_op_node.rhs.accept(self);
112                                    return AstOutput {
113                                        node: format!("{} += {}", lhs_str.node, rhs_str.node),
114                                        comments: self.merge_comments(vec![
115                                            base_comments.clone(),
116                                            lhs_str.comments.clone(),
117                                            rhs_str.comments,
118                                        ]),
119                                    };
120                                }
121                            }
122                        }
123                    }
124                    BinOpType::Sub => {
125                        if let ExprKind::Literal(lit) = bin_op_node.rhs.clone() {
126                            if let LiteralNode::Number(num) = lit.as_ref() {
127                                if *num == 1 {
128                                    return AstOutput {
129                                        node: format!("{}--", lhs_str.node),
130                                        comments: self.merge_comments(vec![
131                                            base_comments.clone(),
132                                            lhs_str.comments.clone(),
133                                        ]),
134                                    };
135                                } else {
136                                    let rhs_str = bin_op_node.rhs.accept(self);
137                                    return AstOutput {
138                                        node: format!("{} -= {}", lhs_str.node, rhs_str.node),
139                                        comments: self.merge_comments(vec![
140                                            base_comments.clone(),
141                                            lhs_str.comments.clone(),
142                                            rhs_str.comments,
143                                        ]),
144                                    };
145                                }
146                            }
147                        }
148                    }
149                    _ => {}
150                }
151            }
152        }
153
154        // Step 3: Default assignment.
155        let prev_context = self.context;
156        self.context = self.context.with_expr_root(true);
157        let rhs_str = stmt_node.rhs.accept(self);
158        self.context = prev_context;
159        AstOutput {
160            node: format!("{} = {}", lhs_str.node, rhs_str.node),
161            comments: self.merge_comments(vec![base_comments, lhs_str.comments, rhs_str.comments]),
162        }
163    }
164
165    /// Visits a virtual branch node.
166    fn visit_virtual_branch(
167        &mut self,
168        node: &P<crate::decompiler::ast::vbranch::VirtualBranchNode>,
169    ) -> Self::Output {
170        // Put out `goto` statement
171        let mut s = String::new();
172        s.push_str("goto ");
173        s.push_str(&node.branch().to_string());
174
175        AstOutput {
176            node: s,
177            comments: node.metadata().comments().clone(),
178        }
179    }
180
181    /// Visits an expression node.
182    fn visit_expr(&mut self, node: &ExprKind) -> AstOutput {
183        match node {
184            ExprKind::Literal(literal) => literal.accept(self),
185            ExprKind::BinOp(bin_op) => bin_op.accept(self),
186            ExprKind::UnaryOp(unary_op) => unary_op.accept(self),
187            ExprKind::FunctionCall(func_call) => func_call.accept(self),
188            ExprKind::Array(array) => array.accept(self),
189            ExprKind::New(new_node) => new_node.accept(self),
190            ExprKind::NewArray(new_array) => new_array.accept(self),
191            ExprKind::MemberAccess(member_access) => member_access.accept(self),
192            ExprKind::Identifier(identifier) => identifier.accept(self),
193            ExprKind::ArrayAccess(array_access) => array_access.accept(self),
194            ExprKind::Phi(phi) => phi.accept(self),
195            ExprKind::Range(range) => range.accept(self),
196        }
197    }
198
199    /// Visits an array node.
200    fn visit_array(&mut self, node: &ArrayKind) -> AstOutput {
201        match node {
202            ArrayKind::MergedArray(array) => self.visit_array_node(array),
203            ArrayKind::PhiArray(phi_array) => self.visit_phi_array(phi_array),
204        }
205    }
206
207    /// Visits an array node.
208    fn visit_array_node(&mut self, node: &P<ArrayNode>) -> AstOutput {
209        let mut s = String::new();
210        let mut comments = node.metadata().comments().clone();
211        s.push('{');
212        for (i, elem) in node.elements.iter().enumerate() {
213            let elem_out = elem.accept(self);
214            s.push_str(&elem_out.node);
215            comments.extend(elem_out.comments);
216            if i < node.elements.len() - 1 {
217                s.push_str(", ");
218            }
219        }
220        s.push('}');
221        AstOutput { node: s, comments }
222    }
223
224    /// Visits a phi array node.
225    fn visit_phi_array(&mut self, node: &P<PhiArrayNode>) -> AstOutput {
226        // output will look like { elem1, elem2, elem3, ... phi }
227        let mut s = String::new();
228        let mut comments = node.metadata().comments().clone();
229        s.push('{');
230        for (i, elem) in node.elements.iter().enumerate() {
231            let elem_out = elem.accept(self);
232            s.push_str(&elem_out.node);
233            comments.extend(elem_out.comments);
234            if i < node.elements.len() - 1 {
235                s.push_str(", ");
236            }
237        }
238        s.push_str(", ...");
239        let phi_out = node.phi.accept(self);
240        s.push_str(&phi_out.node);
241        comments.extend(phi_out.comments);
242        s.push('}');
243        AstOutput { node: s, comments }
244    }
245
246    /// Visits an array access node.
247    fn visit_array_access(&mut self, node: &P<ArrayAccessNode>) -> AstOutput {
248        let array_str = node.arr.accept(self);
249        let index_str = node.index.accept(self);
250        AstOutput {
251            node: format!("{}[{}]", array_str.node, index_str.node),
252            comments: self.merge_comments(vec![
253                node.metadata().comments().clone(),
254                array_str.comments,
255                index_str.comments,
256            ]),
257        }
258    }
259
260    /// Visits a binary operation node.
261    fn visit_bin_op(&mut self, node: &P<BinaryOperationNode>) -> AstOutput {
262        let base_comments = node.metadata().comments().clone();
263        let prev_context = self.context;
264        self.context = self.context.with_expr_root(false);
265        let lhs_str = node.lhs.accept(self);
266        let rhs_str = node.rhs.accept(self);
267        self.context = prev_context;
268        let op_str = node.op_type.to_string();
269        if self.context.expr_root {
270            AstOutput {
271                node: format!("{} {} {}", lhs_str.node, op_str, rhs_str.node),
272                comments: self.merge_comments(vec![
273                    base_comments,
274                    lhs_str.comments,
275                    rhs_str.comments,
276                ]),
277            }
278        } else {
279            AstOutput {
280                node: format!("({} {} {})", lhs_str.node, op_str, rhs_str.node),
281                comments: self.merge_comments(vec![
282                    base_comments,
283                    lhs_str.comments,
284                    rhs_str.comments,
285                ]),
286            }
287        }
288    }
289
290    /// Visits a unary operation node.
291    fn visit_unary_op(&mut self, node: &P<UnaryOperationNode>) -> AstOutput {
292        let base_comments = node.metadata().comments().clone();
293        let prev_context = self.context;
294        self.context = self.context.with_expr_root(false);
295        let operand_str = node.operand.accept(self);
296        self.context = prev_context;
297        let op_str = node.op_type.to_string();
298        if self.context.expr_root {
299            AstOutput {
300                node: format!("{}{}", op_str, operand_str.node),
301                comments: self.merge_comments(vec![
302                    base_comments,
303                    node.metadata().comments().clone(),
304                    operand_str.comments,
305                ]),
306            }
307        } else {
308            AstOutput {
309                node: format!("({}{})", op_str, operand_str.node),
310                comments: self.merge_comments(vec![
311                    base_comments,
312                    node.metadata().comments().clone(),
313                    operand_str.comments,
314                ]),
315            }
316        }
317    }
318
319    /// Visits an identifier node.
320    fn visit_identifier(&mut self, node: &P<IdentifierNode>) -> AstOutput {
321        let mut s = node.id().clone();
322        if self.context.include_ssa_versions {
323            if let Some(ssa_version) = node.ssa_version {
324                s.push_str(&format!("#{}", ssa_version));
325            }
326        }
327        AstOutput {
328            node: s,
329            comments: node.metadata().comments().clone(),
330        }
331    }
332
333    /// Visits a literal node.
334    fn visit_literal(&mut self, node: &P<LiteralNode>) -> AstOutput {
335        match node.as_ref() {
336            LiteralNode::String(s) => {
337                let escaped = escape_string(s);
338                AstOutput {
339                    node: format!("\"{}\"", escaped),
340                    comments: node.metadata().comments().clone(),
341                }
342            }
343            LiteralNode::Number(n) => {
344                if self.context.format_number_hex {
345                    AstOutput {
346                        node: format!("0x{:x}", n),
347                        comments: node.metadata().comments().clone(),
348                    }
349                } else {
350                    AstOutput {
351                        node: n.to_string(),
352                        comments: node.metadata().comments().clone(),
353                    }
354                }
355            }
356            LiteralNode::Float(f) => AstOutput {
357                node: f.to_string(),
358                comments: node.metadata().comments().clone(),
359            },
360            LiteralNode::Boolean(b) => AstOutput {
361                node: b.to_string(),
362                comments: node.metadata().comments().clone(),
363            },
364            LiteralNode::Null => AstOutput {
365                node: "null".to_string(),
366                comments: node.metadata().comments().clone(),
367            },
368        }
369    }
370
371    /// Visits a member access node.
372    fn visit_member_access(&mut self, node: &P<MemberAccessNode>) -> AstOutput {
373        let lhs_str = node.lhs.accept(self);
374        let rhs_str = node.rhs.accept(self);
375        AstOutput {
376            node: format!("{}.{}", lhs_str.node, rhs_str.node),
377            comments: self.merge_comments(vec![
378                node.metadata().comments().clone(),
379                lhs_str.comments,
380                rhs_str.comments,
381            ]),
382        }
383    }
384
385    /// Visits a function call node.
386    fn visit_function_call(&mut self, node: &P<FunctionCallNode>) -> AstOutput {
387        // We only care about node.arguments.
388        match &node.arguments {
389            ArrayKind::MergedArray(array_node) => {
390                // The merged array contains the function name as element 0,
391                // and subsequent elements are the arguments.
392                let elements = &array_node.elements;
393                if elements.is_empty() {
394                    // Should not happen, but if it does, return an empty call.
395                    return AstOutput {
396                        node: "fn()".to_string(),
397                        comments: Vec::new(),
398                    };
399                }
400
401                // Visit the first element as the function name.
402                let name_out = elements[0].accept(self);
403                let mut s = String::new();
404                s.push_str(name_out.node.as_str());
405                s.push('(');
406
407                // Visit the rest as arguments.
408                let mut arg_comments = Vec::new();
409                // Use skip(1) because element 0 is the function name.
410                let arg_outputs: Vec<AstOutput> = elements
411                    .iter()
412                    .skip(1)
413                    .map(|arg| arg.accept(self))
414                    .collect();
415                for (i, arg_out) in arg_outputs.iter().enumerate() {
416                    s.push_str(&arg_out.node);
417                    arg_comments.extend(arg_out.comments.clone());
418                    if i < arg_outputs.len() - 1 {
419                        s.push_str(", ");
420                    }
421                }
422                s.push(')');
423
424                // Merge comments from the function name, the node metadata, and the arguments.
425                let comments = self.merge_comments(vec![
426                    node.metadata().comments().clone(),
427                    name_out.comments,
428                    arg_comments,
429                ]);
430                AstOutput { node: s, comments }
431            }
432            ArrayKind::PhiArray(phi_array_node) => {
433                // For a phi array, we output a special string.
434                // Here, we call accept on the PhiArray node to generate its contents,
435                // then wrap that in a "phi_fn_call(...)" string.
436                let call_out = phi_array_node.accept(self);
437                let s = format!("phi_fn_call({})", call_out.node);
438                AstOutput {
439                    node: s,
440                    comments: self.merge_comments(vec![
441                        node.metadata().comments().clone(),
442                        call_out.comments,
443                    ]),
444                }
445            }
446        }
447    }
448
449    /// Visits a function node.
450    fn visit_function(&mut self, node: &P<FunctionNode>) -> AstOutput {
451        let mut comments = node.metadata().comments().clone();
452
453        // Handle anonymous functions by emitting the body only.
454        if node.name().is_none() {
455            let mut s = String::new();
456            for stmt in node.body().instructions.iter() {
457                let stmt_out = stmt.accept(self);
458                // Emit any comments.
459                for comment in stmt_out.comments.iter() {
460                    s.push_str(&self.emit_indent());
461                    s.push_str("// ");
462                    s.push_str(comment);
463                    s.push('\n');
464                }
465                s.push_str(&stmt_out.node);
466                s.push('\n');
467            }
468            return AstOutput { node: s, comments };
469        }
470
471        // At this point we know the function has a name.
472        let name = node.name().as_ref().unwrap();
473        let mut s = String::new();
474        s.push_str(&format!("function {}(", name));
475
476        // Now, since node.params is of type ArrayKind, we match on it.
477        match &node.params() {
478            ArrayKind::MergedArray(array_node) => {
479                // The MergedArray's elements are the parameters.
480                let params = &array_node.elements;
481                for (i, param) in params.iter().enumerate() {
482                    let param_out = param.accept(self);
483                    comments.extend(param_out.comments);
484                    s.push_str(&param_out.node);
485                    if i < params.len() - 1 {
486                        s.push_str(", ");
487                    }
488                }
489            }
490            ArrayKind::PhiArray(phi_array_node) => {
491                // For a PhiArray we indicate that the parameters are unresolved.
492                let phi_out = phi_array_node.accept(self);
493                s.push_str(&format!("phi_params({})", phi_out.node));
494                comments.extend(phi_out.comments);
495            }
496        }
497        s.push(')');
498
499        // Visit the function body.
500        let block_output = node.body().accept(self);
501        s.push_str(&block_output.node);
502
503        AstOutput {
504            node: s,
505            comments: self.merge_comments(vec![comments, block_output.comments]),
506        }
507    }
508
509    /// Visits a return node.
510    fn visit_return(&mut self, node: &P<ReturnNode>) -> AstOutput {
511        let child = node.ret.accept(self);
512        let mut s = String::new();
513        s.push_str("return ");
514        s.push_str(&child.node);
515        AstOutput {
516            node: s,
517            comments: self.merge_comments(vec![node.metadata().comments().clone(), child.comments]),
518        }
519    }
520
521    /// Visits a block node.
522    fn visit_block(&mut self, node: &P<BlockNode>) -> AstOutput {
523        let mut s = String::new();
524        if self.context.indent_style == IndentStyle::Allman {
525            s.push('\n');
526            s.push_str(&self.emit_indent());
527            s.push_str("{\n");
528        } else {
529            s.push_str(" {\n");
530        }
531        let old_context = self.context;
532        self.context = self.context.with_indent();
533        if !node.instructions.is_empty() {
534            for stmt in node.instructions.iter() {
535                let stmt_out = stmt.accept(self);
536                // First emit any comments.
537                for comment in stmt_out.comments.iter() {
538                    s.push_str(&self.emit_indent());
539                    s.push_str("// ");
540                    s.push_str(comment);
541                    s.push('\n');
542                }
543                // Then emit the statement.
544                s.push_str(&self.emit_indent());
545                s.push_str(&stmt_out.node);
546                s.push('\n');
547            }
548        }
549        self.context = old_context;
550        s.push_str(&self.emit_indent());
551        s.push('}');
552        AstOutput {
553            node: s,
554            comments: node.metadata().comments().clone(),
555        }
556    }
557
558    /// Visits a control flow node.
559    fn visit_control_flow(&mut self, node: &P<ControlFlowNode>) -> AstOutput {
560        let mut s = String::new();
561        let mut base_comments = node.metadata().comments().clone();
562        let name = match node.ty() {
563            ControlFlowType::If => "if",
564            ControlFlowType::Else => "else",
565            ControlFlowType::ElseIf => "else if",
566            ControlFlowType::With => "with",
567            ControlFlowType::While => "while",
568            ControlFlowType::For => "for",
569            ControlFlowType::DoWhile => "do",
570        };
571        if *node.ty() == ControlFlowType::DoWhile {
572            s.push_str(name);
573            let body_out = node.body().accept(self);
574            s.push(' ');
575            s.push_str(&body_out.node);
576            s.push_str(" while (");
577            if let Some(condition) = node.condition() {
578                let condition_out = condition.accept(self);
579                s.push_str(&condition_out.node);
580                base_comments.extend(condition_out.comments.clone());
581            }
582            s.push_str(");");
583            AstOutput {
584                node: s,
585                comments: self.merge_comments(vec![base_comments, body_out.comments]),
586            }
587        } else {
588            s.push_str(name);
589            if let Some(condition) = node.condition() {
590                let condition_out = condition.accept(self);
591                s.push_str(" (");
592                s.push_str(&condition_out.node);
593                s.push_str(") ");
594                base_comments.extend(condition_out.comments.clone());
595            }
596            let body_out = node.body().accept(self);
597            s.push_str(&body_out.node);
598            AstOutput {
599                node: s,
600                comments: self.merge_comments(vec![base_comments, body_out.comments]),
601            }
602        }
603    }
604
605    /// Visits a phi node.
606    fn visit_phi(&mut self, node: &P<PhiNode>) -> AstOutput {
607        let mut s = String::new();
608        s.push_str("phi<idx=");
609        s.push_str(&node.idx().to_string());
610        s.push_str(", regions=(");
611        for (i, region) in node.regions().iter().enumerate() {
612            s.push_str(&region.0.to_string());
613            if i < node.regions().len() - 1 {
614                s.push_str(", ");
615            }
616        }
617        s.push_str(")>");
618        AstOutput {
619            node: s,
620            comments: node.metadata().comments().clone(),
621        }
622    }
623
624    /// Visits a new node
625    fn visit_new(&mut self, node: &P<crate::decompiler::ast::new::NewNode>) -> AstOutput {
626        let type_out = node.new_type.accept(self);
627        let arg_out = node.arg.accept(self);
628        // TODO: if type_out is a string literal, we shouldn't put out the quotes.
629        AstOutput {
630            node: format!("new {}({})", type_out.node, arg_out.node),
631            comments: self.merge_comments(vec![
632                node.metadata().comments().clone(),
633                type_out.comments,
634                arg_out.comments,
635            ]),
636        }
637    }
638
639    /// Visits a new array node
640    fn visit_new_array(
641        &mut self,
642        node: &P<crate::decompiler::ast::new_array::NewArrayNode>,
643    ) -> AstOutput {
644        let arg_out = node.arg.accept(self);
645        AstOutput {
646            node: format!("new [{}]", arg_out.node),
647            comments: self
648                .merge_comments(vec![node.metadata().comments().clone(), arg_out.comments]),
649        }
650    }
651
652    /// Visits a range node
653    fn visit_range(&mut self, node: &P<crate::decompiler::ast::range::RangeNode>) -> AstOutput {
654        let start_out = node.start.accept(self);
655        let end_out = node.end.accept(self);
656        AstOutput {
657            node: format!("<{}, {}>", start_out.node, end_out.node),
658            comments: self.merge_comments(vec![
659                node.metadata().comments().clone(),
660                start_out.comments,
661                end_out.comments,
662            ]),
663        }
664    }
665}
gbf_core/decompiler/ast/visitors/emitter.rs

gbf_core/decompiler/ast/visitors/
emitter.rs
