factorio_exporter/
exporter.rs

1use std::{
2    collections::HashMap,
3    fs::File,
4    io::Write,
5    path::Path,
6    process::{Command, Output},
7};
8
9use convert_case::{Case, Casing};
10use indoc::writedoc;
11use itertools::Itertools;
12use regex::Captures;
13use regex_macro::regex;
14use serde_derive::Deserialize;
15use serde_yaml::Value;
16use tempfile::TempDir;
17use tracing::{debug, error, info};
18
19use crate::{
20    api::{Api, HasAttributes},
21    internal::{
22        mod_controller::{ModController, ModManifestBuilder},
23        script_generator::ScriptGenerator,
24    },
25    FactorioExporterError::{self, FactorioExecutionError},
26    Result,
27};
28
29const CONFIG: &str = "config.ini";
30const SAVE: &str = "save.zip";
31const MOD_NAME: &str = "factorio_exporter";
32const MOD_VERSION: &str = "0.0.1";
33const MODS_DIR: &str = "mods";
34
35/// Main class for orchestrating the export.
36pub struct FactorioExporter<'a> {
37    factorio_binary: &'a Path,
38    api: &'a Api,
39    locale: &'a str,
40    temp_dir: TempDir,
41    mod_controller: ModController,
42    export_icons: bool,
43}
44
45impl FactorioExporter<'_> {
46    /// Creates and configures a new `FactorioExporter` instance.
47    ///
48    /// # Arguments
49    ///
50    /// * `factorio_binary` - File system path of a Factorio binary. This can be
51    ///   any variant of the binary, full, headless, or demo.
52    /// * `api` - The definition of the Factorio API, as loaded by
53    ///   [`load_api`](super::load_api).
54    /// * `locale` - Locale code to use for translated strings.
55    /// * `export_icons` - Whether icon paths should be collected in the data
56    ///   phase and patched into the prototype definitions using a heuristic.
57    pub fn new<'a>(
58        factorio_binary: &'a Path,
59        api: &'a Api,
60        locale: &'a str,
61        export_icons: bool,
62    ) -> Result<FactorioExporter<'a>> {
63        let temp_dir = tempfile::Builder::new().prefix(MOD_NAME).tempdir()?;
64        let mod_controller = ModController::new(temp_dir.path().join(MODS_DIR));
65        Ok(FactorioExporter {
66            factorio_binary,
67            api,
68            locale,
69            temp_dir,
70            mod_controller,
71            export_icons,
72        })
73    }
74}
75
76const ARGS: &[&str] = &["--config", CONFIG];
77
78impl FactorioExporter<'_> {
79    /// Export the prototype definitions from Factorio and partially deserialize
80    /// them into a [`serde_yaml::Value`] object, which can easily deserialized
81    /// of serialized into other data types further.
82    ///
83    /// This function executes Factorio twice, once to create a save file, and a
84    /// second time to execute an exporter mod that does the heavy-lifting. The
85    /// process uses a temporary directory, so that the main Factorio
86    /// installation is not touched. Any existing Factorio configuration,
87    /// including installed mods are therefore ignored.
88    pub fn export(&self) -> Result<Value> {
89        self.create_exec_dir()?;
90        self.create_exporter_mod()?;
91
92        info!("create an empty save file");
93        self.run_factorio(&["--create", SAVE, "--mod-directory", "none"])?;
94
95        info!("execute Factorio to export prototypes");
96        #[rustfmt::skip]
97        let output = self.run_factorio(&[
98            "--benchmark", SAVE,
99            "--benchmark-ticks", "1",
100            "--benchmark-runs", "1",
101            "--instrument-mod", MOD_NAME,
102        ])?;
103
104        info!("parse Factorio output");
105        self.parse_output(&String::from_utf8_lossy(&output.stdout))
106    }
107
108    fn create_exec_dir(&self) -> Result<()> {
109        let config = self.temp_dir.path().join(CONFIG);
110
111        debug!("creating config file: {:?}", config);
112        writedoc!(
113            File::create(config)?,
114            r#"
115                [path]
116                read-data=__PATH__executable__/../../data
117                write-data=.
118                [general]
119                locale={}
120            "#,
121            self.locale
122        )?;
123        Ok(())
124    }
125
126    /// Install mods into the temporary execution directory before exporting.
127    /// This allows exporting additional items, recipes, and all other changes
128    /// that the mods make to be part of the export.
129    ///
130    /// No particular checks are made that the dependencies of the specified
131    /// mods can be resolved. This is the responsibility of the caller.
132    /// Otherwise Factorio will probably not start.
133    ///
134    /// # Arguments
135    ///
136    /// * `mods` - A list of file system paths that point to Factorio mods in
137    ///   `.zip` format.
138    pub fn install_mods<I, P>(&self, mods: I) -> Result<()>
139    where
140        I: IntoIterator<Item = P>,
141        P: AsRef<Path>,
142    {
143        info!("installing mods");
144
145        for m in mods {
146            debug!("installing mod: {:?}", m.as_ref());
147            self.mod_controller.add_mod(m.as_ref())?;
148        }
149        Ok(())
150    }
151
152    fn create_exporter_mod(&self) -> Result<()> {
153        let attrs = self.api.classes["LuaGameScript"]
154            .attributes()
155            .iter()
156            .copied()
157            .filter(|attr| attr.name.ends_with("prototypes"))
158            .collect_vec();
159
160        self.mod_controller
161            .create_mod(
162                ModManifestBuilder::default()
163                    .name(MOD_NAME)
164                    .version(MOD_VERSION)
165                    .title("Factorio Exporter")
166                    .author("Michael Forster <email@michael-forster.de")
167                    .build()
168                    .unwrap(),
169            )?
170            .add_file("export.lua", include_str!("../lua/export.lua"))?
171            .add_file(
172                "instrument-after-data.lua",
173                include_str!("../lua/instrument-after-data.lua"),
174            )?
175            .add_file("instrument-control.lua", include_str!("../lua/instrument-control.lua"))?
176            .add_file("prototypes.lua", &ScriptGenerator::new(self.api).generate("game", attrs))?;
177
178        Ok(())
179    }
180
181    fn run_factorio(&self, args: &[&str]) -> Result<Output> {
182        if !self.factorio_binary.is_file() {
183            return Err(FactorioExporterError::FileNotFoundError {
184                file: self.factorio_binary.into(),
185            });
186        }
187
188        let mut binary = Command::new(self.factorio_binary);
189        let command = &mut binary.current_dir(&self.temp_dir).args(ARGS).args(args);
190
191        debug!("executing command: {:?}", command);
192
193        let output = command.output()?;
194
195        if !output.status.success() {
196            let stdout = String::from_utf8_lossy(&output.stdout).to_string();
197            let stderr = String::from_utf8_lossy(&output.stderr).to_string();
198
199            error!("STDOUT\n{}", &stdout);
200            error!("STDERR\n{}", &stderr);
201
202            return Err(FactorioExecutionError { stdout, stderr });
203        }
204
205        Ok(output)
206    }
207
208    fn find_section<'a>(output: &'a str, marker: &str) -> Result<&'a str> {
209        let start_marker = format!("<{marker}>");
210        let start = output.find(&start_marker).ok_or_else(|| {
211            FactorioExporterError::FactorioOutputError {
212                message: format!("Didn't find {start_marker} marker"),
213                output: output.into(),
214            }
215        })?;
216
217        let stop_marker = format!("</{marker}>");
218        let stop = output.find(&stop_marker).ok_or_else(|| {
219            FactorioExporterError::FactorioOutputError {
220                message: format!("Didn't find {stop_marker} marker"),
221                output: output.into(),
222            }
223        })?;
224
225        Ok(&output[start + start_marker.len()..stop])
226    }
227
228    fn parse_output(&self, s: &str) -> Result<Value> {
229        // Unfortunately we have no control over the string printed by Lua's
230        // `localised_print`. There can be single/double quotes or new lines in
231        // there. Neither JSON nor YAML can deal with that well. YAML could if we
232        // had a way to control the indentation, but we don't. So, let's solve it
233        // the hacky way: post-processing.
234
235        debug!("parse prototype output");
236        let re = regex!(r"(?s)<STRING>(.*?)</STRING>");
237        let sanitized = re.replace_all(Self::find_section(s, "EXPORT")?, |caps: &Captures| {
238            format!("'{}'", &caps[1].replace('\n', "\\n").replace('\'', "''"))
239        });
240        let mut data: Value = serde_yaml::from_str::<Value>(&sanitized)?;
241
242        debug!(
243            "found {} items, {} recipes, {} technologies",
244            data["item_prototypes"].as_mapping().unwrap().len(),
245            data["recipe_prototypes"].as_mapping().unwrap().len(),
246            data["technology_prototypes"].as_mapping().unwrap().len()
247        );
248
249        // Icon paths are not available in Factorio's runtime stage, so we must
250        // resort to getting them in the data stage. Unfortunately data
251        // structures in the data stage are a bit messy, so we need to apply
252        // some heuristics to map icons into the prototypes that we get in the
253        // runtime stage. We add an icon property to a prototype if it's `name`
254        // and `object_name` or `type` match the section names and section
255        // element names in `data.raw`.
256
257        if self.export_icons {
258            debug!("parse icons output");
259
260            let icons: Vec<Icon> = serde_yaml::from_str(Self::find_section(s, "ICONS")?)?;
261            let icons: HashMap<(&str, &str), &str> =
262                icons.iter().map(|icon| ((icon.name, icon.section), icon.path)).collect();
263
264            debug!("patch {} icons into prototypes", icons.len());
265
266            let object_name_pattern = regex!("Lua(.*)Prototype");
267            for (_, section) in data.as_mapping_mut().expect("root should be a mapping") {
268                if let Value::Mapping(section) = section {
269                    for (name, el) in section {
270                        let name = name.as_str().expect("key should be a string");
271                        let val = el.as_mapping_mut().expect("value should be mapping");
272
273                        if let Some(path) = val
274                            .get("type")
275                            .and_then(|value| value.as_str())
276                            .and_then(|ty| icons.get(&(name, ty)).map(|r| r.to_string()))
277                            .or_else(|| {
278                                val.get("object_name")
279                                    .and_then(|value| value.as_str())
280                                    .and_then(|name| {
281                                        object_name_pattern
282                                            .captures(name)
283                                            .map(|captures| (&captures[1]).to_case(Case::Kebab))
284                                    })
285                                    .and_then(|ty| icons.get(&(name, &ty)).map(|r| r.to_string()))
286                            })
287                        {
288                            val.insert("icon".into(), (*path).into());
289                        };
290                    }
291                }
292            }
293        }
294
295        Ok(data)
296    }
297}
298
299#[derive(Debug, Deserialize)]
300struct Icon<'a> {
301    section: &'a str,
302    name: &'a str,
303    path: &'a str,
304}