Make more annotations select-able (#2180)

This is a follow-up to #2139 making more annotations support
`select`-ing.

5 files changed, 108 insertions, 60 deletions

diff --git a/crate_universe/src/config.rs b/crate_universe/src/config.rs
index 245dbf7d..c4af4d03 100644
--- a/crate_universe/src/config.rs
+++ b/crate_universe/src/config.rs
@@ -144,7 +144,7 @@ impl From<GitReference> for Commitish {
 }
 /// A value which may either be a plain String, or a dict of platform triples
 /// (or other cfg expressions understood by `crate::context::platforms::resolve_cfg_platforms`) to Strings.
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
 #[serde(untagged)]
 pub enum StringOrSelect {
     Value(String),
@@ -179,11 +179,11 @@ pub struct CrateAnnotations {
 
     /// Additional data to pass to
     /// [deps](https://bazelbuild.github.io/rules_rust/defs.html#rust_library-deps) attribute.
-    pub deps: Option<BTreeSet<String>>,
+    pub deps: Option<BTreeSet<StringOrSelect>>,
 
     /// Additional data to pass to
     /// [proc_macro_deps](https://bazelbuild.github.io/rules_rust/defs.html#rust_library-proc_macro_deps) attribute.
-    pub proc_macro_deps: Option<BTreeSet<String>>,
+    pub proc_macro_deps: Option<BTreeSet<StringOrSelect>>,
 
     /// Additional data to pass to  the target's
     /// [crate_features](https://bazelbuild.github.io/rules_rust/defs.html#rust_library-crate_features) attribute.
@@ -222,19 +222,19 @@ pub struct CrateAnnotations {
 
     /// Additional dependencies to pass to a build script's
     /// [deps](https://bazelbuild.github.io/rules_rust/cargo.html#cargo_build_script-deps) attribute.
-    pub build_script_deps: Option<BTreeSet<String>>,
+    pub build_script_deps: Option<BTreeSet<StringOrSelect>>,
 
     /// Additional data to pass to a build script's
     /// [proc_macro_deps](https://bazelbuild.github.io/rules_rust/cargo.html#cargo_build_script-proc_macro_deps) attribute.
-    pub build_script_proc_macro_deps: Option<BTreeSet<String>>,
+    pub build_script_proc_macro_deps: Option<BTreeSet<StringOrSelect>>,
 
     /// Additional data to pass to a build script's
     /// [build_script_data](https://bazelbuild.github.io/rules_rust/cargo.html#cargo_build_script-data) attribute.
-    pub build_script_data: Option<BTreeSet<String>>,
+    pub build_script_data: Option<BTreeSet<StringOrSelect>>,
 
     /// Additional data to pass to a build script's
     /// [tools](https://bazelbuild.github.io/rules_rust/cargo.html#cargo_build_script-tools) attribute.
-    pub build_script_tools: Option<BTreeSet<String>>,
+    pub build_script_tools: Option<BTreeSet<StringOrSelect>>,
 
     /// An optional glob pattern to set on the
     /// [build_script_data](https://bazelbuild.github.io/rules_rust/cargo.html#cargo_build_script-build_script_env) attribute.
@@ -246,7 +246,7 @@ pub struct CrateAnnotations {
 
     /// Additional rustc_env flags to pass to a build script's
     /// [rustc_env](https://bazelbuild.github.io/rules_rust/cargo.html#cargo_build_script-rustc_env) attribute.
-    pub build_script_rustc_env: Option<BTreeMap<String, String>>,
+    pub build_script_rustc_env: Option<BTreeMap<String, StringOrSelect>>,
 
     /// Additional labels to pass to a build script's
     /// [toolchains](https://bazel.build/reference/be/common-definitions#common-attributes) attribute.
@@ -374,7 +374,7 @@ pub struct AnnotationsProvidedByPackage {
     pub rustc_env_files: Option<BTreeSet<String>>,
     pub rustc_flags: Option<Vec<String>>,
     pub build_script_env: Option<BTreeMap<String, StringOrSelect>>,
-    pub build_script_rustc_env: Option<BTreeMap<String, String>>,
+    pub build_script_rustc_env: Option<BTreeMap<String, StringOrSelect>>,
     pub build_script_rundir: Option<String>,
     pub additive_build_file_content: Option<String>,
     pub extra_aliased_targets: Option<BTreeMap<String, String>>,
diff --git a/crate_universe/src/context/crate_context.rs b/crate_universe/src/context/crate_context.rs
index 8488eb6d..d4301e89 100644
--- a/crate_universe/src/context/crate_context.rs
+++ b/crate_universe/src/context/crate_context.rs
@@ -5,7 +5,7 @@ use std::collections::{BTreeMap, BTreeSet};
 use cargo_metadata::{Node, Package, PackageId};
 use serde::{Deserialize, Serialize};
 
-use crate::config::{CrateId, GenBinaries, StringOrSelect};
+use crate::config::{CrateId, GenBinaries};
 use crate::metadata::{CrateAnnotation, Dependency, PairredExtras, SourceAnnotation};
 use crate::utils::sanitize_module_name;
 use crate::utils::starlark::{Glob, SelectList, SelectMap, SelectStringDict, SelectStringList};
@@ -122,8 +122,8 @@ pub struct CommonAttributes {
     #[serde(skip_serializing_if = "SelectList::is_empty")]
     pub deps: SelectList<CrateDependency>,
 
-    #[serde(skip_serializing_if = "BTreeSet::is_empty")]
-    pub extra_deps: BTreeSet<String>,
+    #[serde(skip_serializing_if = "SelectList::is_empty")]
+    pub extra_deps: SelectList<String>,
 
     #[serde(skip_serializing_if = "SelectList::is_empty")]
     pub deps_dev: SelectList<CrateDependency>,
@@ -136,8 +136,8 @@ pub struct CommonAttributes {
     #[serde(skip_serializing_if = "SelectList::is_empty")]
     pub proc_macro_deps: SelectList<CrateDependency>,
 
-    #[serde(skip_serializing_if = "BTreeSet::is_empty")]
-    pub extra_proc_macro_deps: BTreeSet<String>,
+    #[serde(skip_serializing_if = "SelectList::is_empty")]
+    pub extra_proc_macro_deps: SelectList<String>,
 
     #[serde(skip_serializing_if = "SelectList::is_empty")]
     pub proc_macro_deps_dev: SelectList<CrateDependency>,
@@ -200,8 +200,8 @@ pub struct BuildScriptAttributes {
     #[serde(skip_serializing_if = "SelectList::is_empty")]
     pub deps: SelectList<CrateDependency>,
 
-    #[serde(skip_serializing_if = "BTreeSet::is_empty")]
-    pub extra_deps: BTreeSet<String>,
+    #[serde(skip_serializing_if = "SelectList::is_empty")]
+    pub extra_deps: SelectStringList,
 
     // TODO: refactor a crate with a build.rs file from two into three bazel
     // rules in order to deduplicate link_dep information. Currently as the
@@ -223,8 +223,8 @@ pub struct BuildScriptAttributes {
     #[serde(skip_serializing_if = "SelectList::is_empty")]
     pub link_deps: SelectList<CrateDependency>,
 
-    #[serde(skip_serializing_if = "BTreeSet::is_empty")]
-    pub extra_link_deps: BTreeSet<String>,
+    #[serde(skip_serializing_if = "SelectList::is_empty")]
+    pub extra_link_deps: SelectStringList,
 
     #[serde(skip_serializing_if = "SelectStringDict::is_empty")]
     pub build_script_env: SelectStringDict,
@@ -232,8 +232,8 @@ pub struct BuildScriptAttributes {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub rundir: Option<String>,
 
-    #[serde(skip_serializing_if = "BTreeSet::is_empty")]
-    pub extra_proc_macro_deps: BTreeSet<String>,
+    #[serde(skip_serializing_if = "SelectList::is_empty")]
+    pub extra_proc_macro_deps: SelectStringList,
 
     #[serde(skip_serializing_if = "SelectList::is_empty")]
     pub proc_macro_deps: SelectList<CrateDependency>,
@@ -494,12 +494,14 @@ impl CrateContext {
 
             // Deps
             if let Some(extra) = &crate_extra.deps {
-                self.common_attrs.extra_deps = extra.clone();
+                self.common_attrs.extra_deps.extend(extra.iter().cloned());
             }
 
             // Proc macro deps
             if let Some(extra) = &crate_extra.proc_macro_deps {
-                self.common_attrs.extra_proc_macro_deps = extra.clone();
+                self.common_attrs
+                    .extra_proc_macro_deps
+                    .extend(extra.iter().cloned());
             }
 
             // Compile data
@@ -564,26 +566,22 @@ impl CrateContext {
             if let Some(attrs) = &mut self.build_script_attrs {
                 // Deps
                 if let Some(extra) = &crate_extra.build_script_deps {
-                    attrs.extra_deps = extra.clone();
+                    attrs.extra_deps.extend(extra.iter().cloned())
                 }
 
                 // Proc macro deps
                 if let Some(extra) = &crate_extra.build_script_proc_macro_deps {
-                    attrs.extra_proc_macro_deps = extra.clone();
+                    attrs.extra_proc_macro_deps.extend(extra.iter().cloned());
                 }
 
                 // Data
                 if let Some(extra) = &crate_extra.build_script_data {
-                    for data in extra {
-                        attrs.data.insert(data.clone(), None);
-                    }
+                    attrs.data.extend(extra.iter().cloned());
                 }
 
                 // Tools
                 if let Some(extra) = &crate_extra.build_script_tools {
-                    for data in extra {
-                        attrs.tools.insert(data.clone(), None);
-                    }
+                    attrs.tools.extend(extra.iter().cloned());
                 }
 
                 // Toolchains
@@ -600,29 +598,16 @@ impl CrateContext {
 
                 // Rustc env
                 if let Some(extra) = &crate_extra.build_script_rustc_env {
-                    attrs.rustc_env.extend(extra.clone(), None);
+                    attrs
+                        .rustc_env
+                        .extend_from_string_or_select(extra.iter().map(clone_tuple));
                 }
 
                 // Build script env
                 if let Some(extra) = &crate_extra.build_script_env {
-                    for (key, value) in extra {
-                        match value {
-                            StringOrSelect::Value(value) => {
-                                attrs
-                                    .build_script_env
-                                    .insert(key.clone(), value.clone(), None);
-                            }
-                            StringOrSelect::Select(select) => {
-                                for (select_key, value) in select {
-                                    attrs.build_script_env.insert(
-                                        key.clone(),
-                                        value.clone(),
-                                        Some(select_key.clone()),
-                                    );
-                                }
-                            }
-                        }
-                    }
+                    attrs
+                        .build_script_env
+                        .extend_from_string_or_select(extra.iter().map(clone_tuple))
                 }
 
                 if let Some(rundir) = &crate_extra.build_script_rundir {
@@ -771,6 +756,11 @@ impl CrateContext {
     }
 }
 
+fn clone_tuple<V1: Clone, V2: Clone>(t: (&V1, &V2)) -> (V1, V2) {
+    let (v1, v2) = t;
+    (v1.clone(), v2.clone())
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
diff --git a/crate_universe/src/rendering.rs b/crate_universe/src/rendering.rs
index 90bbcf1a..4a0884ba 100644
--- a/crate_universe/src/rendering.rs
+++ b/crate_universe/src/rendering.rs
@@ -385,13 +385,13 @@ impl Renderer {
             deps: self
                 .make_deps(
                     attrs.map_or(&empty_deps, |attrs| &attrs.deps),
-                    attrs.map_or(&empty_set, |attrs| &attrs.extra_deps),
+                    attrs.map_or(&empty_list, |attrs| &attrs.extra_deps),
                 )
                 .remap_configurations(platforms),
             link_deps: self
                 .make_deps(
                     attrs.map_or(&empty_deps, |attrs| &attrs.link_deps),
-                    attrs.map_or(&empty_set, |attrs| &attrs.extra_link_deps),
+                    attrs.map_or(&empty_list, |attrs| &attrs.extra_link_deps),
                 )
                 .remap_configurations(platforms),
             edition: krate.common_attrs.edition.clone(),
@@ -400,7 +400,7 @@ impl Renderer {
             proc_macro_deps: self
                 .make_deps(
                     attrs.map_or(&empty_deps, |attrs| &attrs.proc_macro_deps),
-                    attrs.map_or(&empty_set, |attrs| &attrs.extra_proc_macro_deps),
+                    attrs.map_or(&empty_list, |attrs| &attrs.extra_proc_macro_deps),
                 )
                 .remap_configurations(platforms),
             rundir: attrs.and_then(|attrs| attrs.rundir.clone()),
@@ -628,14 +628,12 @@ impl Renderer {
     fn make_deps(
         &self,
         deps: &SelectList<CrateDependency>,
-        extra_deps: &BTreeSet<String>,
+        extra_deps: &SelectList<String>,
     ) -> SelectList<String> {
         let mut deps = deps
             .clone()
             .map(|dep| self.crate_label(&dep.id.name, &dep.id.version, &dep.target));
-        for extra_dep in extra_deps {
-            deps.insert(extra_dep.clone(), None);
-        }
+        deps.extend(extra_deps.into_iter());
         deps
     }
 
diff --git a/crate_universe/src/utils/starlark/select.rs b/crate_universe/src/utils/starlark/select.rs
index 091057a0..a6049b6b 100644
--- a/crate_universe/src/utils/starlark/select.rs
+++ b/crate_universe/src/utils/starlark/select.rs
@@ -1,6 +1,7 @@
 use std::collections::{btree_set, BTreeMap, BTreeSet};
 use std::iter::{once, FromIterator};
 
+use crate::config::StringOrSelect;
 use serde::ser::{SerializeMap, SerializeTupleStruct, Serializer};
 use serde::{Deserialize, Serialize};
 use serde_starlark::{FunctionCall, LineComment, MULTILINE};
@@ -92,6 +93,42 @@ impl<T: Ord> SelectList<T> {
     }
 }
 
+impl SelectList<String> {
+    pub fn extend<Iter: Iterator<Item = StringOrSelect>>(&mut self, values: Iter) {
+        for value in values {
+            match value {
+                StringOrSelect::Value(value) => {
+                    self.insert(value, None);
+                }
+                StringOrSelect::Select(select) => {
+                    for (select_key, value) in select {
+                        self.insert(value.clone(), Some(select_key.clone()));
+                    }
+                }
+            }
+        }
+    }
+}
+
+impl IntoIterator for &SelectList<String> {
+    type Item = StringOrSelect;
+    type IntoIter = <Vec<StringOrSelect> as IntoIterator>::IntoIter;
+    fn into_iter(self) -> Self::IntoIter {
+        let mut all_values = Vec::with_capacity(self.common.len() + self.selects.len());
+        for value in &self.common {
+            all_values.push(StringOrSelect::Value(value.clone()))
+        }
+        for (key, values) in &self.selects {
+            for value in values {
+                let mut map = BTreeMap::new();
+                map.insert(key.clone(), value.clone());
+                all_values.push(StringOrSelect::Select(map))
+            }
+        }
+        all_values.into_iter()
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
 pub struct WithOriginalConfigurations<T> {
     value: T,
@@ -371,6 +408,26 @@ impl<T: Ord> SelectDict<T> {
     }
 }
 
+impl SelectDict<String> {
+    pub fn extend_from_string_or_select<Iter: Iterator<Item = (String, StringOrSelect)>>(
+        &mut self,
+        values: Iter,
+    ) {
+        for (key, value) in values {
+            match value {
+                StringOrSelect::Value(value) => {
+                    self.insert(key, value, None);
+                }
+                StringOrSelect::Select(select) => {
+                    for (select_key, value) in select {
+                        self.insert(key.clone(), value, Some(select_key));
+                    }
+                }
+            }
+        }
+    }
+}
+
 impl<T: Ord + Clone> SelectDict<T> {
     /// Generates a new SelectDict re-keyed by the given configuration mapping.
     /// This mapping maps from configurations in the current SelectDict to sets
diff --git a/examples/crate_universe/multi_package/cargo-bazel-lock.json b/examples/crate_universe/multi_package/cargo-bazel-lock.json
index fa78b29f..1b5aa5ec 100644
--- a/examples/crate_universe/multi_package/cargo-bazel-lock.json
+++ b/examples/crate_universe/multi_package/cargo-bazel-lock.json
@@ -1,5 +1,5 @@
 {
-  "checksum": "3f78847a3b168aa8bb22b43e50ed441e484f0c0b4196a2871d8f5f0e632e2388",
+  "checksum": "73b1eb70f910e037343d0afa28af9739b96806eff40e0708404e8ec749848ed8",
   "crates": {
     "aho-corasick 0.7.20": {
       "name": "aho-corasick",
@@ -1584,9 +1584,12 @@
             ]
           }
         },
-        "extra_deps": [
-          "@m_pkgs__curl//:curl"
-        ],
+        "extra_deps": {
+          "common": [
+            "@m_pkgs__curl//:curl"
+          ],
+          "selects": {}
+        },
         "edition": "2018",
         "version": "0.4.60+curl-7.88.1"
       },