xref: /aosp_15_r20/external/bazelbuild-rules_rust/crate_universe/src/utils/starlark/select_dict.rs (revision d4726bddaa87cc4778e7472feed243fa4b6c267f)
1 use std::collections::{BTreeMap, BTreeSet};
2 use std::fmt::Debug;
3 
4 use serde::ser::{SerializeMap, Serializer};
5 use serde::Serialize;
6 use serde_starlark::{FunctionCall, MULTILINE};
7 
8 use crate::select::{Select, SelectableOrderedValue, SelectableValue};
9 use crate::utils::starlark::{
10     looks_like_bazel_configuration_label, NoMatchingPlatformTriples, WithOriginalConfigurations,
11 };
12 
13 #[derive(Debug, PartialEq, Eq)]
14 pub(crate) struct SelectDict<U, T>
15 where
16     U: SelectableOrderedValue,
17     T: SelectableValue,
18 {
19     // Invariant: keys in this map are not in any of the inner maps of `selects`.
20     common: BTreeMap<U, T>,
21     // Invariant: none of the inner maps are empty.
22     selects: BTreeMap<String, BTreeMap<U, WithOriginalConfigurations<T>>>,
23     // Elements from the `Select` whose configuration did not get mapped to any
24     // new configuration. They could be ignored, but are preserved here to
25     // generate comments that help the user understand what happened.
26     unmapped: BTreeMap<String, BTreeMap<U, T>>,
27 }
28 
29 impl<U, T> SelectDict<U, T>
30 where
31     U: SelectableOrderedValue,
32     T: SelectableValue,
33 {
34     /// Re-keys the provided Select by the given configuration mapping.
35     /// This mapping maps from configurations in the input Select to sets
36     /// of configurations in the output SelectDict.
new( select: Select<BTreeMap<U, T>>, platforms: &BTreeMap<String, BTreeSet<String>>, ) -> Self37     pub(crate) fn new(
38         select: Select<BTreeMap<U, T>>,
39         platforms: &BTreeMap<String, BTreeSet<String>>,
40     ) -> Self {
41         let (common, selects) = select.into_parts();
42 
43         // Map new configuration -> WithOriginalConfigurations(value, old configurations).
44         let mut remapped: BTreeMap<String, BTreeMap<U, WithOriginalConfigurations<T>>> =
45             BTreeMap::new();
46         // Map unknown configuration -> value.
47         let mut unmapped: BTreeMap<String, BTreeMap<U, T>> = BTreeMap::new();
48 
49         for (original_configuration, entries) in selects {
50             match platforms.get(&original_configuration) {
51                 Some(configurations) => {
52                     for configuration in configurations {
53                         for (key, value) in &entries {
54                             remapped
55                                 .entry(configuration.clone())
56                                 .or_default()
57                                 .entry(key.clone())
58                                 .or_insert_with(|| WithOriginalConfigurations {
59                                     value: value.clone(),
60                                     original_configurations: BTreeSet::new(),
61                                 })
62                                 .original_configurations
63                                 .insert(original_configuration.clone());
64                         }
65                     }
66                 }
67                 None => {
68                     for (key, value) in entries {
69                         if looks_like_bazel_configuration_label(&original_configuration) {
70                             remapped
71                                 .entry(original_configuration.clone())
72                                 .or_default()
73                                 .entry(key)
74                                 .or_insert_with(|| WithOriginalConfigurations {
75                                     value: value.clone(),
76                                     original_configurations: BTreeSet::new(),
77                                 })
78                                 .original_configurations
79                                 .insert(original_configuration.clone());
80                         } else {
81                             unmapped
82                                 .entry(original_configuration.clone())
83                                 .or_default()
84                                 .insert(key, value);
85                         };
86                     }
87                 }
88             }
89         }
90 
91         Self {
92             common,
93             selects: remapped,
94             unmapped,
95         }
96     }
97 
is_empty(&self) -> bool98     pub(crate) fn is_empty(&self) -> bool {
99         self.common.is_empty() && self.selects.is_empty() && self.unmapped.is_empty()
100     }
101 }
102 
103 impl<U, T> Serialize for SelectDict<U, T>
104 where
105     U: SelectableOrderedValue,
106     T: SelectableValue,
107 {
serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> where S: Serializer,108     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
109     where
110         S: Serializer,
111     {
112         // If there are no platform-specific entries, we output just an ordinary
113         // dict.
114         //
115         // If there are platform-specific ones, we use the following. Ideally it
116         // could be done as `dicts.add({...}, select({...}))` but bazel_skylib's
117         // dicts.add does not support selects.
118         //
119         //     select({
120         //         "configuration": {
121         //             "common-key": "common-value",
122         //             "plat-key": "plat-value",  # cfg(whatever)
123         //         },
124         //         "//conditions:default": {
125         //             "common-key": "common-value",
126         //         },
127         //     })
128         //
129         // If there are unmapped entries, we include them like this:
130         //
131         //     selects.with_unmapped({
132         //         "configuration": {
133         //             "common-key": "common-value",
134         //             "plat-key": "plat-value",  # cfg(whatever)
135         //         },
136         //         "//conditions:default": {
137         //             "common-key": "common-value",
138         //         },
139         //         selects.NO_MATCHING_PLATFORM_TRIPLES: {
140         //             "cfg(obscure): {
141         //                 "unmapped-key": "unmapped-value",
142         //             },
143         //         },
144         //     })
145 
146         if self.selects.is_empty() && self.unmapped.is_empty() {
147             return self.common.serialize(serializer);
148         }
149 
150         struct SelectInner<'a, U, T>(&'a SelectDict<U, T>)
151         where
152             U: SelectableOrderedValue,
153             T: SelectableValue;
154 
155         impl<'a, U, T> Serialize for SelectInner<'a, U, T>
156         where
157             U: SelectableOrderedValue,
158             T: SelectableValue,
159         {
160             fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
161             where
162                 S: Serializer,
163             {
164                 let mut map = serializer.serialize_map(Some(MULTILINE))?;
165                 for (configuration, dict) in &self.0.selects {
166                     #[derive(Serialize)]
167                     #[serde(untagged)]
168                     enum Either<'a, T> {
169                         Common(&'a T),
170                         Selects(&'a WithOriginalConfigurations<T>),
171                     }
172 
173                     let mut combined = BTreeMap::new();
174                     combined.extend(
175                         self.0
176                             .common
177                             .iter()
178                             .map(|(key, value)| (key, Either::Common(value))),
179                     );
180                     combined.extend(
181                         dict.iter()
182                             .map(|(key, value)| (key, Either::Selects(value))),
183                     );
184                     map.serialize_entry(configuration, &combined)?;
185                 }
186                 map.serialize_entry("//conditions:default", &self.0.common)?;
187                 if !self.0.unmapped.is_empty() {
188                     struct SelectUnmapped<'a, U, T>(&'a BTreeMap<String, BTreeMap<U, T>>)
189                     where
190                         U: SelectableOrderedValue,
191                         T: SelectableValue;
192 
193                     impl<'a, U, T> Serialize for SelectUnmapped<'a, U, T>
194                     where
195                         U: SelectableOrderedValue,
196                         T: SelectableValue,
197                     {
198                         fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
199                         where
200                             S: Serializer,
201                         {
202                             let mut map = serializer.serialize_map(Some(MULTILINE))?;
203                             for (cfg, dict) in self.0.iter() {
204                                 map.serialize_entry(cfg, dict)?;
205                             }
206                             map.end()
207                         }
208                     }
209 
210                     map.serialize_entry(
211                         &NoMatchingPlatformTriples,
212                         &SelectUnmapped(&self.0.unmapped),
213                     )?;
214                 }
215                 map.end()
216             }
217         }
218 
219         let function = if self.unmapped.is_empty() {
220             "select"
221         } else {
222             "selects.with_unmapped"
223         };
224 
225         FunctionCall::new(function, [SelectInner(self)]).serialize(serializer)
226     }
227 }
228 
229 #[cfg(test)]
230 mod test {
231     use super::*;
232 
233     use indoc::indoc;
234 
235     #[test]
empty_select_dict()236     fn empty_select_dict() {
237         let select_dict: SelectDict<String, String> =
238             SelectDict::new(Default::default(), &Default::default());
239 
240         let expected_starlark = indoc! {r#"
241             {}
242         "#};
243 
244         assert_eq!(
245             select_dict.serialize(serde_starlark::Serializer).unwrap(),
246             expected_starlark,
247         );
248     }
249 
250     #[test]
no_platform_specific_select_dict()251     fn no_platform_specific_select_dict() {
252         let mut select: Select<BTreeMap<String, String>> = Select::default();
253         select.insert(("Greeting".to_owned(), "Hello".to_owned()), None);
254 
255         let select_dict = SelectDict::new(select, &Default::default());
256 
257         let expected_starlark = indoc! {r#"
258             {
259                 "Greeting": "Hello",
260             }
261         "#};
262 
263         assert_eq!(
264             select_dict.serialize(serde_starlark::Serializer).unwrap(),
265             expected_starlark,
266         );
267     }
268 
269     #[test]
only_platform_specific_select_dict()270     fn only_platform_specific_select_dict() {
271         let mut select: Select<BTreeMap<String, String>> = Select::default();
272         select.insert(
273             ("Greeting".to_owned(), "Hello".to_owned()),
274             Some("platform".to_owned()),
275         );
276 
277         let platforms = BTreeMap::from([(
278             "platform".to_owned(),
279             BTreeSet::from(["platform".to_owned()]),
280         )]);
281 
282         let select_dict = SelectDict::new(select, &platforms);
283 
284         let expected_starlark = indoc! {r#"
285             select({
286                 "platform": {
287                     "Greeting": "Hello",  # platform
288                 },
289                 "//conditions:default": {},
290             })
291         "#};
292 
293         assert_eq!(
294             select_dict.serialize(serde_starlark::Serializer).unwrap(),
295             expected_starlark,
296         );
297     }
298 
299     #[test]
mixed_select_dict()300     fn mixed_select_dict() {
301         let mut select: Select<BTreeMap<String, String>> = Select::default();
302         select.insert(
303             ("Greeting".to_owned(), "Hello".to_owned()),
304             Some("platform".to_owned()),
305         );
306         select.insert(("Message".to_owned(), "Goodbye".to_owned()), None);
307 
308         let platforms = BTreeMap::from([(
309             "platform".to_owned(),
310             BTreeSet::from(["platform".to_owned()]),
311         )]);
312 
313         let select_dict = SelectDict::new(select, &platforms);
314 
315         let expected_starlark = indoc! {r#"
316             select({
317                 "platform": {
318                     "Greeting": "Hello",  # platform
319                     "Message": "Goodbye",
320                 },
321                 "//conditions:default": {
322                     "Message": "Goodbye",
323                 },
324             })
325         "#};
326 
327         assert_eq!(
328             select_dict.serialize(serde_starlark::Serializer).unwrap(),
329             expected_starlark,
330         );
331     }
332 
333     #[test]
remap_select_dict_configurations()334     fn remap_select_dict_configurations() {
335         let mut select: Select<BTreeMap<String, String>> = Select::default();
336         select.insert(
337             ("dep-a".to_owned(), "a".to_owned()),
338             Some("cfg(macos)".to_owned()),
339         );
340         select.insert(
341             ("dep-b".to_owned(), "b".to_owned()),
342             Some("cfg(macos)".to_owned()),
343         );
344         select.insert(
345             ("dep-d".to_owned(), "d".to_owned()),
346             Some("cfg(macos)".to_owned()),
347         );
348         select.insert(
349             ("dep-a".to_owned(), "a".to_owned()),
350             Some("cfg(x86_64)".to_owned()),
351         );
352         select.insert(
353             ("dep-c".to_owned(), "c".to_owned()),
354             Some("cfg(x86_64)".to_owned()),
355         );
356         select.insert(
357             ("dep-e".to_owned(), "e".to_owned()),
358             Some("cfg(pdp11)".to_owned()),
359         );
360         select.insert(("dep-d".to_owned(), "d".to_owned()), None);
361         select.insert(
362             ("dep-f".to_owned(), "f".to_owned()),
363             Some("@platforms//os:magic".to_owned()),
364         );
365         select.insert(
366             ("dep-g".to_owned(), "g".to_owned()),
367             Some("//another:platform".to_owned()),
368         );
369 
370         let platforms = BTreeMap::from([
371             (
372                 "cfg(macos)".to_owned(),
373                 BTreeSet::from(["x86_64-macos".to_owned(), "aarch64-macos".to_owned()]),
374             ),
375             (
376                 "cfg(x86_64)".to_owned(),
377                 BTreeSet::from(["x86_64-linux".to_owned(), "x86_64-macos".to_owned()]),
378             ),
379         ]);
380 
381         let select_dict = SelectDict::new(select, &platforms);
382 
383         let expected = SelectDict {
384             common: BTreeMap::from([("dep-d".to_string(), "d".to_owned())]),
385             selects: BTreeMap::from([
386                 (
387                     "x86_64-macos".to_owned(),
388                     BTreeMap::from([
389                         (
390                             "dep-a".to_string(),
391                             WithOriginalConfigurations {
392                                 value: "a".to_owned(),
393                                 original_configurations: BTreeSet::from([
394                                     "cfg(macos)".to_owned(),
395                                     "cfg(x86_64)".to_owned(),
396                                 ]),
397                             },
398                         ),
399                         (
400                             "dep-b".to_string(),
401                             WithOriginalConfigurations {
402                                 value: "b".to_owned(),
403                                 original_configurations: BTreeSet::from(["cfg(macos)".to_owned()]),
404                             },
405                         ),
406                         (
407                             "dep-c".to_string(),
408                             WithOriginalConfigurations {
409                                 value: "c".to_owned(),
410                                 original_configurations: BTreeSet::from(["cfg(x86_64)".to_owned()]),
411                             },
412                         ),
413                     ]),
414                 ),
415                 (
416                     "aarch64-macos".to_owned(),
417                     BTreeMap::from([
418                         (
419                             "dep-a".to_string(),
420                             WithOriginalConfigurations {
421                                 value: "a".to_owned(),
422                                 original_configurations: BTreeSet::from(["cfg(macos)".to_owned()]),
423                             },
424                         ),
425                         (
426                             "dep-b".to_string(),
427                             WithOriginalConfigurations {
428                                 value: "b".to_owned(),
429                                 original_configurations: BTreeSet::from(["cfg(macos)".to_owned()]),
430                             },
431                         ),
432                     ]),
433                 ),
434                 (
435                     "x86_64-linux".to_owned(),
436                     BTreeMap::from([
437                         (
438                             "dep-a".to_string(),
439                             WithOriginalConfigurations {
440                                 value: "a".to_owned(),
441                                 original_configurations: BTreeSet::from(["cfg(x86_64)".to_owned()]),
442                             },
443                         ),
444                         (
445                             "dep-c".to_string(),
446                             WithOriginalConfigurations {
447                                 value: "c".to_owned(),
448                                 original_configurations: BTreeSet::from(["cfg(x86_64)".to_owned()]),
449                             },
450                         ),
451                     ]),
452                 ),
453                 (
454                     "@platforms//os:magic".to_owned(),
455                     BTreeMap::from([(
456                         "dep-f".to_string(),
457                         WithOriginalConfigurations {
458                             value: "f".to_owned(),
459                             original_configurations: BTreeSet::from([
460                                 "@platforms//os:magic".to_owned()
461                             ]),
462                         },
463                     )]),
464                 ),
465                 (
466                     "//another:platform".to_owned(),
467                     BTreeMap::from([(
468                         "dep-g".to_string(),
469                         WithOriginalConfigurations {
470                             value: "g".to_owned(),
471                             original_configurations: BTreeSet::from([
472                                 "//another:platform".to_owned()
473                             ]),
474                         },
475                     )]),
476                 ),
477             ]),
478             unmapped: BTreeMap::from([(
479                 "cfg(pdp11)".to_owned(),
480                 BTreeMap::from([("dep-e".to_string(), "e".to_owned())]),
481             )]),
482         };
483 
484         assert_eq!(select_dict, expected);
485 
486         let expected_starlark = indoc! {r#"
487             selects.with_unmapped({
488                 "//another:platform": {
489                     "dep-d": "d",
490                     "dep-g": "g",  # //another:platform
491                 },
492                 "@platforms//os:magic": {
493                     "dep-d": "d",
494                     "dep-f": "f",  # @platforms//os:magic
495                 },
496                 "aarch64-macos": {
497                     "dep-a": "a",  # cfg(macos)
498                     "dep-b": "b",  # cfg(macos)
499                     "dep-d": "d",
500                 },
501                 "x86_64-linux": {
502                     "dep-a": "a",  # cfg(x86_64)
503                     "dep-c": "c",  # cfg(x86_64)
504                     "dep-d": "d",
505                 },
506                 "x86_64-macos": {
507                     "dep-a": "a",  # cfg(macos), cfg(x86_64)
508                     "dep-b": "b",  # cfg(macos)
509                     "dep-c": "c",  # cfg(x86_64)
510                     "dep-d": "d",
511                 },
512                 "//conditions:default": {
513                     "dep-d": "d",
514                 },
515                 selects.NO_MATCHING_PLATFORM_TRIPLES: {
516                     "cfg(pdp11)": {
517                         "dep-e": "e",
518                     },
519                 },
520             })
521         "#};
522 
523         assert_eq!(
524             select_dict.serialize(serde_starlark::Serializer).unwrap(),
525             expected_starlark,
526         );
527     }
528 }
529