xref: /aosp_15_r20/external/bazelbuild-rules_rust/crate_universe/src/utils/starlark/select_set.rs (revision d4726bddaa87cc4778e7472feed243fa4b6c267f)
1 use std::collections::{BTreeMap, BTreeSet};
2 use std::fmt::Debug;
3 
4 use serde::ser::{SerializeMap, SerializeTupleStruct, Serializer};
5 use serde::Serialize;
6 use serde_starlark::{FunctionCall, MULTILINE};
7 
8 use crate::select::{Select, SelectableOrderedValue};
9 use crate::utils::starlark::serialize::MultilineArray;
10 use crate::utils::starlark::{
11     looks_like_bazel_configuration_label, NoMatchingPlatformTriples, WithOriginalConfigurations,
12 };
13 
14 #[derive(Debug, PartialEq, Eq)]
15 pub(crate) struct SelectSet<T>
16 where
17     T: SelectableOrderedValue,
18 {
19     // Invariant: any T in `common` is not anywhere in `selects`.
20     common: BTreeSet<T>,
21     // Invariant: none of the sets are empty.
22     selects: BTreeMap<String, BTreeSet<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, BTreeSet<T>>,
27 }
28 
29 impl<T> SelectSet<T>
30 where
31     T: SelectableOrderedValue,
32 {
33     /// Re-keys the provided Select by the given configuration mapping.
34     /// This mapping maps from configurations in the input Select to sets of
35     /// configurations in the output SelectSet.
new( select: Select<BTreeSet<T>>, platforms: &BTreeMap<String, BTreeSet<String>>, ) -> Self36     pub(crate) fn new(
37         select: Select<BTreeSet<T>>,
38         platforms: &BTreeMap<String, BTreeSet<String>>,
39     ) -> Self {
40         let (common, selects) = select.into_parts();
41 
42         // Map new configuration -> value -> old configurations.
43         let mut remapped: BTreeMap<String, BTreeMap<T, BTreeSet<String>>> = BTreeMap::new();
44         // Map unknown configuration -> value.
45         let mut unmapped: BTreeMap<String, BTreeSet<T>> = BTreeMap::new();
46 
47         for (original_configuration, values) in selects {
48             match platforms.get(&original_configuration) {
49                 Some(configurations) => {
50                     for configuration in configurations {
51                         for value in &values {
52                             remapped
53                                 .entry(configuration.clone())
54                                 .or_default()
55                                 .entry(value.clone())
56                                 .or_default()
57                                 .insert(original_configuration.clone());
58                         }
59                     }
60                 }
61                 None => {
62                     if looks_like_bazel_configuration_label(&original_configuration) {
63                         let destination =
64                             remapped.entry(original_configuration.clone()).or_default();
65                         for value in values {
66                             destination
67                                 .entry(value)
68                                 .or_default()
69                                 .insert(original_configuration.clone());
70                         }
71                     } else {
72                         unmapped
73                             .entry(original_configuration.clone())
74                             .or_default()
75                             .extend(values.into_iter());
76                     };
77                 }
78             }
79         }
80 
81         Self {
82             common,
83             selects: remapped
84                 .into_iter()
85                 .map(|(new_configuration, value_to_original_configuration)| {
86                     (
87                         new_configuration,
88                         value_to_original_configuration
89                             .into_iter()
90                             .map(
91                                 |(value, original_configurations)| WithOriginalConfigurations {
92                                     value,
93                                     original_configurations,
94                                 },
95                             )
96                             .collect(),
97                     )
98                 })
99                 .collect(),
100             unmapped,
101         }
102     }
103 
104     /// Determine whether or not the select should be serialized
is_empty(&self) -> bool105     pub(crate) fn is_empty(&self) -> bool {
106         self.common.is_empty() && self.selects.is_empty() && self.unmapped.is_empty()
107     }
108 }
109 
110 impl<T> Serialize for SelectSet<T>
111 where
112     T: SelectableOrderedValue,
113 {
serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> where S: Serializer,114     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
115     where
116         S: Serializer,
117     {
118         // Output looks like:
119         //
120         //     [
121         //         "common...",
122         //     ] + select({
123         //         "configuration": [
124         //             "value...",  # cfg(whatever)
125         //         ],
126         //         "//conditions:default": [],
127         //     })
128         //
129         // The common part and select are each omitted if they are empty (except
130         // if the entire thing is empty, in which case we serialize the common
131         // part to get an empty array).
132         //
133         // If there are unmapped entries, we include them like this:
134         //
135         //     [
136         //         "common...",
137         //     ] + selects.with_unmapped({
138         //         "configuration": [
139         //             "value...",  # cfg(whatever)
140         //         ],
141         //         "//conditions:default": [],
142         //         selects.NO_MATCHING_PLATFORM_TRIPLES: {
143         //             "cfg(obscure)": [
144         //                 "value...",
145         //             ],
146         //         },
147         //     })
148 
149         let mut plus = serializer.serialize_tuple_struct("+", MULTILINE)?;
150 
151         if !self.common.is_empty() || self.selects.is_empty() && self.unmapped.is_empty() {
152             plus.serialize_field(&MultilineArray(&self.common))?;
153         }
154 
155         if !self.selects.is_empty() || !self.unmapped.is_empty() {
156             struct SelectInner<'a, T: Ord>(&'a SelectSet<T>)
157             where
158                 T: SelectableOrderedValue;
159 
160             impl<'a, T> Serialize for SelectInner<'a, T>
161             where
162                 T: SelectableOrderedValue,
163             {
164                 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
165                 where
166                     S: Serializer,
167                 {
168                     let mut map = serializer.serialize_map(Some(MULTILINE))?;
169                     for (cfg, value) in &self.0.selects {
170                         map.serialize_entry(cfg, &MultilineArray(value))?;
171                     }
172                     map.serialize_entry("//conditions:default", &[] as &[T])?;
173                     if !self.0.unmapped.is_empty() {
174                         struct SelectUnmapped<'a, T>(&'a BTreeMap<String, BTreeSet<T>>)
175                         where
176                             T: SelectableOrderedValue;
177 
178                         impl<'a, T> Serialize for SelectUnmapped<'a, T>
179                         where
180                             T: SelectableOrderedValue,
181                         {
182                             fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
183                             where
184                                 S: Serializer,
185                             {
186                                 let mut map = serializer.serialize_map(Some(MULTILINE))?;
187                                 for (cfg, values) in self.0.iter() {
188                                     map.serialize_entry(cfg, &MultilineArray(values))?;
189                                 }
190                                 map.end()
191                             }
192                         }
193 
194                         map.serialize_entry(
195                             &NoMatchingPlatformTriples,
196                             &SelectUnmapped(&self.0.unmapped),
197                         )?;
198                     }
199                     map.end()
200                 }
201             }
202 
203             let function = if self.unmapped.is_empty() {
204                 "select"
205             } else {
206                 "selects.with_unmapped"
207             };
208 
209             plus.serialize_field(&FunctionCall::new(function, [SelectInner(self)]))?;
210         }
211 
212         plus.end()
213     }
214 }
215 
216 #[cfg(test)]
217 mod test {
218     use super::*;
219 
220     use indoc::indoc;
221 
222     #[test]
empty_select_set()223     fn empty_select_set() {
224         let select_set: SelectSet<String> = SelectSet::new(Default::default(), &Default::default());
225 
226         let expected_starlark = indoc! {r#"
227             []
228         "#};
229 
230         assert_eq!(
231             select_set.serialize(serde_starlark::Serializer).unwrap(),
232             expected_starlark,
233         );
234     }
235 
236     #[test]
no_platform_specific_select_set()237     fn no_platform_specific_select_set() {
238         let mut select: Select<BTreeSet<String>> = Select::default();
239         select.insert("Hello".to_owned(), None);
240 
241         let select_set = SelectSet::new(select, &Default::default());
242 
243         let expected_starlark = indoc! {r#"
244             [
245                 "Hello",
246             ]
247         "#};
248 
249         assert_eq!(
250             select_set.serialize(serde_starlark::Serializer).unwrap(),
251             expected_starlark,
252         );
253     }
254 
255     #[test]
only_platform_specific_select_set()256     fn only_platform_specific_select_set() {
257         let mut select: Select<BTreeSet<String>> = Select::default();
258         select.insert("Hello".to_owned(), Some("platform".to_owned()));
259 
260         let platforms = BTreeMap::from([(
261             "platform".to_owned(),
262             BTreeSet::from(["platform".to_owned()]),
263         )]);
264 
265         let select_set = SelectSet::new(select, &platforms);
266 
267         let expected_starlark = indoc! {r#"
268             select({
269                 "platform": [
270                     "Hello",  # platform
271                 ],
272                 "//conditions:default": [],
273             })
274         "#};
275 
276         assert_eq!(
277             select_set.serialize(serde_starlark::Serializer).unwrap(),
278             expected_starlark,
279         );
280     }
281 
282     #[test]
mixed_select_set()283     fn mixed_select_set() {
284         let mut select: Select<BTreeSet<String>> = Select::default();
285         select.insert("Hello".to_owned(), Some("platform".to_owned()));
286         select.insert("Goodbye".to_owned(), None);
287 
288         let platforms = BTreeMap::from([(
289             "platform".to_owned(),
290             BTreeSet::from(["platform".to_owned()]),
291         )]);
292 
293         let select_set = SelectSet::new(select, &platforms);
294 
295         let expected_starlark = indoc! {r#"
296             [
297                 "Goodbye",
298             ] + select({
299                 "platform": [
300                     "Hello",  # platform
301                 ],
302                 "//conditions:default": [],
303             })
304         "#};
305 
306         assert_eq!(
307             select_set.serialize(serde_starlark::Serializer).unwrap(),
308             expected_starlark,
309         );
310     }
311 
312     #[test]
remap_select_set_configurations()313     fn remap_select_set_configurations() {
314         let mut select: Select<BTreeSet<String>> = Select::default();
315         select.insert("dep-a".to_owned(), Some("cfg(macos)".to_owned()));
316         select.insert("dep-b".to_owned(), Some("cfg(macos)".to_owned()));
317         select.insert("dep-d".to_owned(), Some("cfg(macos)".to_owned()));
318         select.insert("dep-a".to_owned(), Some("cfg(x86_64)".to_owned()));
319         select.insert("dep-c".to_owned(), Some("cfg(x86_64)".to_owned()));
320         select.insert("dep-e".to_owned(), Some("cfg(pdp11)".to_owned()));
321         select.insert("dep-d".to_owned(), None);
322         select.insert("dep-f".to_owned(), Some("@platforms//os:magic".to_owned()));
323         select.insert("dep-g".to_owned(), Some("//another:platform".to_owned()));
324 
325         let platforms = BTreeMap::from([
326             (
327                 "cfg(macos)".to_owned(),
328                 BTreeSet::from(["x86_64-macos".to_owned(), "aarch64-macos".to_owned()]),
329             ),
330             (
331                 "cfg(x86_64)".to_owned(),
332                 BTreeSet::from(["x86_64-linux".to_owned(), "x86_64-macos".to_owned()]),
333             ),
334         ]);
335 
336         let select_set = SelectSet::new(select, &platforms);
337 
338         let expected = SelectSet {
339             common: BTreeSet::from(["dep-d".to_owned()]),
340             selects: BTreeMap::from([
341                 (
342                     "x86_64-macos".to_owned(),
343                     BTreeSet::from([
344                         WithOriginalConfigurations {
345                             value: "dep-a".to_owned(),
346                             original_configurations: BTreeSet::from([
347                                 "cfg(macos)".to_owned(),
348                                 "cfg(x86_64)".to_owned(),
349                             ]),
350                         },
351                         WithOriginalConfigurations {
352                             value: "dep-b".to_owned(),
353                             original_configurations: BTreeSet::from(["cfg(macos)".to_owned()]),
354                         },
355                         WithOriginalConfigurations {
356                             value: "dep-c".to_owned(),
357                             original_configurations: BTreeSet::from(["cfg(x86_64)".to_owned()]),
358                         },
359                     ]),
360                 ),
361                 (
362                     "aarch64-macos".to_owned(),
363                     BTreeSet::from([
364                         WithOriginalConfigurations {
365                             value: "dep-a".to_owned(),
366                             original_configurations: BTreeSet::from(["cfg(macos)".to_owned()]),
367                         },
368                         WithOriginalConfigurations {
369                             value: "dep-b".to_owned(),
370                             original_configurations: BTreeSet::from(["cfg(macos)".to_owned()]),
371                         },
372                     ]),
373                 ),
374                 (
375                     "x86_64-linux".to_owned(),
376                     BTreeSet::from([
377                         WithOriginalConfigurations {
378                             value: "dep-a".to_owned(),
379                             original_configurations: BTreeSet::from(["cfg(x86_64)".to_owned()]),
380                         },
381                         WithOriginalConfigurations {
382                             value: "dep-c".to_owned(),
383                             original_configurations: BTreeSet::from(["cfg(x86_64)".to_owned()]),
384                         },
385                     ]),
386                 ),
387                 (
388                     "@platforms//os:magic".to_owned(),
389                     BTreeSet::from([WithOriginalConfigurations {
390                         value: "dep-f".to_owned(),
391                         original_configurations: BTreeSet::from(
392                             ["@platforms//os:magic".to_owned()],
393                         ),
394                     }]),
395                 ),
396                 (
397                     "//another:platform".to_owned(),
398                     BTreeSet::from([WithOriginalConfigurations {
399                         value: "dep-g".to_owned(),
400                         original_configurations: BTreeSet::from(["//another:platform".to_owned()]),
401                     }]),
402                 ),
403             ]),
404             unmapped: BTreeMap::from([(
405                 "cfg(pdp11)".to_owned(),
406                 BTreeSet::from(["dep-e".to_owned()]),
407             )]),
408         };
409 
410         assert_eq!(select_set, expected);
411 
412         let expected_starlark = indoc! {r#"
413             [
414                 "dep-d",
415             ] + selects.with_unmapped({
416                 "//another:platform": [
417                     "dep-g",  # //another:platform
418                 ],
419                 "@platforms//os:magic": [
420                     "dep-f",  # @platforms//os:magic
421                 ],
422                 "aarch64-macos": [
423                     "dep-a",  # cfg(macos)
424                     "dep-b",  # cfg(macos)
425                 ],
426                 "x86_64-linux": [
427                     "dep-a",  # cfg(x86_64)
428                     "dep-c",  # cfg(x86_64)
429                 ],
430                 "x86_64-macos": [
431                     "dep-a",  # cfg(macos), cfg(x86_64)
432                     "dep-b",  # cfg(macos)
433                     "dep-c",  # cfg(x86_64)
434                 ],
435                 "//conditions:default": [],
436                 selects.NO_MATCHING_PLATFORM_TRIPLES: {
437                     "cfg(pdp11)": [
438                         "dep-e",
439                     ],
440                 },
441             })
442         "#};
443 
444         assert_eq!(
445             select_set.serialize(serde_starlark::Serializer).unwrap(),
446             expected_starlark,
447         );
448     }
449 }
450