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, .. }
14        | Statement::ForNumeric { body, .. }
15        | Statement::While { body, .. } => body_has_continue(body),
16        _ => false,
17    })
18}
19
20impl LuaGenerator {
21    /// Generate code for a given [`Statement`].
22    #[allow(clippy::too_many_lines)]
23    pub(crate) fn generate_statement(
24        &mut self,
25        statement: &Statement,
26        module: Option<&Module>,
27        module_name: Option<&str>,
28        scope: Scope,
29    ) -> LuaGeneratorResult<()> {
30        match statement {
31            Statement::FunctionDecl(function) => {
32                self.generate_function(function, module, scope, module_name)?;
33            }
34            Statement::StructDecl(struct_decl) => {
35                self.generate_struct(struct_decl, module, scope, module_name)?;
36            }
37            Statement::EnumDecl(enum_decl) => {
38                self.generate_enum(enum_decl, module, scope, module_name)?;
39            }
40            Statement::VariableDecl {
41                name,
42                value,
43                source_type,
44                ..
45            } => {
46                let value = self.generate_expression(value);
47                let type_comment = self.variable_type_comment(source_type.as_deref());
48                let line = match (scope, module_name) {
49                    (Scope::Public, Some(module_name)) => {
50                        format!("{module_name}.{name}{type_comment} = {value}")
51                    }
52                    _ => format!("local {name}{type_comment} = {value}"),
53                };
54                self.write_line(&line);
55            }
56            Statement::Assignment { target, value } => {
57                let target = self.generate_expression(target);
58                let value = self.generate_expression(value);
59                self.write_line(&format!("{target} = {value}"));
60            }
61            Statement::Conditional {
62                condition,
63                then_block,
64                else_block,
65            } => {
66                let condition = self.generate_expression(condition);
67                self.write_line(&format!("if {condition} then"));
68
69                self.indent_level += 1;
70                for statement in then_block {
71                    self.generate_statement(statement, module, module_name, Scope::Private)?;
72                }
73                self.indent_level -= 1;
74
75                if !else_block.is_empty() {
76                    self.write_line("else");
77                    self.indent_level += 1;
78                    for statement in else_block {
79                        self.generate_statement(statement, module, module_name, Scope::Private)?;
80                    }
81                    self.indent_level -= 1;
82                }
83                self.write_line("end");
84            }
85            Statement::Return(value) => {
86                let line = value.as_ref().map_or_else(
87                    || "return".to_string(),
88                    |value| format!("return {}", self.generate_expression(value)),
89                );
90                self.write_line(&line);
91            }
92            Statement::Expr(expression) => {
93                self.write_line(&self.generate_expression(expression));
94            }
95            Statement::ForIn {
96                var,
97                iter,
98                body,
99                ipairs,
100            } => {
101                self.loop_depth += 1;
102                let depth = self.loop_depth;
103                let iter = self.generate_expression(iter);
104                let iterator = if *ipairs { "ipairs" } else { "pairs" };
105                self.write_line(&format!("for _, {var} in {iterator}({iter}) do"));
106                self.indent_level += 1;
107                for stmt in body {
108                    self.generate_statement(stmt, module, module_name, Scope::Private)?;
109                }
110                if body_has_continue(body) {
111                    self.write_line(&format!("::__continue_{depth}::"));
112                }
113                self.indent_level -= 1;
114                self.write_line("end");
115                self.loop_depth -= 1;
116            }
117            Statement::ForNumeric {
118                var,
119                start,
120                limit,
121                body,
122            } => {
123                self.loop_depth += 1;
124                let depth = self.loop_depth;
125                let start = self.generate_expression(start);
126                let limit = self.generate_expression(limit);
127                self.write_line(&format!("for {var} = {start}, {limit} do"));
128                self.indent_level += 1;
129                for stmt in body {
130                    self.generate_statement(stmt, module, module_name, Scope::Private)?;
131                }
132                if body_has_continue(body) {
133                    self.write_line(&format!("::__continue_{depth}::"));
134                }
135                self.indent_level -= 1;
136                self.write_line("end");
137                self.loop_depth -= 1;
138            }
139            Statement::While { condition, body } => {
140                self.loop_depth += 1;
141                let depth = self.loop_depth;
142                let condition = self.generate_expression(condition);
143                self.write_line(&format!("while {condition} do"));
144                self.indent_level += 1;
145                for stmt in body {
146                    self.generate_statement(stmt, module, module_name, Scope::Private)?;
147                }
148                if body_has_continue(body) {
149                    self.write_line(&format!("::__continue_{depth}::"));
150                }
151                self.indent_level -= 1;
152                self.write_line("end");
153                self.loop_depth -= 1;
154            }
155            Statement::Continue => {
156                self.write_line(&format!("goto __continue_{}", self.loop_depth));
157            }
158            Statement::Break => {
159                self.write_line("break");
160            }
161        }
162
163        Ok(())
164    }
165}