Skip to content

Commit 5a1e622

Browse files
author
Damien LACHAUME / PALO-IT
committed
Add capability to only generate code for 'ids' for mithril-build-script::FakeAggregatorData
1 parent e3abce5 commit 5a1e622

File tree

2 files changed

+127
-33
lines changed

2 files changed

+127
-33
lines changed

‎mithril-build-script/src/fake_aggregator.rs‎

Lines changed: 124 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ pub type FileContent = String;
1313
/// List items are just their corresponding file content loaded in memory.
1414
/// Individual items are btreemap with the source filename as key and the file content as value.
1515
#[derive(Debug, Default)]
16-
pub struct DataFolder {
16+
pub struct FakeAggregatorData {
1717
epoch_settings: FileContent,
1818

1919
certificates_list: FileContent,
@@ -30,14 +30,14 @@ pub struct DataFolder {
3030
ctx_proofs: BTreeMap<ArtifactId, FileContent>,
3131
}
3232

33-
impl DataFolder {
33+
impl FakeAggregatorData {
3434
pub fn load_from_folder(folder: &Path) -> Self {
3535
fn extract_artifact_id(filename: &str, prefix: &str) -> String {
3636
let id_with_extension = filename.strip_prefix(prefix).unwrap();
3737
id_with_extension.strip_suffix(".json").unwrap().to_string()
3838
}
3939

40-
let mut data_folder = DataFolder::default();
40+
let mut data = FakeAggregatorData::default();
4141

4242
for entry in list_json_files_in_folder(folder) {
4343
let filename = entry.file_name().to_string_lossy().to_string();
@@ -50,94 +50,130 @@ impl DataFolder {
5050

5151
match filename.as_str() {
5252
"epoch-settings.json" => {
53-
data_folder.epoch_settings = file_content;
53+
data.epoch_settings = file_content;
5454
}
5555
"mithril-stake-distributions.json" => {
56-
data_folder.msds_list = file_content;
56+
data.msds_list = file_content;
5757
}
5858
"snapshots.json" => {
59-
data_folder.snapshots_list = file_content;
59+
data.snapshots_list = file_content;
6060
}
6161
"certificates.json" => {
62-
data_folder.certificates_list = file_content;
62+
data.certificates_list = file_content;
6363
}
6464
"ctx-commitments.json" => {
65-
data_folder.ctx_commitments_list = file_content;
65+
data.ctx_commitments_list = file_content;
6666
}
6767
_ if filename.starts_with("mithril-stake-distribution-") => {
68-
data_folder.individual_msds.insert(
68+
data.individual_msds.insert(
6969
extract_artifact_id(&filename, "mithril-stake-distribution-"),
7070
file_content,
7171
);
7272
}
7373
_ if filename.starts_with("snapshot-") => {
74-
data_folder
75-
.individual_snapshots
74+
data.individual_snapshots
7675
.insert(extract_artifact_id(&filename, "snapshot-"), file_content);
7776
}
7877
_ if filename.starts_with("certificate-") => {
79-
data_folder
80-
.individual_certificates
78+
data.individual_certificates
8179
.insert(extract_artifact_id(&filename, "certificate-"), file_content);
8280
}
8381
_ if filename.starts_with("ctx-commitment-") => {
84-
data_folder.individual_ctx_commitments.insert(
82+
data.individual_ctx_commitments.insert(
8583
extract_artifact_id(&filename, "ctx-commitment-"),
8684
file_content,
8785
);
8886
}
8987
_ if filename.starts_with("ctx-proof-") => {
90-
data_folder
91-
.ctx_proofs
88+
data.ctx_proofs
9289
.insert(extract_artifact_id(&filename, "ctx-proof-"), file_content);
9390
}
9491
// unknown file
9592
_ => {}
9693
}
9794
}
9895

99-
data_folder
96+
data
10097
}
10198

102-
pub fn generate_code_for_all_data(self) -> String {
103-
format!(
104-
"use std::collections::BTreeMap;
99+
pub fn generate_code_for_ids(self) -> String {
100+
Self::assemble_code(
101+
&[
102+
generate_ids_array(
103+
"snapshot_digests",
104+
BTreeSet::from_iter(self.individual_snapshots.keys().cloned()),
105+
),
106+
generate_ids_array(
107+
"msd_hashes",
108+
BTreeSet::from_iter(self.individual_msds.keys().cloned()),
109+
),
110+
generate_ids_array(
111+
"certificate_hashes",
112+
BTreeSet::from_iter(self.individual_certificates.keys().cloned()),
113+
),
114+
generate_ids_array(
115+
"ctx_commitment_hashes",
116+
BTreeSet::from_iter(self.individual_ctx_commitments.keys().cloned()),
117+
),
118+
generate_ids_array(
119+
"proof_transaction_hashes",
120+
BTreeSet::from_iter(self.ctx_proofs.keys().cloned()),
121+
),
122+
],
123+
false,
124+
)
125+
}
105126

106-
{}
107-
",
108-
[
127+
pub fn generate_code_for_all_data(self) -> String {
128+
Self::assemble_code(
129+
&[
109130
generate_list_getter("epoch_settings", self.epoch_settings),
110131
generate_ids_array(
111132
"snapshot_digests",
112-
BTreeSet::from_iter(self.individual_snapshots.keys().cloned())
133+
BTreeSet::from_iter(self.individual_snapshots.keys().cloned()),
113134
),
114135
generate_artifact_getter("snapshots", self.individual_snapshots),
115136
generate_list_getter("snapshot_list", self.snapshots_list),
116137
generate_ids_array(
117138
"msd_hashes",
118-
BTreeSet::from_iter(self.individual_msds.keys().cloned())
139+
BTreeSet::from_iter(self.individual_msds.keys().cloned()),
119140
),
120141
generate_artifact_getter("msds", self.individual_msds),
121142
generate_list_getter("msd_list", self.msds_list),
122143
generate_ids_array(
123144
"certificate_hashes",
124-
BTreeSet::from_iter(self.individual_certificates.keys().cloned())
145+
BTreeSet::from_iter(self.individual_certificates.keys().cloned()),
125146
),
126147
generate_artifact_getter("certificates", self.individual_certificates),
127148
generate_list_getter("certificate_list", self.certificates_list),
128149
generate_ids_array(
129150
"ctx_commitment_hashes",
130-
BTreeSet::from_iter(self.individual_ctx_commitments.keys().cloned())
151+
BTreeSet::from_iter(self.individual_ctx_commitments.keys().cloned()),
131152
),
132153
generate_artifact_getter("ctx_commitments", self.individual_ctx_commitments),
133154
generate_list_getter("ctx_commitments_list", self.ctx_commitments_list),
134155
generate_ids_array(
135156
"proof_transaction_hashes",
136-
BTreeSet::from_iter(self.ctx_proofs.keys().cloned())
157+
BTreeSet::from_iter(self.ctx_proofs.keys().cloned()),
137158
),
138159
generate_artifact_getter("ctx_proofs", self.ctx_proofs),
139-
]
140-
.join(
160+
],
161+
true,
162+
)
163+
}
164+
165+
fn assemble_code(functions_code: &[String], include_use_btree_map: bool) -> String {
166+
format!(
167+
"{}{}
168+
",
169+
if include_use_btree_map {
170+
"use std::collections::BTreeMap;
171+
172+
"
173+
} else {
174+
""
175+
},
176+
functions_code.join(
141177
"
142178
143179
"
@@ -231,4 +267,62 @@ pub fn generate_ids_array(array_name: &str, ids: BTreeSet<ArtifactId>) -> String
231267
#[cfg(test)]
232268
mod tests {
233269
use super::*;
270+
271+
#[test]
272+
fn generate_ids_array_with_empty_data() {
273+
assert_eq!(
274+
generate_ids_array("snapshots_digests", BTreeSet::new()),
275+
"pub(crate) const fn snapshots_digests<'a>() -> [&'a str; 0] {
276+
[
277+
]
278+
}"
279+
);
280+
}
281+
282+
#[test]
283+
fn generate_ids_array_with_non_empty_data() {
284+
assert_eq!(
285+
generate_ids_array(
286+
"snapshots_digests",
287+
BTreeSet::from_iter(["abc".to_string(), "def".to_string(), "hij".to_string()])
288+
),
289+
r#"pub(crate) const fn snapshots_digests<'a>() -> [&'a str; 3] {
290+
[
291+
"abc",
292+
"def",
293+
"hij",
294+
]
295+
}"#
296+
);
297+
}
298+
299+
#[test]
300+
fn assemble_code_with_btree_use() {
301+
assert_eq!(
302+
"use std::collections::BTreeMap;
303+
304+
fn a() {}
305+
306+
fn b() {}
307+
",
308+
FakeAggregatorData::assemble_code(
309+
&["fn a() {}".to_string(), "fn b() {}".to_string()],
310+
true
311+
)
312+
)
313+
}
314+
315+
#[test]
316+
fn assemble_code_without_btree_use() {
317+
assert_eq!(
318+
"fn a() {}
319+
320+
fn b() {}
321+
",
322+
FakeAggregatorData::assemble_code(
323+
&["fn a() {}".to_string(), "fn b() {}".to_string()],
324+
false
325+
)
326+
)
327+
}
234328
}

‎mithril-test-lab/mithril-aggregator-fake/build.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
// build.rs
22

3-
use mithril_build_script::fake_aggregator::DataFolder;
3+
use mithril_build_script::fake_aggregator::FakeAggregatorData;
44
use std::env;
55
use std::fs;
66
use std::path::Path;
@@ -10,8 +10,8 @@ fn main() {
1010
let dest_path = Path::new(&out_dir).join("imported_data.rs");
1111

1212
let data_folder_path: &Path = Path::new("./default_data");
13-
let data_folder = DataFolder::load_from_folder(data_folder_path);
14-
let generated_code = data_folder.generate_code_for_all_data();
13+
let data = FakeAggregatorData::load_from_folder(data_folder_path);
14+
let generated_code = data.generate_code_for_all_data();
1515
fs::write(dest_path, generated_code).unwrap();
1616

1717
println!("cargo:rerun-if-changed=default_data/");

0 commit comments

Comments
 (0)