gbf_core/decompiler/ast/
control_flow.rs

1#![deny(missing_docs)]
2
3use gbf_macros::AstNodeTransform;
4use serde::{Deserialize, Serialize};
5
6use super::{
7    AstKind, AstVisitable, block::BlockNode, expr::ExprKind, ptr::P, visitors::AstVisitor,
8};
9
10/// Represents the type of control flow node.
11#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
12pub enum ControlFlowType {
13    /// Represents an if control flow node.
14    If,
15    /// Represents an else if control flow node.
16    ElseIf,
17    /// Represents an else control flow node.
18    Else,
19    /// Represents a with control flow node.
20    With,
21    /// Represents a while control flow node.
22    While,
23    /// Represents a for control flow node.
24    For,
25    /// Represents a DoWhile control flow node.
26    DoWhile,
27}
28
29/// Represents a metadata node in the AST
30#[derive(Debug, Clone, Serialize, Deserialize, Eq, AstNodeTransform)]
31#[convert_to(AstKind::ControlFlow)]
32pub struct ControlFlowNode {
33    ty: ControlFlowType,
34    expr: Option<ExprKind>,
35    then_block: P<BlockNode>,
36}
37
38impl ControlFlowNode {
39    /// Creates a new `ControlFlowNode` with the given name, condition, and body.
40    ///
41    /// # Arguments
42    /// - `ty`: The type of the control flow node.
43    /// - `condition`: The condition of the control flow node.
44    /// - `body`: The body of the control flow node.
45    ///
46    /// # Returns
47    /// A new `ControlFlowNode`.
48    pub fn new<E, V>(ty: ControlFlowType, condition: Option<E>, body: Vec<V>) -> Self
49    where
50        E: Into<ExprKind>,
51        V: Into<AstKind>,
52    {
53        Self {
54            ty,
55            expr: condition.map(|e| e.into()),
56            then_block: BlockNode::new(body).into(),
57        }
58    }
59
60    /// Returns the condition of the ControlFlowNode.
61    pub fn condition(&self) -> &Option<ExprKind> {
62        &self.expr
63    }
64
65    /// Returns the body of the ControlFlowNode.
66    pub fn body(&self) -> &P<BlockNode> {
67        &self.then_block
68    }
69
70    /// Returns the type of the ControlFlowNode.
71    pub fn ty(&self) -> &ControlFlowType {
72        &self.ty
73    }
74}
75
76// == Other implementations for literal ==
77impl AstVisitable for P<ControlFlowNode> {
78    fn accept<V: AstVisitor>(&self, visitor: &mut V) -> V::Output {
79        visitor.visit_control_flow(self)
80    }
81}
82
83impl PartialEq for ControlFlowNode {
84    fn eq(&self, other: &Self) -> bool {
85        self.ty == other.ty && self.expr == other.expr && self.then_block == other.then_block
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92    use crate::decompiler::ast::{
93        AstNodeError, bin_op::BinOpType, emit, new_bin_op, new_fn, new_id, new_member_access,
94        new_num, new_return, new_str,
95    };
96
97    #[test]
98    fn test_control_flow_node() -> Result<(), AstNodeError> {
99        /* if (foo.bar == "baz")  { return 1; } */
100        let condition = new_member_access(new_id("foo"), new_id("bar"))?;
101        let condition = new_bin_op(condition, new_str("baz"), BinOpType::Equal)?;
102        let body = BlockNode::new(vec![new_return(new_num(1))]);
103        let control_flow = ControlFlowNode::new(ControlFlowType::If, Some(condition), vec![body]);
104        assert_eq!(control_flow.ty(), &ControlFlowType::If);
105        assert!(control_flow.condition().is_some());
106        assert_eq!(control_flow.body().instructions.len(), 1);
107        Ok(())
108    }
109
110    #[test]
111    fn test_control_flow_node_emit() -> Result<(), AstNodeError> {
112        /* if (foo.bar == "baz")  { return 1; } */
113        let condition = new_member_access(new_id("foo"), new_id("bar"))?;
114        let condition = new_bin_op(condition, new_str("baz"), BinOpType::Equal)?;
115        let body = vec![new_return(new_num(1))];
116        let control_flow = ControlFlowNode::new(ControlFlowType::If, Some(condition), body);
117        let function = new_fn(
118            Some("onCreated".to_string()),
119            vec![new_id("test")],
120            vec![control_flow],
121        );
122        let output = emit(function);
123        assert_eq!(
124            output,
125            "function onCreated(test)\n{\n    if (foo.bar == \"baz\") \n    {\n        return 1;\n    }\n}"
126        );
127        Ok(())
128    }
129
130    #[test]
131    fn test_control_flow_else_emit() -> Result<(), AstNodeError> {
132        /* if (foo.bar == "baz")  { return 1; } else { return 2; } */
133        let condition = new_member_access(new_id("foo"), new_id("bar"))?;
134        let condition = new_bin_op(condition, new_str("baz"), BinOpType::Equal)?;
135        let body = vec![new_return(new_num(1))];
136        let else_body = vec![new_return(new_num(2))];
137        let control_flow = ControlFlowNode::new(ControlFlowType::If, Some(condition), body);
138        let else_control_flow =
139            ControlFlowNode::new(ControlFlowType::Else, None::<ExprKind>, else_body);
140        let function = new_fn(
141            Some("onCreated".to_string()),
142            vec![new_id("test")],
143            vec![control_flow, else_control_flow],
144        );
145        let output = emit(function);
146        assert_eq!(
147            output,
148            "function onCreated(test)\n{\n    if (foo.bar == \"baz\") \n    {\n        return 1;\n    }\n    else\n    {\n        return 2;\n    }\n}"
149        );
150        Ok(())
151    }
152}