[cargo-zerocopy] Auto-install targets, more components (#3118)

Auto-install any `--target` arguments. Auto-install all of the following
components: `rust-src`, `rustfmt`, `clippy`, `miri` (nightly only)

gherrit-pr-id: Gj5ggnjidkmm6c3foelh4ou7sbi55y4qr

Author: joshlf

tools/cargo-zerocopy/src/main.rs

@@ -20,8 +20,9 @@
 // Cargo.toml section.
 
 use std::{
+    collections::HashSet,
     env, fmt, fs,
-    io::{self, Read as _},
+    io::{self, BufRead as _, Write as _},
     process::{self, Command, Output},
 };
 
@@ -126,44 +127,143 @@ fn get_toolchain_versions() -> Versions {
     }
 }
 
-fn is_toolchain_installed(versions: &Versions, name: &str) -> Result<bool, Error> {
-    let version = versions.get(name)?;
-    let output = rustup(["run", version, "cargo", "version"], None).output().unwrap();
-    if output.status.success() {
-        let output = rustup([&format!("+{version}"), "component", "list"], None).output_or_exit();
-        let stdout = String::from_utf8(output.stdout).unwrap();
-        Ok(stdout.contains("rust-src (installed)"))
-    } else {
-        Ok(false)
+fn ensure_installed_or_exit(
+    is_installed: impl FnOnce() -> Result<bool, Error>,
+    install: impl FnOnce() -> Result<(), Error>,
+    missing_item_desc: &str,
+    prompt: &str,
+) -> Result<(), Error> {
+    if is_installed()? {
+        return Ok(());
     }
-}
 
-fn install_toolchain_or_exit(versions: &Versions, name: &str) -> Result<(), Error> {
-    eprintln!("[cargo-zerocopy] missing either toolchain '{name}' or component 'rust-src'");
+    eprintln!("[cargo-zerocopy] {missing_item_desc}");
     if env::var("GITHUB_RUN_ID").is_ok() {
         eprintln!("[cargo-zerocopy] detected GitHub Actions environment; auto-installing without waiting for confirmation");
     } else if env::var("CARGO_ZEROCOPY_AUTO_INSTALL_TOOLCHAIN").is_ok() {
         eprintln!("[cargo-zerocopy] detected CARGO_ZEROCOPY_AUTO_INSTALL_TOOLCHAIN environment variable; auto-installing without waiting for confirmation");
     } else {
-        eprintln!("[cargo-zerocopy] set CARGO_ZEROCOPY_AUTO_INSTALL_TOOLCHAIN=1 to always install toolchains without prompting");
+        eprintln!("[cargo-zerocopy] set CARGO_ZEROCOPY_AUTO_INSTALL_TOOLCHAIN=1 to always install toolchains and targets without prompting");
         loop {
-            eprint!("[cargo-zerocopy] would you like to install toolchain '{name}' and component 'rust-src' via 'rustup' (y/n)? ");
-            let mut input = [0];
-            io::stdin().read_exact(&mut input).unwrap();
-            match input[0] as char {
-                'y' | 'Y' => break,
-                'n' | 'N' => process::exit(1),
-                _ => (),
+            eprint!("[cargo-zerocopy] {prompt} (y/n)? ");
+            io::stderr().flush().unwrap();
+            let mut line = String::new();
+            io::stdin().lock().read_line(&mut line).unwrap();
+            let input = line.trim().to_lowercase();
+            if input.starts_with('y') {
+                break;
+            } else if input.starts_with('n') {
+                process::exit(1);
             }
         }
     }
 
-    let version = versions.get(name)?;
-    rustup(["toolchain", "install", version, "-c", "rust-src"], None).execute();
+    install()?;
 
     Ok(())
 }
 
+fn install_toolchain_or_exit(versions: &Versions, name: &str) -> Result<(), Error> {
+    let version = versions.get(name)?.to_string();
+    let is_nightly = version.contains("nightly");
+
+    ensure_installed_or_exit(
+        || {
+            let output = rustup(["run", &version, "cargo", "version"], None).output();
+            let output = match output {
+                Ok(o) => o,
+                Err(e) => {
+                    eprintln!("[cargo-zerocopy] failed to run rustup: {e}");
+                    process::exit(1);
+                }
+            };
+            if output.status.success() {
+                let output =
+                    rustup([&format!("+{version}"), "component", "list"], None).output_or_exit();
+                let stdout = String::from_utf8(output.stdout).unwrap();
+                let is_installed =
+                    |c| stdout.lines().any(|l| l.starts_with(c) && l.contains("(installed)"));
+                let mut installed =
+                    is_installed("rust-src") && is_installed("rustfmt") && is_installed("clippy");
+                if is_nightly {
+                    installed = installed && is_installed("miri");
+                }
+                Ok(installed)
+            } else {
+                Ok(false)
+            }
+        },
+        || {
+            let mut args = vec![
+                "toolchain",
+                "install",
+                &version,
+                "-c",
+                "rust-src",
+                "-c",
+                "rustfmt",
+                "-c",
+                "clippy",
+            ];
+            if is_nightly {
+                args.push("-c");
+                args.push("miri");
+            }
+            rustup(args, None).execute();
+            Ok(())
+        },
+        &format!(
+            "missing toolchain '{name}' or one of its components (rust-src, rustfmt, clippy{})",
+            if is_nightly { ", miri" } else { "" }
+        ),
+        &format!("would you like to install toolchain '{name}' and its components via 'rustup'"),
+    )
+}
+
+fn install_targets_or_exit(version: &str, targets: &[String]) -> Result<(), Error> {
+    // Avoid running `rustup` in the common case that no `--target` arguments
+    // are provided.
+    if targets.is_empty() {
+        return Ok(());
+    }
+
+    let output = rustup(["target", "list", "--toolchain", version], None).output_or_exit();
+    let stdout = String::from_utf8(output.stdout).unwrap();
+    let mut installed = HashSet::new();
+    let mut available = HashSet::new();
+
+    for line in stdout.lines() {
+        let mut parts = line.split_whitespace();
+        if let Some(target) = parts.next() {
+            available.insert(target.to_string());
+            if parts.next() == Some("(installed)") {
+                installed.insert(target.to_string());
+            }
+        }
+    }
+
+    let to_install = targets
+        .iter()
+        .filter(|target| {
+            !installed.contains(target.as_str()) && available.contains(target.as_str())
+        })
+        .cloned()
+        .collect::<Vec<_>>();
+
+    let to_install_str = to_install.join(", ");
+    ensure_installed_or_exit(
+        || Ok(to_install.is_empty()),
+        || {
+            let mut args = vec!["target", "add", "--toolchain", version];
+            args.extend(to_install.iter().map(|s| s.as_str()));
+            rustup(args, None).execute();
+            Ok(())
+        },
+        &format!("missing target(s): {to_install_str}"),
+        &format!("would you like to install target(s) '{to_install_str}' via 'rustup'"),
+    )
+}
+
 fn get_rustflags(name: &str) -> String {
     // See #1792 for context on zerocopy_derive_union_into_bytes.
     let mut flags =
@@ -223,10 +323,40 @@ fn delegate_cargo() -> Result<(), Error> {
             if let Some(name) = arg.strip_prefix('+') {
                 let version = versions.get(name)?;
 
-                if !is_toolchain_installed(&versions, name)? {
-                    install_toolchain_or_exit(&versions, name)?;
+                install_toolchain_or_exit(&versions, name)?;
+
+                let mut targets = Vec::new();
+                if let Ok(t) = env::var("CARGO_BUILD_TARGET") {
+                    targets.push(t);
+                }
+
+                let args_vec = args.collect::<Vec<_>>();
+                let mut i = 0;
+                while i < args_vec.len() {
+                    let arg = &args_vec[i];
+                    if arg == "--" {
+                        break;
+                    }
+
+                    if arg == "--target" {
+                        if i + 1 < args_vec.len() {
+                            targets.push(args_vec[i + 1].clone());
+                            i += 1;
+                        }
+                    } else if let Some(t) = arg.strip_prefix("--target=") {
+                        targets.push(t.to_string());
+                    }
+                    i += 1;
                 }
 
+                targets.retain(|t| !t.ends_with(".json"));
+                targets.sort();
+                targets.dedup();
+
+                install_targets_or_exit(version, &targets)?;
+
+                let mut args = args_vec.into_iter();
+
                 let env_rustflags = env::vars()
                     .filter_map(|(k, v)| if k == "RUSTFLAGS" { Some(v) } else { None })
                     .next()