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.keys().cloned(),
144                    ),
145                ),
146                generate_ids_array(
147                    "proof_transaction_hashes",
148                    BTreeSet::from_iter(self.cardano_transaction_proofs.keys().cloned()),
149                ),
150            ],
151            false,
152        )
153    }
154
155    pub fn generate_code_for_all_data(self) -> String {
156        Self::assemble_code(
157            &[
158                generate_list_getter("epoch_settings", self.epoch_settings),
159                generate_ids_array(
160                    "snapshot_digests",
161                    BTreeSet::from_iter(self.individual_snapshots.keys().cloned()),
162                ),
163                generate_artifact_getter("snapshots", self.individual_snapshots),
164                generate_list_getter("snapshot_list", self.snapshots_list),
165                generate_ids_array(
166                    "mithril_stake_distribution_hashes",
167                    BTreeSet::from_iter(
168                        self.individual_mithril_stake_distributions.keys().cloned(),
169                    ),
170                ),
171                generate_artifact_getter(
172                    "mithril_stake_distributions",
173                    self.individual_mithril_stake_distributions,
174                ),
175                generate_list_getter(
176                    "mithril_stake_distribution_list",
177                    self.mithril_stake_distributions_list,
178                ),
179                generate_ids_array(
180                    "cardano_stake_distribution_hashes",
181                    BTreeSet::from_iter(
182                        self.individual_cardano_stake_distributions.keys().cloned(),
183                    ),
184                ),
185                generate_ids_array(
186                    "cardano_stake_distribution_epochs",
187                    BTreeSet::from_iter(extract_cardano_stake_distribution_epochs(
188                        &self.individual_cardano_stake_distributions,
189                    )),
190                ),
191                generate_artifact_getter(
192                    "cardano_stake_distributions",
193                    self.individual_cardano_stake_distributions,
194                ),
195                generate_list_getter(
196                    "cardano_stake_distribution_list",
197                    self.cardano_stake_distributions_list,
198                ),
199                generate_ids_array(
200                    "certificate_hashes",
201                    BTreeSet::from_iter(self.individual_certificates.keys().cloned()),
202                ),
203                generate_ids_array(
204                    "cardano_database_snapshot_hashes",
205                    BTreeSet::from_iter(self.individual_cardano_database_snapshots.keys().cloned()),
206                ),
207                generate_artifact_getter(
208                    "cardano_database_snapshots",
209                    self.individual_cardano_database_snapshots,
210                ),
211                generate_list_getter(
212                    "cardano_database_snapshot_list",
213                    self.cardano_database_snapshots_list,
214                ),
215                generate_artifact_getter("certificates", self.individual_certificates),
216                generate_list_getter("certificate_list", self.certificates_list),
217                generate_ids_array(
218                    "cardano_transaction_snapshot_hashes",
219                    BTreeSet::from_iter(
220                        self.individual_cardano_transaction_snapshots.keys().cloned(),
221                    ),
222                ),
223                generate_artifact_getter(
224                    "cardano_transaction_snapshots",
225                    self.individual_cardano_transaction_snapshots,
226                ),
227                generate_list_getter(
228                    "cardano_transaction_snapshots_list",
229                    self.cardano_transaction_snapshots_list,
230                ),
231                generate_ids_array(
232                    "proof_transaction_hashes",
233                    BTreeSet::from_iter(self.cardano_transaction_proofs.keys().cloned()),
234                ),
235                generate_artifact_getter(
236                    "cardano_transaction_proofs",
237                    self.cardano_transaction_proofs,
238                ),
239            ],
240            true,
241        )
242    }
243
244    fn assemble_code(functions_code: &[String], include_use_btree_map: bool) -> String {
245        format!(
246            "{}{}
247",
248            if include_use_btree_map {
249                "use std::collections::BTreeMap;
250
251"
252            } else {
253                ""
254            },
255            functions_code.join(
256                "
257
258"
259            )
260        )
261    }
262
263    fn read_artifacts_json_file(json_file: &Path) -> BTreeMap<ArtifactId, FileContent> {
264        let file = File::open(json_file).unwrap();
265        let parsed_json: serde_json::Value = serde_json::from_reader(&file).unwrap();
266
267        let json_object = parsed_json.as_object().unwrap();
268        let res: Result<Vec<_>, _> = json_object
269            .iter()
270            .map(|(key, value)| extract_artifact_id_and_content(key, value))
271            .collect();
272
273        BTreeMap::from_iter(res.unwrap())
274    }
275}
276
277pub fn extract_cardano_stake_distribution_epochs(
278    individual_csds: &BTreeMap<ArtifactId, FileContent>,
279) -> Vec<String> {
280    individual_csds
281        .values()
282        .map(|content| {
283            let json_value: serde_json::Value =
284                serde_json::from_str(content).unwrap_or_else(|err| {
285                    panic!("Failed to parse JSON in csd content: {content}\nError: {err}");
286                });
287
288            json_value
289                .get("epoch")
290                .and_then(|epoch| epoch.as_u64().map(|s| s.to_string()))
291                .unwrap_or_else(|| {
292                    panic!("Epoch not found or invalid in csd content: {content}");
293                })
294        })
295        .collect()
296}
297
298fn extract_artifact_id_and_content(
299    key: &String,
300    value: &serde_json::Value,
301) -> Result<(ArtifactId, FileContent), String> {
302    let json_content = serde_json::to_string_pretty(value).map_err(|e| e.to_string())?;
303    Ok((key.to_owned(), json_content))
304}
305
306pub fn list_json_files_in_folder(folder: &Path) -> impl Iterator<Item = fs::DirEntry> + '_ {
307    crate::list_files_in_folder(folder)
308        .filter(|e| e.file_name().to_string_lossy().ends_with(".json"))
309}
310
311// pub(crate) fn $fun_name()() -> BTreeMap<String, String>
312pub fn generate_artifact_getter(
313    fun_name: &str,
314    source_jsons: BTreeMap<ArtifactId, FileContent>,
315) -> String {
316    let mut artifacts_list = String::new();
317
318    for (artifact_id, file_content) in source_jsons {
319        write!(
320            artifacts_list,
321            r###"
322        (
323            "{artifact_id}",
324            r#"{file_content}"#
325        ),"###
326        )
327        .unwrap();
328    }
329
330    format!(
331        r###"pub(crate) fn {fun_name}() -> BTreeMap<String, String> {{
332    [{artifacts_list}
333    ]
334    .into_iter()
335    .map(|(k, v)| (k.to_owned(), v.to_owned()))
336    .collect()
337}}"###
338    )
339}
340
341/// pub(crate) fn $fun_name() -> &'static str
342pub fn generate_list_getter(fun_name: &str, source_json: FileContent) -> String {
343    format!(
344        r###"pub(crate) fn {fun_name}() -> &'static str {{
345    r#"{source_json}"#
346}}"###
347    )
348}
349
350/// pub(crate) fn $array_name() -> [&'a str; $ids.len]
351pub fn generate_ids_array(array_name: &str, ids: BTreeSet<ArtifactId>) -> String {
352    let mut ids_list = String::new();
353
354    for id in &ids {
355        write!(
356            ids_list,
357            r#"
358        "{id}","#
359        )
360        .unwrap();
361    }
362
363    format!(
364        r###"pub(crate) const fn {}<'a>() -> [&'a str; {}] {{
365    [{}
366    ]
367}}"###,
368        array_name,
369        ids.len(),
370        ids_list,
371    )
372}
373
374#[cfg(test)]
375mod tests {
376    use crate::get_temp_dir;
377
378    use super::*;
379
380    #[test]
381    fn generate_ids_array_with_empty_data() {
382        assert_eq!(
383            generate_ids_array("snapshots_digests", BTreeSet::new()),
384            "pub(crate) const fn snapshots_digests<'a>() -> [&'a str; 0] {
385    [
386    ]
387}"
388        );
389    }
390
391    #[test]
392    fn generate_ids_array_with_non_empty_data() {
393        assert_eq!(
394            generate_ids_array(
395                "snapshots_digests",
396                BTreeSet::from_iter(["abc".to_string(), "def".to_string(), "hij".to_string()])
397            ),
398            r#"pub(crate) const fn snapshots_digests<'a>() -> [&'a str; 3] {
399    [
400        "abc",
401        "def",
402        "hij",
403    ]
404}"#
405        );
406    }
407
408    #[test]
409    fn assemble_code_with_btree_use() {
410        assert_eq!(
411            "use std::collections::BTreeMap;
412
413fn a() {}
414
415fn b() {}
416",
417            FakeAggregatorData::assemble_code(
418                &["fn a() {}".to_string(), "fn b() {}".to_string()],
419                true
420            )
421        )
422    }
423
424    #[test]
425    fn assemble_code_without_btree_use() {
426        assert_eq!(
427            "fn a() {}
428
429fn b() {}
430",
431            FakeAggregatorData::assemble_code(
432                &["fn a() {}".to_string(), "fn b() {}".to_string()],
433                false
434            )
435        )
436    }
437
438    #[test]
439    fn parse_artifacts_json_into_btree_of_key_and_pretty_sub_json() {
440        let dir = get_temp_dir("read_artifacts_json_file");
441        let file = dir.join("test.json");
442        fs::write(
443            &file,
444            r#"{
445    "hash1": { "name": "artifact1" },
446    "hash2": { "name": "artifact2" }
447}"#,
448        )
449        .unwrap();
450
451        let id_per_json = FakeAggregatorData::read_artifacts_json_file(&file);
452
453        let expected = BTreeMap::from([
454            (
455                "hash1".to_string(),
456                r#"{
457  "name": "artifact1"
458}"#
459                .to_string(),
460            ),
461            (
462                "hash2".to_string(),
463                r#"{
464  "name": "artifact2"
465}"#
466                .to_string(),
467            ),
468        ]);
469        assert_eq!(expected, id_per_json);
470    }
471
472    #[test]
473    fn extract_csd_epochs_with_valid_data() {
474        let mut csds = BTreeMap::new();
475        csds.insert(
476            "csd-123".to_string(),
477            r#"{"hash": "csd-123", "epoch": 123}"#.to_string(),
478        );
479        csds.insert(
480            "csd-456".to_string(),
481            r#"{"hash": "csd-456", "epoch": 456}"#.to_string(),
482        );
483
484        let epochs = extract_cardano_stake_distribution_epochs(&csds);
485
486        assert_eq!(epochs, vec![123.to_string(), 456.to_string()]);
487    }
488
489    #[test]
490    #[should_panic(expected = "Failed to parse JSON in csd content")]
491    fn extract_csd_epochs_with_invalid_json() {
492        let mut csds = BTreeMap::new();
493        csds.insert(
494            "csd-123".to_string(),
495            r#""hash": "csd-123", "epoch": "123"#.to_string(),
496        );
497
498        extract_cardano_stake_distribution_epochs(&csds);
499    }
500
501    #[test]
502    #[should_panic(expected = "Epoch not found or invalid in csd content")]
503    fn test_extract_csd_epochs_with_missing_epoch() {
504        let mut csds = BTreeMap::new();
505        csds.insert("csd-123".to_string(), r#"{"hash": "csd-123"}"#.to_string());
506
507        extract_cardano_stake_distribution_epochs(&csds);
508    }
509
510    #[test]
511    fn test_extract_csd_epochs_with_empty_map() {
512        let csds = BTreeMap::new();
513
514        let epochs = extract_cardano_stake_distribution_epochs(&csds);
515
516        assert!(epochs.is_empty());
517    }
518}