Skip to main content

factorio_codegen/generator/
table.rs

1use factorio_ir::{
2    enumeration::{Enum, EnumVariantFields},
3    expression::Expression,
4    function::Function,
5    module::Module,
6    scope::Scope,
7    structure::Struct,
8};
9
10use crate::{LuaGenerator, LuaGeneratorError, LuaGeneratorResult};
11
12impl LuaGenerator {
13    /// Generate a `struct`.
14    pub(crate) fn generate_struct(
15        &mut self,
16        struct_decl: &Struct,
17        module: Option<&Module>,
18        scope: Scope,
19        module_name: Option<&str>,
20    ) -> LuaGeneratorResult<()> {
21        if scope == Scope::Private && module_name.is_some() {
22            return Err(LuaGeneratorError::StructLocalAndExported(
23                struct_decl.name.clone(),
24            ));
25        }
26
27        if let Some(module) = module
28            && module.is_imported_type_extension(struct_decl)
29        {
30            let table_path = module
31                .imported_item_local(&struct_decl.name)
32                .ok_or_else(|| {
33                    LuaGeneratorError::FailedToGetTablePathForStruct(struct_decl.name.clone())
34                })?
35                .to_string();
36
37            for (name, value) in &struct_decl.constants {
38                let value = self.generate_expression(value);
39                self.write_line(&format!("{table_path}.{name} = {value}"));
40            }
41
42            for method in &struct_decl.methods {
43                self.generate_table_method(method, &struct_decl.name, &table_path)?;
44            }
45
46            return Ok(());
47        }
48
49        let table_path = match (scope, module_name) {
50            (Scope::Public, Some(module_name)) => {
51                format!("{module_name}.{}", struct_decl.name)
52            }
53            (Scope::Private | Scope::Public, None) => struct_decl.name.clone(),
54            (Scope::Private, Some(_)) => unreachable!(),
55        };
56
57        let prefix = if scope == Scope::Private && module_name.is_none() {
58            if self.forward_declared_locals.contains(&struct_decl.name) {
59                ""
60            } else {
61                "local "
62            }
63        } else {
64            ""
65        };
66
67        self.write_doc_comments(struct_decl.doc.as_deref());
68
69        if self.debug_level_at_least(0)
70            && let Some(debug) = &struct_decl.debug
71        {
72            self.write_line(&format!("-- {}", debug.header_comment));
73        }
74
75        self.write_line(&format!("{prefix}{table_path} = {{}}"));
76
77        for (name, value) in &struct_decl.constants {
78            let value = self.generate_expression(value);
79            self.write_line(&format!("{table_path}.{name} = {value}"));
80        }
81
82        for method in &struct_decl.methods {
83            self.generate_table_method(method, &struct_decl.name, &table_path)?;
84        }
85
86        Ok(())
87    }
88
89    /// Generate a user-defined tagged-table enum and its method table.
90    pub(crate) fn generate_enum(
91        &mut self,
92        enum_decl: &Enum,
93        _module: Option<&Module>,
94        scope: Scope,
95        module_name: Option<&str>,
96    ) -> LuaGeneratorResult<()> {
97        if scope == Scope::Private && module_name.is_some() {
98            return Err(LuaGeneratorError::StructLocalAndExported(
99                enum_decl.name.clone(),
100            ));
101        }
102
103        let table_path = match (scope, module_name) {
104            (Scope::Public, Some(module_name)) => format!("{module_name}.{}", enum_decl.name),
105            (Scope::Private | Scope::Public, None) => enum_decl.name.clone(),
106            (Scope::Private, Some(_)) => unreachable!(),
107        };
108        let prefix = if scope == Scope::Private && module_name.is_none() {
109            if self.forward_declared_locals.contains(&enum_decl.name) {
110                ""
111            } else {
112                "local "
113            }
114        } else {
115            ""
116        };
117
118        self.write_doc_comments(enum_decl.doc.as_deref());
119        self.write_line(&format!("{prefix}{table_path} = {{}}"));
120
121        for variant in &enum_decl.variants {
122            if matches!(variant.fields, EnumVariantFields::Unit) {
123                let value = self.generate_expression(&Expression::EnumLiteral {
124                    enum_name: enum_decl.name.clone(),
125                    variant: variant.name.clone(),
126                    fields: vec![],
127                });
128                self.write_line(&format!("{table_path}.{} = {value}", variant.name));
129            }
130        }
131        for (name, value) in &enum_decl.constants {
132            let value = self.generate_expression(value);
133            self.write_line(&format!("{table_path}.{name} = {value}"));
134        }
135        for method in &enum_decl.methods {
136            self.generate_table_method(method, &enum_decl.name, &table_path)?;
137        }
138        Ok(())
139    }
140
141    /// Generate a method found on a `struct`
142    pub(crate) fn generate_table_method(
143        &mut self,
144        func: &Function,
145        struct_name: &str,
146        table_path: &str,
147    ) -> LuaGeneratorResult<()> {
148        let function_uses_self = func
149            .params
150            .first()
151            .is_some_and(|parameter| parameter.name == "self");
152
153        let params = if function_uses_self {
154            func.params
155                .iter()
156                .skip(1)
157                .map(|parameter| self.format_parameter(parameter))
158                .collect::<Vec<_>>()
159                .join(", ")
160        } else {
161            func.params
162                .iter()
163                .map(|parameter| self.format_parameter(parameter))
164                .collect::<Vec<_>>()
165                .join(", ")
166        };
167
168        let return_comment = self.function_return_comment(
169            func.debug
170                .as_ref()
171                .and_then(|debug| debug.return_type.as_deref()),
172        );
173        let separator = if function_uses_self { ":" } else { "." };
174
175        self.write_doc_comments(func.doc.as_deref());
176
177        if self.debug_level_at_least(0)
178            && let Some(debug) = &func.debug
179        {
180            self.write_line(&format!("-- {}", debug.header_comment));
181        }
182
183        self.write_line(&format!(
184            "function {table_path}{separator}{}({params}){return_comment}",
185            func.name
186        ));
187
188        self.struct_table_context = Some((struct_name.to_string(), table_path.to_string()));
189        self.indent_level += 1;
190        self.generate_block(&func.body, None)?;
191        self.indent_level -= 1;
192        self.struct_table_context = None;
193
194        self.write_line("end");
195
196        Ok(())
197    }
198}