Skip to main content

gmt_fem_code_builder/
lib.rs

1/*!
2# GMT FEM Code Generator
3
4Generate code based on the input and output tables of the GMT Finite Element Model (FEM)
5
6The crate is used by other crates to generate interfaces to the FEM inputs and outputs.
7
8[gmt-fem](https://crates.io/crates/gmt-fem) build script calls upon [generate_fem] to generate the
9FEM [Inputs](https://docs.rs/gmt-fem/latest/gmt_fem/fem_io/enum.Inputs.html) and [Outputs](https://docs.rs/gmt-fem/latest/gmt_fem/fem_io/enum.Outputs.html) enums.
10
11[gmt_dos-clients_io](https://crates.io/crates/gmt_dos-clients_io) build script invokes [generate_io] to generate the UIDs `Enums` that matches the [inputs and outputs](https://docs.rs/gmt_dos-clients_io/latest/gmt_dos_clients_io/gmt_fem/index.html) of the FEM.
12
13[gmt_dos-clients_fem](https://crates.io/crates/gmt_dos-clients_fem) build script uses [generate_interface] to generate the traits implementation for the FEM inputs and outputs `Enum`s in [gmt-dos-clients_io](https://docs.rs/gmt_dos-clients_io/latest/gmt_dos_clients_io/gmt_fem/index.html).
14
15Invoking [rustc_config] in a build script retrieves the FEM inputs and outputs from the inputs and outputs tables using [io_names], and creates some compilation flags to enable conditional compilation according to the availability of some inputs or outputs.
16
17The compilation flags that [rustc_config] creates are:
18 - `mount`
19 - `m1`
20 - `m1_hp_force_extension`
21 - `m2`
22 - `top-end="ASM"`
23 - `top-end="FSM"`
24 - `m2_rbm="MCM2Lcl6D"`
25 - `m2_rbm="MCM2Lcl"`
26 - `cfd2021`
27 - `cfd2025`
28 - `ground_acceleration`
29
30The full path to the FEM data **must be** set to the environment variable `FEM_REPO`.
31*/
32
33use std::{
34    env,
35    fs::{self, File},
36    io::Read,
37    path::Path,
38};
39
40use apache_arrow::{
41    self as arrow,
42    array::{LargeStringArray, StringArray},
43    record_batch::RecordBatchReader,
44};
45use bytes::Bytes;
46use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
47use zip::ZipArchive;
48
49#[derive(thiserror::Error, Debug)]
50pub enum Error {
51    #[error("No suitable record in file")]
52    NoRecord,
53    #[error("No suitable data in file")]
54    NoData,
55    #[error("Cannot read arrow table")]
56    ReadArrow(#[from] arrow::error::ArrowError),
57    #[error("Cannot read parquet file")]
58    ReadParquet(#[from] parquet::errors::ParquetError),
59    #[error("Cannot find archive in zip file")]
60    Zip(#[from] zip::result::ZipError),
61    #[error("Cannot read zip file content")]
62    ReadZip(#[from] std::io::Error),
63}
64
65mod names;
66pub use names::{Name, Names};
67mod io;
68pub(crate) use io::IO;
69mod get_io;
70pub(crate) use get_io::GetIO;
71
72use apache_arrow::datatypes::Schema;
73use apache_arrow::record_batch::RecordBatch;
74use std::sync::Arc;
75
76// Parse the Arrow table
77fn get_data(
78    field: &str,
79    fem_io: &str,
80    schema: Arc<Schema>,
81    table: &RecordBatch,
82) -> Option<Vec<String>> {
83    let (idx, _) = schema.column_with_name(field).expect(&format!(
84        r#"failed to get {}puts "{}" index with field:\n{:}"#,
85        fem_io,
86        field,
87        schema.field_with_name(field).unwrap()
88    ));
89    match schema.field_with_name(field).unwrap().data_type() {
90        arrow::datatypes::DataType::Utf8 => table
91            .column(idx)
92            .as_any()
93            .downcast_ref::<StringArray>()
94            .expect(&format!(
95                r#"failed to get {}puts "group" data at index #{} from field\n{:}"#,
96                fem_io,
97                idx,
98                schema.field_with_name("group").unwrap()
99            ))
100            .iter()
101            .map(|x| x.map(|x| x.to_owned()))
102            .collect(),
103        arrow::datatypes::DataType::LargeUtf8 => table
104            .column(idx)
105            .as_any()
106            .downcast_ref::<LargeStringArray>()
107            .expect(&format!(
108                r#"failed to get {}puts "group" data at index #{} from field\n{:}"#,
109                fem_io,
110                idx,
111                schema.field_with_name("group").unwrap()
112            ))
113            .iter()
114            .map(|x| x.map(|x| x.to_owned()))
115            .collect(),
116        other => panic!(
117            r#"Expected "Uft8" or "LargeUtf8" datatype, found {}"#,
118            other
119        ),
120    }
121}
122
123// Read the fields
124fn get_fem_io(zip_file: &mut ZipArchive<File>, fem_io: &str) -> Result<Names, Error> {
125    println!("FEM_{}PUTS", fem_io.to_uppercase());
126    let Ok(mut input_file) = zip_file.by_name(&format!(
127        "rust/modal_state_space_model_2ndOrder_{}.parquet",
128        fem_io
129    )) else {
130        panic!(
131            r#"cannot find "rust/modal_state_space_model_2ndOrder_{}.parquet" in archive"#,
132            fem_io
133        )
134    };
135    let mut contents: Vec<u8> = Vec::new();
136    input_file.read_to_end(&mut contents)?;
137
138    let Ok(parquet_reader) = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(contents)) else {
139        panic!("failed to create `ParquetRecordBatchReaderBuilder`")
140    };
141    let Ok(parquet_reader) = parquet_reader.with_batch_size(2048).build() else {
142        panic!("failed to create `ParquetRecordBatchReader`")
143    };
144    let schema = parquet_reader.schema();
145
146    parquet_reader
147        .map(|maybe_table| {
148            if let Ok(table) = maybe_table {
149                get_data("group", fem_io, schema.clone(), &table)
150                    .zip(get_data("description", fem_io, schema.clone(), &table))
151                    .ok_or(Error::NoData)
152            } else {
153                Err(Error::NoRecord)
154            }
155        })
156        .collect::<Result<Vec<_>, Error>>()
157        .map(|data| {
158            let (n, d): (Vec<_>, Vec<_>) = data.into_iter().unzip();
159            let n: Vec<_> = n.into_iter().flatten().collect();
160            let d: Vec<_> = d.into_iter().flatten().collect();
161            (n, d)
162        })
163        // .map(|data| data.into_iter().flatten().collect::<Vec<_>>())
164        .map(|data| {
165            let (name, description) = data;
166            let mut data_iter = name.into_iter();
167            let mut description_iter = description.into_iter();
168            let mut name = data_iter.next().unwrap();
169            let mut names: Vec<Name> = vec![Name::from(&name)];
170            names
171                .last_mut()
172                .map(|name| name.push_description(description_iter.next().unwrap()));
173            loop {
174                match data_iter.next() {
175                    Some(data) if data == name => {
176                        names
177                            .last_mut()
178                            .map(|name| name.push_description(description_iter.next().unwrap()));
179                    }
180                    Some(data) => {
181                        name = data;
182                        names.push(name.as_str().into());
183                        names
184                            .last_mut()
185                            .map(|name| name.push_description(description_iter.next().unwrap()));
186                    }
187                    None => break,
188                }
189            }
190            names.into_iter().collect()
191        })
192}
193
194/// Returns the list of inputs and outputs of the FEM
195pub fn io_names(from_crate: &str) -> std::result::Result<(Names, Names), Error> {
196    Ok(if let Ok(fem_repo) = env::var("FEM_REPO") {
197        // Gets the FEM repository
198        println!(
199            "cargo:warning={}: generating FEM/Actors interface code based on the FEM inputs and outputs tables in {}",
200            from_crate, fem_repo
201        );
202        // Opens the mat file
203        let path = Path::new(&fem_repo);
204        let Ok(file) = File::open(path.join("modal_state_space_model_2ndOrder.zip")) else {
205            panic!("Cannot find `modal_state_space_model_2ndOrder.zip` in `FEM_REPO`");
206        };
207        let mut zip_file = zip::ZipArchive::new(file)?;
208
209        let Ok(input_names) = get_fem_io(&mut zip_file, "in") else {
210            panic!("failed to parse FEM inputs variables")
211        };
212        let Ok(output_names) = get_fem_io(&mut zip_file, "out") else {
213            panic!("failed to parse FEM outputs variables")
214        };
215        (input_names, output_names)
216    } else {
217        println!(
218            "cargo:warning=the FEM_REPO environment variable is not set, using dummy inputs and outputs instead"
219        );
220        let (inputs, outputs): (Vec<_>, Vec<_>) = (1..=5)
221            .map(|i| {
222                (
223                    String::from(format!("In{i}")),
224                    String::from(format!("Out{i}")),
225                )
226            })
227            .unzip();
228        (inputs.into_iter().collect(), outputs.into_iter().collect())
229    })
230}
231
232/// Generate the code for [gmt_dos-clients_fem](https://crates.io/crates/gmt_dos-clients_fem) interfaces
233pub fn generate_interface(from_crate: &str) -> anyhow::Result<()> {
234    let (input_names, output_names): (Names, Names) = io_names(from_crate)?;
235
236    let out_dir = env::var_os("OUT_DIR").unwrap();
237    let dest_path = Path::new(&out_dir);
238
239    fs::write(
240        dest_path.join("fem_get_in.rs"),
241        format!("{}", GetIO::new("In", &input_names)),
242    )?;
243    fs::write(
244        dest_path.join("fem_get_out.rs"),
245        format!("{}", GetIO::new("Out", &output_names)),
246    )?;
247
248    fs::write(
249        dest_path.join("fem_inputs.rs"),
250        input_names
251            .iter()
252            .map(|name| format!("{}", name.impl_enum_variant_for_io("Inputs")))
253            .collect::<Vec<String>>()
254            .join("\n"),
255    )?;
256    fs::write(
257        dest_path.join("fem_outputs.rs"),
258        output_names
259            .iter()
260            .map(|name| format!("{}", name.impl_enum_variant_for_io("Outputs")))
261            .collect::<Vec<String>>()
262            .join("\n"),
263    )?;
264
265    rustc_config(from_crate, Some((input_names, output_names)))?;
266
267    println!("cargo:rerun-if-env-changed=FEM_REPO");
268    Ok(())
269}
270
271/// Generate the list of inputs and outputs of the FEM as UIDs `Enum` in [gmt_dos-clients_io](https://crates.io/crates/gmt_dos-clients_io)
272pub fn generate_io(from_crate: &str) -> anyhow::Result<()> {
273    let (input_names, output_names): (Names, Names) = io_names(from_crate)?;
274
275    let out_dir = env::var_os("OUT_DIR").unwrap();
276    let dest_path = Path::new(&out_dir);
277
278    fs::write(
279        dest_path.join("fem_actors_inputs.rs"),
280        format!("{}", input_names),
281    )?;
282    fs::write(
283        dest_path.join("fem_actors_outputs.rs"),
284        format!("{}", output_names),
285    )?;
286
287    println!("cargo:rerun-if-env-changed=FEM_REPO");
288    Ok(())
289}
290
291/**
292Creates rustc compilation flags
293
294The compilations flags are created according to the availability of some inputs and outputs.
295If no inputs and outputs are given, then [io_names] is used to retrieve all of them.
296
297The compilation flags that [rustc_config] creates are:
298 - `mount`
299 - `m1`
300 - `m1_hp_force_extension`
301 - `m2`
302 - `top-end="ASM"`
303 - `top-end="FSM"`
304 - `m2_rbm="MCM2Lcl6D"`
305 - `m2_rbm="MCM2Lcl"`
306 - `cfd2021`
307 - `cfd2025`
308 - `ground_acceleration`
309*/
310pub fn rustc_config(from_crate: &str, io: Option<(Names, Names)>) -> anyhow::Result<()> {
311    if option_env!("FEM_REPO").is_some() {
312        println!("cargo::rustc-cfg=fem");
313        let (input_names, output_names): (Names, Names) = match io {
314            Some(io) => Ok(io),
315            None => io_names(from_crate),
316        }?;
317        if input_names.find("MCM2S1VCDeltaF").is_some() {
318            println!("cargo::warning={}: ASM top-end", from_crate);
319            println!(r#"cargo::rustc-cfg=topend="ASM""#)
320        }
321        if input_names.find("MCM2PZTF").is_some() {
322            println!("cargo::warning={}: FSM top-end", from_crate);
323            println!(r#"cargo::rustc-cfg=topend="FSM""#);
324        }
325        match (
326            input_names.find("MCM2S1VCDeltaF"),
327            input_names.find("MCM2PZTF"),
328            input_names.find("MCM2SmHexF"),
329        ) {
330            (Some(_), None, Some(_)) => {
331                println!("cargo::warning={}: ASMS inputs", from_crate);
332                println!(r#"cargo::rustc-cfg=m2"#)
333            }
334            (None, Some(_), Some(_)) => {
335                println!("cargo::warning={}: FSMS inputs", from_crate);
336                println!(r#"cargo::rustc-cfg=m2"#)
337            }
338            _ => (),
339        };
340        if input_names.find("CFD2021106F").is_some() {
341            println!("cargo::warning={}: 2021 CFD inputs", from_crate);
342            println!(r#"cargo::rustc-cfg=cfd2021"#)
343        }
344        if input_names.find("CFD2025046F").is_some() {
345            println!("cargo::warning={}: 2025 CFD inputs", from_crate);
346            println!(r#"cargo::rustc-cfg=cfd2025"#)
347        }
348        match (
349            input_names.find("OSSAzDriveTorque"),
350            input_names.find("OSSElDriveTorque"),
351            input_names.find("OSSRotDriveTorque"),
352            output_names.find("OSSAzEncoderAngle"),
353            output_names.find("OSSElEncoderAngle"),
354            output_names.find("OSSRotEncoderAngle"),
355        ) {
356            (Some(_), Some(_), Some(_), Some(_), Some(_), Some(_)) => {
357                println!("cargo::warning={}: Mount inputs and outputs", from_crate);
358                println!(r#"cargo::rustc-cfg=mount"#)
359            }
360            _ => (),
361        };
362        match (
363            input_names.find("OSSHarpointDeltaF"),
364            input_names.find("M1ActuatorsSegment1"),
365            output_names.find("OSSHardpointD"),
366            output_names.find("OSSM1Lcl"),
367            output_names.find("M1Segment1AxialD"),
368            output_names.find("OSSHardpointForce"),
369            input_names.find("OSSHardpointExtension"),
370        ) {
371            (Some(_), Some(_), Some(_), Some(_), Some(_), None, None) => {
372                println!("cargo::warning={}: M1 inputs and outputs", from_crate);
373                println!(r#"cargo::rustc-cfg=m1"#)
374            }
375            (Some(_), Some(_), Some(_), Some(_), Some(_), Some(_), Some(_)) => {
376                println!(
377                    "cargo::warning={}: M1 inputs and outputs with hardpoints forces and extensions",
378                    from_crate
379                );
380                println!(r#"cargo::rustc-cfg=m1_hp_force_extension"#);
381                println!(r#"cargo::rustc-cfg=m1"#)
382            }
383            _ => (),
384        };
385        if input_names.find("OSS00GroundAcc").is_some() {
386            println!("cargo::warning={}: OSS00GroundAcc input", from_crate);
387            println!(r#"cargo::rustc-cfg=ground_acceleration"#)
388        }
389        if output_names.find("MCM2Lcl6D").is_some() {
390            println!("cargo::warning={}: MCM2Lcl6D as M2 RBM output", from_crate);
391            println!(r#"cargo::rustc-cfg=m2_rbm="MCM2Lcl6D""#)
392        }
393        if output_names.find("MCM2Lcl").is_some() {
394            println!("cargo::warning={}: MCM2Lcl as M2 RBM output", from_crate);
395            println!(r#"cargo::rustc-cfg=m2_rbm="MCM2Lcl""#)
396        }
397    }
398    Ok(())
399}
400
401/// Generate the FEM [Inputs](https://docs.rs/gmt-fem/latest/gmt_fem/fem_io/enum.Inputs.html) and [Outputs](https://docs.rs/gmt-fem/latest/gmt_fem/fem_io/enum.Outputs.html) enums for [gmt-fem](https://crates.io/crates/gmt-fem)
402pub fn generate_fem(from_crate: &str) -> anyhow::Result<()> {
403    let (input_names, output_names): (Names, Names) = io_names(from_crate)?;
404
405    let out_dir = env::var_os("OUT_DIR").unwrap();
406    let dest_path = Path::new(&out_dir);
407
408    fs::write(
409        dest_path.join("fem_inputs.rs"),
410        format!("{}", IO::new("Inputs", &input_names)),
411    )?;
412    fs::write(
413        dest_path.join("fem_outputs.rs"),
414        format!("{}", IO::new("Outputs", &output_names)),
415    )?;
416
417    println!("cargo:rerun-if-env-changed=FEM_REPO");
418
419    Ok(())
420}