mithril_build_script/
fake_aggregator.rs

1use std::collections::BTreeMap;
2use std::collections::BTreeSet;
3use std::fmt::Write as _;
4use std::fs;
5use std::fs::File;
6use std::path::Path;
7
8use serde_json;
9
10pub type ArtifactId = String;
11pub type FileContent = String;
12
13/// In memory representation of a folder containing data imported using the `scripts/import.sh` script
14/// of the fake aggregator.
15#[derive(Debug, Default)]
16pub struct FakeAggregatorData {
17    status: FileContent,
18
19    epoch_settings: FileContent,
20
21    certificates_list: FileContent,
22    individual_certificates: BTreeMap<ArtifactId, FileContent>,
23
24    snapshots_list: FileContent,
25    individual_snapshots: BTreeMap<ArtifactId, FileContent>,
26
27    mithril_stake_distributions_list: FileContent,
28    individual_mithril_stake_distributions: BTreeMap<ArtifactId, FileContent>,
29
30    cardano_transaction_snapshots_list: FileContent,
31    individual_cardano_transaction_snapshots: BTreeMap<ArtifactId, FileContent>,
32    cardano_transaction_proofs: BTreeMap<ArtifactId, FileContent>,
33
34    cardano_stake_distributions_list: FileContent,
35    individual_cardano_stake_distributions: BTreeMap<ArtifactId, FileContent>,
36
37    cardano_database_snapshots_list: FileContent,
38    individual_cardano_database_snapshots: BTreeMap<ArtifactId, FileContent>,
39}
40
41impl FakeAggregatorData {
42    pub fn load_from_folder(folder: &Path) -> Self {
43        let mut data = FakeAggregatorData::default();
44
45        for entry in list_json_files_in_folder(folder) {
46            let filename = entry.file_name().to_string_lossy().to_string();
47            let file_content = fs::read_to_string(entry.path()).unwrap_or_else(|_| {
48                panic!(
49                    "Could not read file content, file_path: {}",
50                    entry.path().display()
51                )
52            });
53
54            match filename.as_str() {
55                "status.json" => {
56                    data.status = file_content;
57                }
58                "epoch-settings.json" => {
59                    data.epoch_settings = file_content;
60                }
61                "mithril-stake-distributions-list.json" => {
62                    data.mithril_stake_distributions_list = file_content;
63                }
64                "snapshots-list.json" => {
65                    data.snapshots_list = file_content;
66                }
67                "cardano-stake-distributions-list.json" => {
68                    data.cardano_stake_distributions_list = file_content;
69                }
70                "cardano-databases-list.json" => {
71                    data.cardano_database_snapshots_list = file_content;
72                }
73                "certificates-list.json" => {
74                    data.certificates_list = file_content;
75                }
76                "ctx-snapshots-list.json" => {
77                    data.cardano_transaction_snapshots_list = file_content;
78                }
79                "mithril-stake-distributions.json" => {
80                    data.individual_mithril_stake_distributions =
81                        Self::read_artifacts_json_file(&entry.path());
82                }
83                "snapshots.json" => {
84                    data.individual_snapshots = Self::read_artifacts_json_file(&entry.path());
85                }
86                "cardano-stake-distributions.json" => {
87                    data.individual_cardano_stake_distributions =
88                        Self::read_artifacts_json_file(&entry.path());
89                }
90                "cardano-databases.json" => {
91                    data.individual_cardano_database_snapshots =
92                        Self::read_artifacts_json_file(&entry.path());
93                }
94                "certificates.json" => {
95                    data.individual_certificates = Self::read_artifacts_json_file(&entry.path());
96                }
97                "ctx-snapshots.json" => {
98                    data.individual_cardano_transaction_snapshots =
99                        Self::read_artifacts_json_file(&entry.path());
100                }
101                "ctx-proofs.json" => {
102                    data.cardano_transaction_proofs = Self::read_artifacts_json_file(&entry.path());
103                }
104                // unknown file
105                _ => {}
106            }
107        }
108
109        data
110    }
111
112    pub fn generate_code_for_ids(self) -> String {
113        Self::assemble_code(
114            &[
115                generate_ids_array(
116                    "snapshot_digests",
117                    BTreeSet::from_iter(self.individual_snapshots.keys().cloned()),
118                ),
119                generate_ids_array(
120                    "mithril_stake_distribution_hashes",
121                    BTreeSet::from_iter(
122                        self.individual_mithril_stake_distributions.keys().cloned(),
123                    ),
124                ),
125                generate_ids_array(
126                    "cardano_stake_distribution_hashes",
127                    BTreeSet::from_iter(
128                        self.individual_cardano_stake_distributions.keys().cloned(),
129                    ),
130                ),
131                generate_ids_array(
132                    "cardano_stake_distribution_epochs",
133                    BTreeSet::from_iter(extract_cardano_stake_distribution_epochs(
134                        &self.individual_cardano_stake_distributions,
135                    )),
136                ),
137                generate_ids_array(
138                    "cardano_database_snapshot_hashes",
139                    BTreeSet::from_iter(self.individual_cardano_database_snapshots.keys().cloned()),
140                ),
141                generate_ids_array(
142                    "certificate_hashes",
143                    BTreeSet::from_iter(self.individual_certificates.keys().cloned()),
144                ),
145                generate_ids_array(
146                    "cardano_transaction_snapshot_hashes",
147                    BTreeSet::from_iter(
148                        self.individual_cardano_transaction_snapshots.keys().cloned(),
149                    ),
150                ),
151                generate_ids_array(
152                    "proof_transaction_hashes",
153                    BTreeSet::from_iter(self.cardano_transaction_proofs.keys().cloned()),
154                ),
155            ],
156            false,
157        )
158    }
159
160    pub fn generate_code_for_all_data(self) -> String {
161        Self::assemble_code(
162            &[
163                generate_list_getter("status", self.status),
164                generate_list_getter("epoch_settings", self.epoch_settings),
165                generate_ids_array(
166                    "snapshot_digests",
167                    BTreeSet::from_iter(self.individual_snapshots.keys().cloned()),
168                ),
169                generate_artifact_getter("snapshots", self.individual_snapshots),
170                generate_list_getter("snapshot_list", self.snapshots_list),
171                generate_ids_array(
172                    "mithril_stake_distribution_hashes",
173                    BTreeSet::from_iter(
174                        self.individual_mithril_stake_distributions.keys().cloned(),
175                    ),
176                ),
177                generate_artifact_getter(
178                    "mithril_stake_distributions",
179                    self.individual_mithril_stake_distributions,
180                ),
181                generate_list_getter(
182                    "mithril_stake_distribution_list",
183                    self.mithril_stake_distributions_list,
184                ),
185                generate_ids_array(
186                    "cardano_stake_distribution_hashes",
187                    BTreeSet::from_iter(
188                        self.individual_cardano_stake_distributions.keys().cloned(),
189                    ),
190                ),
191                generate_ids_array(
192                    "cardano_stake_distribution_epochs",
193                    BTreeSet::from_iter(extract_cardano_stake_distribution_epochs(
194                        &self.individual_cardano_stake_distributions,
195                    )),
196                ),
197                generate_artifact_getter(
198                    "cardano_stake_distributions",
199                    self.individual_cardano_stake_distributions,
200                ),
201                generate_list_getter(
202                    "cardano_stake_distribution_list",
203                    self.cardano_stake_distributions_list,
204                ),
205                generate_ids_array(
206                    "certificate_hashes",
207                    BTreeSet::from_iter(self.individual_certificates.keys().cloned()),
208                ),
209                generate_ids_array(
210                    "cardano_database_snapshot_hashes",
211                    BTreeSet::from_iter(self.individual_cardano_database_snapshots.keys().cloned()),
212                ),
213                generate_artifact_getter(
214                    "cardano_database_snapshots",
215                    self.individual_cardano_database_snapshots,
216                ),
217                generate_list_getter(
218                    "cardano_database_snapshot_list",
219                    self.cardano_database_snapshots_list,
220                ),
221                generate_artifact_getter("certificates", self.individual_certificates),
222                generate_list_getter("certificate_list", self.certificates_list),
223                generate_ids_array(
224                    "cardano_transaction_snapshot_hashes",
225                    BTreeSet::from_iter(
226                        self.individual_cardano_transaction_snapshots.keys().cloned(),
227                    ),
228                ),
229                generate_artifact_getter(
230                    "cardano_transaction_snapshots",
231                    self.individual_cardano_transaction_snapshots,
232                ),
233                generate_list_getter(
234                    "cardano_transaction_snapshots_list",
235                    self.cardano_transaction_snapshots_list,
236                ),
237                generate_ids_array(
238                    "proof_transaction_hashes",
239                    BTreeSet::from_iter(self.cardano_transaction_proofs.keys().cloned()),
240                ),
241                generate_artifact_getter(
242                    "cardano_transaction_proofs",
243                    self.cardano_transaction_proofs,
244                ),
245            ],
246            true,
247        )
248    }
249
250    fn assemble_code(functions_code: &[String], include_use_btree_map: bool) -> String {
251        format!(
252            "{}{}
253",
254            if include_use_btree_map {
255                "use std::collections::BTreeMap;
256
257"
258            } else {
259                ""
260            },
261            functions_code.join(
262                "
263
264"
265            )
266        )
267    }
268
269    fn read_artifacts_json_file(json_file: &Path) -> BTreeMap<ArtifactId, FileContent> {
270        let file = File::open(json_file).unwrap();
271        let parsed_json: serde_json::Value = serde_json::from_reader(&file).unwrap();
272
273        let json_object = parsed_json.as_object().unwrap();
274        let res: Result<Vec<_>, _> = json_object
275            .iter()
276            .map(|(key, value)| extract_artifact_id_and_content(key, value))
277            .collect();
278
279        BTreeMap::from_iter(res.unwrap())
280    }
281}
282
283pub fn extract_cardano_stake_distribution_epochs(
284    individual_csds: &BTreeMap<ArtifactId, FileContent>,
285) -> Vec<String> {
286    individual_csds
287        .values()
288        .map(|content| {
289            let json_value: serde_json::Value =
290                serde_json::from_str(content).unwrap_or_else(|err| {
291                    panic!("Failed to parse JSON in csd content: {content}\nError: {err}");
292                });
293
294            json_value
295                .get("epoch")
296                .and_then(|epoch| epoch.as_u64().map(|s| s.to_string()))
297                .unwrap_or_else(|| {
298                    panic!("Epoch not found or invalid in csd content: {content}");
299                })
300        })
301        .collect()
302}
303
304fn extract_artifact_id_and_content(
305    key: &String,
306    value: &serde_json::Value,
307) -> Result<(ArtifactId, FileContent), String> {
308    let json_content = serde_json::to_string_pretty(value).map_err(|e| e.to_string())?;
309    Ok((key.to_owned(), json_content))
310}
311
312pub fn list_json_files_in_folder(folder: &Path) -> impl Iterator<Item = fs::DirEntry> + '_ {
313    crate::list_files_in_folder(folder)
314        .filter(|e| e.file_name().to_string_lossy().ends_with(".json"))
315}
316
317// pub(crate) fn $fun_name()() -> BTreeMap<String, String>
318pub fn generate_artifact_getter(
319    fun_name: &str,
320    source_jsons: BTreeMap<ArtifactId, FileContent>,
321) -> String {
322    let mut artifacts_list = String::new();
323
324    for (artifact_id, file_content) in source_jsons {
325        write!(
326            artifacts_list,
327            r###"
328        (
329            "{artifact_id}",
330            r#"{file_content}"#
331        ),"###
332        )
333        .unwrap();
334    }
335
336    format!(
337        r###"pub(crate) fn {fun_name}() -> BTreeMap<String, String> {{
338    [{artifacts_list}
339    ]
340    .into_iter()
341    .map(|(k, v)| (k.to_owned(), v.to_owned()))
342    .collect()
343}}"###
344    )
345}
346
347/// pub(crate) fn $fun_name() -> &'static str
348pub fn generate_list_getter(fun_name: &str, source_json: FileContent) -> String {
349    format!(
350        r###"pub(crate) fn {fun_name}() -> &'static str {{
351    r#"{source_json}"#
352}}"###
353    )
354}
355
356/// pub(crate) fn $array_name() -> [&'a str; $ids.len]
357pub fn generate_ids_array(array_name: &str, ids: BTreeSet<ArtifactId>) -> String {
358    let mut ids_list = String::new();
359
360    for id in &ids {
361        write!(
362            ids_list,
363            r#"
364        "{id}","#
365        )
366        .unwrap();
367    }
368
369    format!(
370        r###"pub(crate) const fn {}<'a>() -> [&'a str; {}] {{
371    [{}
372    ]
373}}"###,
374        array_name,
375        ids.len(),
376        ids_list,
377    )
378}
379
380#[cfg(test)]
381mod tests {
382    use crate::get_temp_dir;
383
384    use super::*;
385
386    #[test]
387    fn generate_ids_array_with_empty_data() {
388        assert_eq!(
389            generate_ids_array("snapshots_digests", BTreeSet::new()),
390            "pub(crate) const fn snapshots_digests<'a>() -> [&'a str; 0] {
391    [
392    ]
393}"
394        );
395    }
396
397    #[test]
398    fn generate_ids_array_with_non_empty_data() {
399        assert_eq!(
400            generate_ids_array(
401                "snapshots_digests",
402                BTreeSet::from_iter(["abc".to_string(), "def".to_string(), "hij".to_string()])
403            ),
404            r#"pub(crate) const fn snapshots_digests<'a>() -> [&'a str; 3] {
405    [
406        "abc",
407        "def",
408        "hij",
409    ]
410}"#
411        );
412    }
413
414    #[test]
415    fn assemble_code_with_btree_use() {
416        assert_eq!(
417            "use std::collections::BTreeMap;
418
419fn a() {}
420
421fn b() {}
422",
423            FakeAggregatorData::assemble_code(
424                &["fn a() {}".to_string(), "fn b() {}".to_string()],
425                true
426            )
427        )
428    }
429
430    #[test]
431    fn assemble_code_without_btree_use() {
432        assert_eq!(
433            "fn a() {}
434
435fn b() {}
436",
437            FakeAggregatorData::assemble_code(
438                &["fn a() {}".to_string(), "fn b() {}".to_string()],
439                false
440            )
441        )
442    }
443
444    #[test]
445    fn parse_artifacts_json_into_btree_of_key_and_pretty_sub_json() {
446        let dir = get_temp_dir("read_artifacts_json_file");
447        let file = dir.join("test.json");
448        fs::write(
449            &file,
450            r#"{
451    "hash1": { "name": "artifact1" },
452    "hash2": { "name": "artifact2" }
453}"#,
454        )
455        .unwrap();
456
457        let id_per_json = FakeAggregatorData::read_artifacts_json_file(&file);
458
459        let expected = BTreeMap::from([
460            (
461                "hash1".to_string(),
462                r#"{
463  "name": "artifact1"
464}"#
465                .to_string(),
466            ),
467            (
468                "hash2".to_string(),
469                r#"{
470  "name": "artifact2"
471}"#
472                .to_string(),
473            ),
474        ]);
475        assert_eq!(expected, id_per_json);
476    }
477
478    #[test]
479    fn extract_csd_epochs_with_valid_data() {
480        let mut csds = BTreeMap::new();
481        csds.insert(
482            "csd-123".to_string(),
483            r#"{"hash": "csd-123", "epoch": 123}"#.to_string(),
484        );
485        csds.insert(
486            "csd-456".to_string(),
487            r#"{"hash": "csd-456", "epoch": 456}"#.to_string(),
488        );
489
490        let epochs = extract_cardano_stake_distribution_epochs(&csds);
491
492        assert_eq!(epochs, vec![123.to_string(), 456.to_string()]);
493    }
494
495    #[test]
496    #[should_panic(expected = "Failed to parse JSON in csd content")]
497    fn extract_csd_epochs_with_invalid_json() {
498        let mut csds = BTreeMap::new();
499        csds.insert(
500            "csd-123".to_string(),
501            r#""hash": "csd-123", "epoch": "123"#.to_string(),
502        );
503
504        extract_cardano_stake_distribution_epochs(&csds);
505    }
506
507    #[test]
508    #[should_panic(expected = "Epoch not found or invalid in csd content")]
509    fn test_extract_csd_epochs_with_missing_epoch() {
510        let mut csds = BTreeMap::new();
511        csds.insert("csd-123".to_string(), r#"{"hash": "csd-123"}"#.to_string());
512
513        extract_cardano_stake_distribution_epochs(&csds);
514    }
515
516    #[test]
517    fn test_extract_csd_epochs_with_empty_map() {
518        let csds = BTreeMap::new();
519
520        let epochs = extract_cardano_stake_distribution_epochs(&csds);
521
522        assert!(epochs.is_empty());
523    }
524}