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            "local "
59        } else {
60            ""
61        };
62
63        self.write_doc_comments(struct_decl.doc.as_deref());
64
65        if self.debug_level_at_least(0)
66            && let Some(debug) = &struct_decl.debug
67        {
68            self.write_line(&format!("-- {}", debug.header_comment));
69        }
70
71        self.write_line(&format!("{prefix}{table_path} = {{}}"));
72
73        for (name, value) in &struct_decl.constants {
74            let value = self.generate_expression(value);
75            self.write_line(&format!("{table_path}.{name} = {value}"));
76        }
77
78        for method in &struct_decl.methods {
79            self.generate_table_method(method, &struct_decl.name, &table_path)?;
80        }
81
82        Ok(())
83    }
84
85    /// Generate a user-defined tagged-table enum and its method table.
86    pub(crate) fn generate_enum(
87        &mut self,
88        enum_decl: &Enum,
89        _module: Option<&Module>,
90        scope: Scope,
91        module_name: Option<&str>,
92    ) -> LuaGeneratorResult<()> {
93        if scope == Scope::Private && module_name.is_some() {
94            return Err(LuaGeneratorError::StructLocalAndExported(
95                enum_decl.name.clone(),
96            ));
97        }
98
99        let table_path = match (scope, module_name) {
100            (Scope::Public, Some(module_name)) => format!("{module_name}.{}", enum_decl.name),
101            (Scope::Private | Scope::Public, None) => enum_decl.name.clone(),
102            (Scope::Private, Some(_)) => unreachable!(),
103        };
104        let prefix = if scope == Scope::Private && module_name.is_none() {
105            "local "
106        } else {
107            ""
108        };
109
110        self.write_doc_comments(enum_decl.doc.as_deref());
111        self.write_line(&format!("{prefix}{table_path} = {{}}"));
112
113        for variant in &enum_decl.variants {
114            if matches!(variant.fields, EnumVariantFields::Unit) {
115                let value = self.generate_expression(&Expression::EnumLiteral {
116                    enum_name: enum_decl.name.clone(),
117                    variant: variant.name.clone(),
118                    fields: vec![],
119                });
120                self.write_line(&format!("{table_path}.{} = {value}", variant.name));
121            }
122        }
123        for (name, value) in &enum_decl.constants {
124            let value = self.generate_expression(value);
125            self.write_line(&format!("{table_path}.{name} = {value}"));
126        }
127        for method in &enum_decl.methods {
128            self.generate_table_method(method, &enum_decl.name, &table_path)?;
129        }
130        Ok(())
131    }
132
133    /// Generate a method found on a `struct`
134    pub(crate) fn generate_table_method(
135        &mut self,
136        func: &Function,
137        struct_name: &str,
138        table_path: &str,
139    ) -> LuaGeneratorResult<()> {
140        let function_uses_self = func
141            .params
142            .first()
143            .is_some_and(|parameter| parameter.name == "self");
144
145        let params = if function_uses_self {
146            func.params
147                .iter()
148                .skip(1)
149                .map(|parameter| self.format_parameter(parameter))
150                .collect::<Vec<_>>()
151                .join(", ")
152        } else {
153            func.params
154                .iter()
155                .map(|parameter| self.format_parameter(parameter))
156                .collect::<Vec<_>>()
157                .join(", ")
158        };
159
160        let return_comment = self.function_return_comment(
161            func.debug
162                .as_ref()
163                .and_then(|debug| debug.return_type.as_deref()),
164        );
165        let separator = if function_uses_self { ":" } else { "." };
166
167        self.write_doc_comments(func.doc.as_deref());
168
169        if self.debug_level_at_least(0)
170            && let Some(debug) = &func.debug
171        {
172            self.write_line(&format!("-- {}", debug.header_comment));
173        }
174
175        self.write_line(&format!(
176            "function {table_path}{separator}{}({params}){return_comment}",
177            func.name
178        ));
179
180        self.struct_table_context = Some((struct_name.to_string(), table_path.to_string()));
181        self.indent_level += 1;
182        self.generate_block(&func.body, None)?;
183        self.indent_level -= 1;
184        self.struct_table_context = None;
185
186        self.write_line("end");
187
188        Ok(())
189    }
190}