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