Skip to main content

factorio_codegen/generator/
statement.rs

1use factorio_ir::{module::Module, scope::Scope, statement::Statement};
2
3use crate::{LuaGenerator, LuaGeneratorResult};
4
5fn body_has_continue(body: &[Statement]) -> bool {
6    body.iter().any(|s| match s {
7        Statement::Continue => true,
8        Statement::Conditional {
9            then_block,
10            else_block,
11            ..
12        } => body_has_continue(then_block) || body_has_continue(else_block),
13        Statement::ForIn { body, .. } | Statement::While { body, .. } => body_has_continue(body),
14        _ => false,
15    })
16}
17
18impl LuaGenerator {
19    /// Generate code for a given [`Statement`].
20    pub(crate) fn generate_statement(
21        &mut self,
22        statement: &Statement,
23        module: Option<&Module>,
24        module_name: Option<&str>,
25        scope: Scope,
26    ) -> LuaGeneratorResult<()> {
27        match statement {
28            Statement::FunctionDecl(function) => {
29                self.generate_function(function, module, scope, module_name)?;
30            }
31            Statement::StructDecl(struct_decl) => {
32                self.generate_struct(struct_decl, module, scope, module_name)?;
33            }
34            Statement::VariableDecl {
35                name,
36                value,
37                source_type,
38                ..
39            } => {
40                let value = self.generate_expression(value);
41                let type_comment = self.variable_type_comment(source_type.as_deref());
42                let line = match (scope, module_name) {
43                    (Scope::Public, Some(module_name)) => {
44                        format!("{module_name}.{name}{type_comment} = {value}")
45                    }
46                    _ => format!("local {name}{type_comment} = {value}"),
47                };
48                self.write_line(&line);
49            }
50            Statement::Assignment { target, value } => {
51                let target = self.generate_expression(target);
52                let value = self.generate_expression(value);
53                self.write_line(&format!("{target} = {value}"));
54            }
55            Statement::Conditional {
56                condition,
57                then_block,
58                else_block,
59            } => {
60                let condition = self.generate_expression(condition);
61                self.write_line(&format!("if {condition} then"));
62
63                self.indent_level += 1;
64                for statement in then_block {
65                    self.generate_statement(statement, module, module_name, Scope::Private)?;
66                }
67                self.indent_level -= 1;
68
69                if !else_block.is_empty() {
70                    self.write_line("else");
71                    self.indent_level += 1;
72                    for statement in else_block {
73                        self.generate_statement(statement, module, module_name, Scope::Private)?;
74                    }
75                    self.indent_level -= 1;
76                }
77                self.write_line("end");
78            }
79            Statement::Return(value) => {
80                let line = value.as_ref().map_or_else(
81                    || "return".to_string(),
82                    |value| format!("return {}", self.generate_expression(value)),
83                );
84                self.write_line(&line);
85            }
86            Statement::Expr(expression) => {
87                self.write_line(&self.generate_expression(expression));
88            }
89            Statement::ForIn { var, iter, body } => {
90                self.loop_depth += 1;
91                let depth = self.loop_depth;
92                let iter = self.generate_expression(iter);
93                self.write_line(&format!("for _, {var} in pairs({iter}) do"));
94                self.indent_level += 1;
95                for stmt in body {
96                    self.generate_statement(stmt, module, module_name, Scope::Private)?;
97                }
98                if body_has_continue(body) {
99                    self.write_line(&format!("::__continue_{depth}::"));
100                }
101                self.indent_level -= 1;
102                self.write_line("end");
103                self.loop_depth -= 1;
104            }
105            Statement::While { condition, body } => {
106                self.loop_depth += 1;
107                let depth = self.loop_depth;
108                let condition = self.generate_expression(condition);
109                self.write_line(&format!("while {condition} do"));
110                self.indent_level += 1;
111                for stmt in body {
112                    self.generate_statement(stmt, module, module_name, Scope::Private)?;
113                }
114                if body_has_continue(body) {
115                    self.write_line(&format!("::__continue_{depth}::"));
116                }
117                self.indent_level -= 1;
118                self.write_line("end");
119                self.loop_depth -= 1;
120            }
121            Statement::Continue => {
122                self.write_line(&format!("goto __continue_{}", self.loop_depth));
123            }
124            Statement::Break => {
125                self.write_line("break");
126            }
127        }
128
129        Ok(())
130    }
131}