Extract generations methods in fake aggregator build script to a dedicated crate

Author: Damien LACHAUME / PALO-IT

Cargo.lock

@@ -8,6 +8,7 @@ members = [
   "examples/client-mithril-stake-distribution",
   "examples/client-snapshot",
   "mithril-aggregator",
+  "mithril-build-script",
   "mithril-client",
   "mithril-client-cli",
   "mithril-client-wasm",

Cargo.toml

@@ -0,0 +1,3 @@
+target/
+.DS_Store
+

mithril-build-script/.gitignore

@@ -0,0 +1,10 @@
+[package]
+name = "mithril-build-script"
+version = "0.1.0"
+authors = { workspace = true }
+edition = { workspace = true }
+homepage = { workspace = true }
+license = { workspace = true }
+repository = { workspace = true }
+
+[dependencies]

mithril-build-script/Cargo.toml

@@ -0,0 +1,234 @@
+use std::collections::BTreeMap;
+use std::collections::BTreeSet;
+use std::fmt::Write as _;
+use std::fs;
+use std::path::Path;
+
+pub type ArtifactId = String;
+pub type FileContent = String;
+
+/// In memory representation of a folder containing data imported using the `scripts/import.sh` script
+/// of the fake aggregator.
+///
+/// List items are just their corresponding file content loaded in memory.
+/// Individual items are btreemap with the source filename as key and the file content as value.
+#[derive(Debug, Default)]
+pub struct DataFolder {
+    epoch_settings: FileContent,
+
+    certificates_list: FileContent,
+    individual_certificates: BTreeMap<ArtifactId, FileContent>,
+
+    snapshots_list: FileContent,
+    individual_snapshots: BTreeMap<ArtifactId, FileContent>,
+
+    msds_list: FileContent,
+    individual_msds: BTreeMap<ArtifactId, FileContent>,
+
+    ctx_commitments_list: FileContent,
+    individual_ctx_commitments: BTreeMap<ArtifactId, FileContent>,
+    ctx_proofs: BTreeMap<ArtifactId, FileContent>,
+}
+
+impl DataFolder {
+    pub fn load_from_folder(folder: &Path) -> Self {
+        fn extract_artifact_id(filename: &str, prefix: &str) -> String {
+            let id_with_extension = filename.strip_prefix(prefix).unwrap();
+            id_with_extension.strip_suffix(".json").unwrap().to_string()
+        }
+
+        let mut data_folder = DataFolder::default();
+
+        for entry in list_json_files_in_folder(folder) {
+            let filename = entry.file_name().to_string_lossy().to_string();
+            let file_content = fs::read_to_string(&entry.path()).unwrap_or_else(|_| {
+                panic!(
+                    "Could not read file content, file_path: {}",
+                    entry.path().display()
+                )
+            });
+
+            match filename.as_str() {
+                "epoch-settings.json" => {
+                    data_folder.epoch_settings = file_content;
+                }
+                "mithril-stake-distributions.json" => {
+                    data_folder.msds_list = file_content;
+                }
+                "snapshots.json" => {
+                    data_folder.snapshots_list = file_content;
+                }
+                "certificates.json" => {
+                    data_folder.certificates_list = file_content;
+                }
+                "ctx-commitments.json" => {
+                    data_folder.ctx_commitments_list = file_content;
+                }
+                _ if filename.starts_with("mithril-stake-distribution-") => {
+                    data_folder.individual_msds.insert(
+                        extract_artifact_id(&filename, "mithril-stake-distribution-"),
+                        file_content,
+                    );
+                }
+                _ if filename.starts_with("snapshot-") => {
+                    data_folder
+                        .individual_snapshots
+                        .insert(extract_artifact_id(&filename, "snapshot-"), file_content);
+                }
+                _ if filename.starts_with("certificate-") => {
+                    data_folder
+                        .individual_certificates
+                        .insert(extract_artifact_id(&filename, "certificate-"), file_content);
+                }
+                _ if filename.starts_with("ctx-commitment-") => {
+                    data_folder.individual_ctx_commitments.insert(
+                        extract_artifact_id(&filename, "ctx-commitment-"),
+                        file_content,
+                    );
+                }
+                _ if filename.starts_with("ctx-proof-") => {
+                    data_folder
+                        .ctx_proofs
+                        .insert(extract_artifact_id(&filename, "ctx-proof-"), file_content);
+                }
+                // unknown file
+                _ => {}
+            }
+        }
+
+        data_folder
+    }
+
+    pub fn generate_code_for_all_data(self) -> String {
+        format!(
+            "use std::collections::BTreeMap;
+
+{}
+",
+            [
+                generate_list_getter("epoch_settings", self.epoch_settings),
+                generate_ids_array(
+                    "snapshot_digests",
+                    BTreeSet::from_iter(self.individual_snapshots.keys().cloned())
+                ),
+                generate_artifact_getter("snapshots", self.individual_snapshots),
+                generate_list_getter("snapshot_list", self.snapshots_list),
+                generate_ids_array(
+                    "msd_hashes",
+                    BTreeSet::from_iter(self.individual_msds.keys().cloned())
+                ),
+                generate_artifact_getter("msds", self.individual_msds),
+                generate_list_getter("msd_list", self.msds_list),
+                generate_ids_array(
+                    "certificate_hashes",
+                    BTreeSet::from_iter(self.individual_certificates.keys().cloned())
+                ),
+                generate_artifact_getter("certificates", self.individual_certificates),
+                generate_list_getter("certificate_list", self.certificates_list),
+                generate_ids_array(
+                    "ctx_commitment_hashes",
+                    BTreeSet::from_iter(self.individual_ctx_commitments.keys().cloned())
+                ),
+                generate_artifact_getter("ctx_commitments", self.individual_ctx_commitments),
+                generate_list_getter("ctx_commitments_list", self.ctx_commitments_list),
+                generate_ids_array(
+                    "proof_transaction_hashes",
+                    BTreeSet::from_iter(self.ctx_proofs.keys().cloned())
+                ),
+                generate_artifact_getter("ctx_proofs", self.ctx_proofs),
+            ]
+            .join(
+                "
+
+"
+            )
+        )
+    }
+}
+
+pub fn list_json_files_in_folder(folder: &Path) -> impl Iterator<Item = std::fs::DirEntry> + '_ {
+    fs::read_dir(folder)
+        .unwrap_or_else(|_| panic!("Could not read `{}` dir", folder.display()))
+        .filter_map(move |e| {
+            let entry = e.unwrap_or_else(|_| {
+                panic!("Failed to read a file in the `{}` dir", folder.display())
+            });
+            match entry.file_type() {
+                Ok(file_type) if file_type.is_file() => Some(entry),
+                _ => None,
+            }
+        })
+        .filter(|e| e.file_name().to_string_lossy().ends_with(".json"))
+}
+
+// pub(crate) fn $fun_name()() -> BTreeMap<String, String>
+pub fn generate_artifact_getter(
+    fun_name: &str,
+    source_jsons: BTreeMap<ArtifactId, FileContent>,
+) -> String {
+    let mut artifacts_list = String::new();
+
+    for (artifact_id, file_content) in source_jsons {
+        write!(
+            artifacts_list,
+            r###"
+        (
+            "{}",
+            r#"{}"#
+        ),"###,
+            artifact_id, file_content
+        )
+        .unwrap();
+    }
+
+    format!(
+        r###"pub(crate) fn {}() -> BTreeMap<String, String> {{
+    [{}
+    ]
+    .into_iter()
+    .map(|(k, v)| (k.to_owned(), v.to_owned()))
+    .collect()
+}}"###,
+        fun_name, artifacts_list
+    )
+}
+
+/// pub(crate) fn $fun_name() -> &'static str
+pub fn generate_list_getter(fun_name: &str, source_json: FileContent) -> String {
+    format!(
+        r###"pub(crate) fn {}() -> &'static str {{
+    r#"{}"#
+}}"###,
+        fun_name, source_json
+    )
+}
+
+/// pub(crate) fn $fun_name() -> &'static str
+pub fn generate_ids_array(array_name: &str, ids: BTreeSet<ArtifactId>) -> String {
+    let mut ids_list = String::new();
+
+    for id in &ids {
+        write!(
+            ids_list,
+            r#"
+        "{}","#,
+            id
+        )
+        .unwrap();
+    }
+
+    format!(
+        r###"pub(crate) const fn {}<'a>() -> [&'a str; {}] {{
+    [{}
+    ]
+}}"###,
+        array_name,
+        ids.len(),
+        ids_list,
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+}

mithril-build-script/src/fake_aggregator.rs

@@ -0,0 +1,60 @@
+use std::path::PathBuf;
+
+pub mod fake_aggregator;
+
+pub fn get_package_path(package_name: &str) -> PathBuf {
+    let cargo_pkgid_output = std::process::Command::new(env!("CARGO"))
+        .args(["pkgid", "--quiet", "-p", package_name])
+        .output()
+        .unwrap();
+
+    match cargo_pkgid_output.status.success() {
+        true => {
+            let package_name = std::str::from_utf8(&cargo_pkgid_output.stdout)
+                .unwrap()
+                .trim()
+                .strip_prefix("path+file://")
+                .unwrap()
+                .split('#')
+                .collect::<Vec<_>>()[0];
+
+            PathBuf::from(package_name)
+        }
+        false => {
+            panic!(
+                "cargo pkgid failed: stderr: {}",
+                std::str::from_utf8(&cargo_pkgid_output.stderr)
+                    .unwrap()
+                    .trim()
+            )
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn get_package_path_should_return_path_of_existing_package() {
+        let expected = PathBuf::from("./../mithril-aggregator/")
+            .canonicalize()
+            .unwrap();
+
+        let package_path = get_package_path("mithril-aggregator");
+
+        assert_eq!(package_path, expected);
+    }
+
+    #[test]
+    #[should_panic]
+    fn get_package_path_panic_if_valid_name_of_not_existing_package() {
+        get_package_path("it-does-not-exist");
+    }
+
+    #[test]
+    #[should_panic]
+    fn get_package_path_panic_if_invalid_package_name() {
+        get_package_path("Invalid Package Name ~~~");
+    }
+}

mithril-build-script/src/lib.rs

@@ -33,4 +33,4 @@ reqwest = "0.11.23"
 warp = "0.3.6"
 
 [build-dependencies]
-serde_json = "1.0.113"
+mithril-build-script = { path = "../../mithril-build-script" }