toolchain: refactor, check installed components, install everything as one cmd

Author: Firestar99

crates/cargo-gpu-install/src/install_toolchain.rs

@@ -1,10 +1,16 @@
 //! toolchain installation logic
 
+use crate::user_output;
 use anyhow::Context as _;
 #[cfg(feature = "tty")]
 use crossterm::tty::IsTty as _;
+use std::collections::HashSet;
+use std::process::Command;
+use std::string::FromUtf8Error;
 
-use crate::user_output;
+/// list of required rustup components
+pub const REQUIRED_COMPONENTS: &[&str] =
+    ["rust-src", "rustc-dev", "llvm-tools", "clippy"].as_slice();
 
 /// Use `rustup` to install the toolchain and components, if not already installed.
 ///
@@ -16,87 +22,134 @@ pub fn ensure_toolchain_and_components_exist(
     channel: &str,
     skip_toolchain_install_consent: bool,
 ) -> anyhow::Result<()> {
-    // Check for the required toolchain
-    let output_toolchain_list = std::process::Command::new("rustup")
-        .args(["toolchain", "list"])
-        .output()
-        .context("running rustup command")?;
-    anyhow::ensure!(
-        output_toolchain_list.status.success(),
-        "could not list installed toolchains"
-    );
-    let string_toolchain_list = String::from_utf8_lossy(&output_toolchain_list.stdout);
-    if string_toolchain_list
-        .split_whitespace()
-        .any(|toolchain| toolchain.starts_with(channel))
-    {
-        log::debug!("toolchain {channel} is already installed");
-    } else {
-        let message = format!("Rust {channel} with `rustup`");
+    // While our channel may be `nightly-2024-04-24`, it'll be resolved to the full toolchain name of e.g.
+    // `nightly-2024-04-24-aarch64-unknown-linux-gnu` and that's also what `rustup toolchain list` will print.
+    // Only checking whether the toolchain starts with the channel name may incorrectly pass if you have a toolchain
+    // installed that you're not able to run on your system via `rustup toolchain install --force-non-host ...`.
+    // CMD: `rustc --print host-tuple`
+    // TODO: What if the user has no toolchain installed? You can't query this with rustup sady.
+    let (host_tuple, _) = run_cmd(Command::new("rustc").args(["--print", "host-tuple"]))?;
+    let host_tuple = host_tuple.trim_ascii();
+    let toolchain = format!("{channel}-{host_tuple}");
+
+    if !is_toolchain_installed(&toolchain, host_tuple)? {
+        let message = format!(
+            "toolchain {channel} with components {}",
+            intersperse(", ", REQUIRED_COMPONENTS.iter().copied())
+        );
         get_consent_for_toolchain_install(
             format!("Install {message}").as_ref(),
             skip_toolchain_install_consent,
         )?;
-        crate::user_output!("Installing {message}\n");
-
-        let output_toolchain_add = std::process::Command::new("rustup")
-            .args(["toolchain", "add"])
-            .arg(channel)
-            .stdout(std::process::Stdio::inherit())
-            .stderr(std::process::Stdio::inherit())
-            .output()
-            .context("adding toolchain")?;
-        anyhow::ensure!(
-            output_toolchain_add.status.success(),
-            "could not install required toolchain"
+        user_output!("Installing {message}\n");
+
+        // component list may be out of sync
+        // CMD: `rustup toolchain install nightly-2024-04-24 -c clippy,rust-src,rustc-dev,llvm-tools`
+        run_cmd(
+            Command::new("rustup")
+                .args([
+                    "toolchain",
+                    "install",
+                    &toolchain,
+                    "-c",
+                    &intersperse(",", REQUIRED_COMPONENTS.iter().copied()),
+                ])
+                .stdout(std::process::Stdio::inherit())
+                .stderr(std::process::Stdio::inherit()),
+        )?;
+    }
+
+    Ok(())
+}
+
+/// Returns true if the toolchain and required components are installed.
+fn is_toolchain_installed(toolchain: &str, host_tuple: &str) -> anyhow::Result<bool> {
+    // check if toolchain is installed
+    // CMD: `rustup toolchain list -q`
+    let (list_toolchains, _) = run_cmd(Command::new("rustup").args(["toolchain", "list", "-q"]))?;
+    if !list_toolchains
+        .split_ascii_whitespace()
+        .any(|s| s == toolchain)
+    {
+        log::info!("toolchain {toolchain} is not installed");
+        return Ok(false);
+    }
+
+    // check if required components are installed
+    // NOTE: checking for components will install the toolchain with the default profile, if not already installed!
+    // So we must check beforehand whether the toolchain is installed, to not accidentally install it here.
+    // Passing *just* `-q` will list available components, so add `--installed` for installed components.
+    // CMD: `rustup component list --toolchain nightly-2024-04-24-aarch64-unknown-linux-gnu -q --installed`
+    let (components, _) = run_cmd(Command::new("rustup").args([
+        "component",
+        "list",
+        "--toolchain",
+        toolchain,
+        "-q",
+        "--installed",
+    ]))?;
+
+    // components are listed as:
+    // * `llvm-tools-aarch64-unknown-linux-gnu` and we need to snippet off the host tuple from the end
+    // * `rust-src` since source code isn't target dependent
+    let component_host_suffix = format!("-{host_tuple}");
+    let mut required = REQUIRED_COMPONENTS.iter().copied().collect::<HashSet<_>>();
+    for component in components.split_ascii_whitespace() {
+        required.remove(
+            component
+                .strip_suffix(&component_host_suffix)
+                .unwrap_or(component),
         );
     }
+    if !required.is_empty() {
+        log::info!("components {required:?} missing for toolchain {toolchain}");
+        return Ok(false);
+    }
 
-    // Check for the required components
-    let output_component_list = std::process::Command::new("rustup")
-        .args(["component", "list", "--toolchain"])
-        .arg(channel)
-        .output()
-        .context("getting toolchain list")?;
-    anyhow::ensure!(
-        output_component_list.status.success(),
-        "could not list installed components"
-    );
-    let string_component_list = String::from_utf8_lossy(&output_component_list.stdout);
-    let required_components = ["rust-src", "rustc-dev", "llvm-tools"];
-    let installed_components = string_component_list.lines().collect::<Vec<_>>();
-    let all_components_installed = required_components.iter().all(|component| {
-        installed_components.iter().any(|installed_component| {
-            let is_component = installed_component.starts_with(component);
-            let is_installed = installed_component.ends_with("(installed)");
-            is_component && is_installed
-        })
-    });
-    if all_components_installed {
-        log::debug!("all required components are installed");
-    } else {
-        let message = "toolchain components [rust-src, rustc-dev, llvm-tools] with `rustup`";
-        get_consent_for_toolchain_install(
-            format!("Install {message}").as_ref(),
-            skip_toolchain_install_consent,
-        )?;
-        crate::user_output!("Installing {message}\n");
-
-        let output_component_add = std::process::Command::new("rustup")
-            .args(["component", "add", "--toolchain"])
-            .arg(channel)
-            .args(["rust-src", "rustc-dev", "llvm-tools"])
-            .stdout(std::process::Stdio::inherit())
-            .stderr(std::process::Stdio::inherit())
-            .output()
-            .context("adding rustup component")?;
-        anyhow::ensure!(
-            output_component_add.status.success(),
-            "could not install required components"
+    log::info!("toolchain and required components are already installed");
+    Ok(true)
+}
+
+pub fn run_cmd(cmd: &mut Command) -> anyhow::Result<(String, String)> {
+    let output = cmd.output();
+    let fmt_cmd = || {
+        intersperse(
+            " ",
+            std::iter::once(cmd.get_program())
+                .chain(cmd.get_args())
+                .map(|s| s.to_str().unwrap()),
+        )
+    };
+    let output = output.with_context(|| format!("Failed to launch cmd `{}`", fmt_cmd()))?;
+
+    let utf8_error = |e: FromUtf8Error, kind: &str| {
+        anyhow::anyhow!(
+            "Command `{}` {} contains invalid UTF-8: {} \n {:?}",
+            kind,
+            fmt_cmd(),
+            e.utf8_error(),
+            e.into_bytes()
+        )
+    };
+    let stdout = String::from_utf8(output.stdout).map_err(|e| utf8_error(e, "stdout"))?;
+    let stderr = String::from_utf8(output.stderr).map_err(|e| utf8_error(e, "stderr"))?;
+
+    if !output.status.success() {
+        anyhow::bail!(
+            "Command `{}` failed with {}:\n-- stdout\n{stdout}\n-- stderr\n{stderr}",
+            fmt_cmd(),
+            &output.status,
         );
     }
+    Ok((stdout, stderr))
+}
 
-    Ok(())
+/// Folds an [`Iterator`] of `&str` into a [`String`] while interspersing some `&str` between each element
+#[expect(clippy::string_add, reason = "Deliberately using String::add")]
+fn intersperse<'a>(intersperse: &str, iter: impl Iterator<Item = &'a str>) -> String {
+    let mut s = iter.fold(String::new(), |a, b| a + b + intersperse);
+    s.truncate(s.len() - intersperse.len());
+    s
 }
 
 #[cfg(not(feature = "tty"))]