Use clap dynamic completion for +toolchain

Cargo.toml

@@ -47,7 +47,7 @@ cc = "1"
 chrono = { version = "0.4", default-features = false, features = ["std"] }
 clap = { version = "4", features = ["derive", "wrap_help", "string"] }
 clap-cargo = "0.19.0"
-clap_complete = "4"
+clap_complete = { version = "4", features = ["unstable-dynamic"] }
 console = "0.16"
 effective-limits = "0.5.5"
 enum-map = "3.0.0"

src/cli/rustup_mode.rs

@@ -1,6 +1,7 @@
 use std::{
     borrow::Cow,
     env::consts::EXE_SUFFIX,
+    ffi::OsStr,
     fmt,
     io::{self, Write},
     path::{Path, PathBuf},
@@ -18,7 +19,10 @@ use clap::{
     builder::{PossibleValue, ValueHint},
 };
 use clap_cargo::style::{CONTEXT, ERROR, GOOD, HEADER, TRANSIENT, WARN};
-use clap_complete::Shell;
+use clap_complete::{
+    Shell,
+    engine::{ArgValueCompleter, CompletionCandidate},
+};
 use futures_util::stream::StreamExt;
 use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
 use itertools::Itertools;
@@ -672,6 +676,12 @@ pub async fn main(
     process: &Process,
     console_filter: Handle<EnvFilter, Registry>,
 ) -> Result<ExitCode> {
+    let cfg = &mut Cfg::from_env(current_dir, true, false, process)?;
+    clap_complete::CompleteEnv::with_factory(|| completion_command(cfg))
+        .var("RUSTUP_COMPLETE")
+        .bin("rustup")
+        .complete();
+
     self_update::cleanup_self_updater(process)?;
 
     use clap::error::ErrorKind::*;
@@ -683,7 +693,7 @@ pub async fn main(
         }
         Err(err) if err.kind() == DisplayVersion => {
             write!(process.stdout().lock(), "{}", err.render().ansi())?;
-            display_version(current_dir, process).await?;
+            display_version(cfg).await?;
             return Ok(ExitCode::SUCCESS);
         }
         Err(err) => {
@@ -710,12 +720,8 @@ pub async fn main(
         return Ok(ExitCode::FAILURE);
     };
 
-    let cfg = &mut Cfg::from_env(
-        current_dir,
-        matches.quiet,
-        subcmd.allow_auto_install(),
-        process,
-    )?;
+    cfg.quiet = matches.quiet;
+    cfg.allow_auto_install = subcmd.allow_auto_install();
     cfg.toolchain_override = matches.plus_toolchain;
 
     let should_warn = subcmd.should_warn_empty_setup();
@@ -870,6 +876,26 @@ pub async fn main(
     Ok(exit_code)
 }
 
+fn completion_command(cfg: &Cfg<'_>) -> clap::Command {
+    let toolchains = cfg.list_toolchains().unwrap_or_default();
+    Rustup::command().mut_arg("+toolchain", move |arg| {
+        arg.add(ArgValueCompleter::new(move |current: &OsStr| {
+            let Some(prefix) = current.to_str() else {
+                return Vec::new();
+            };
+            toolchains
+                .iter()
+                .filter_map(|toolchain| {
+                    let candidate = format!("+{toolchain}");
+                    candidate
+                        .starts_with(prefix)
+                        .then(|| CompletionCandidate::new(candidate))
+                })
+                .collect()
+        }))
+    })
+}
+
 async fn default_(
     cfg: &Cfg<'_>,
     toolchain: Option<MaybeResolvableToolchainName>,
@@ -1835,17 +1861,16 @@ fn output_completion_script(
     Ok(ExitCode::SUCCESS)
 }
 
-async fn display_version(current_dir: PathBuf, process: &Process) -> Result<()> {
+async fn display_version(cfg: &mut Cfg<'_>) -> Result<()> {
     info!("this is the version for the rustup toolchain manager, not the rustc compiler");
-    let mut cfg = Cfg::from_env(current_dir, true, false, process)?;
     cfg.toolchain_override = cfg
         .process
         .args()
         .find_map(|arg| arg.strip_prefix('+').map(ResolvableToolchainName::from_str))
         .transpose()?;
 
     match cfg.maybe_ensure_active_toolchain(None).await {
-        Ok(Some((name, _))) => match Toolchain::new(&cfg, name) {
+        Ok(Some((name, _))) => match Toolchain::new(cfg, name) {
             Ok(tc) => info!(
                 "the currently active `rustc` version is `{}`",
                 tc.rustc_version()
@@ -1864,3 +1889,76 @@ async fn display_version(current_dir: PathBuf, process: &Process) -> Result<()>
 
     Ok(())
 }
+
+#[cfg(all(test, feature = "test"))]
+mod tests {
+    use std::{
+        collections::{HashMap, HashSet},
+        ffi::OsString,
+        fs,
+    };
+
+    use super::completion_command;
+    use crate::{config::Cfg, process::TestProcess};
+
+    fn complete(cfg: &Cfg<'_>, args: &[&str], index: usize) -> Vec<String> {
+        let mut command = completion_command(cfg);
+        clap_complete::engine::complete(
+            &mut command,
+            args.iter().map(OsString::from).collect(),
+            index,
+            Some(&cfg.current_dir),
+        )
+        .unwrap()
+        .into_iter()
+        .map(|candidate| candidate.get_value().to_string_lossy().into_owned())
+        .collect()
+    }
+
+    #[test]
+    fn dynamic_completion_distinguishes_toolchains_and_subcommands() {
+        let rustup_home = tempfile::tempdir().unwrap();
+        fs::create_dir_all(rustup_home.path().join("toolchains/custom")).unwrap();
+        let vars = HashMap::from([
+            (
+                "RUSTUP_HOME".to_owned(),
+                rustup_home.path().display().to_string(),
+            ),
+            (
+                "RUSTUP_OVERRIDE_UNIX_FALLBACK_SETTINGS".to_owned(),
+                rustup_home
+                    .path()
+                    .join("missing-settings.toml")
+                    .display()
+                    .to_string(),
+            ),
+        ]);
+        let process = TestProcess::new(rustup_home.path(), &["rustup"], vars, "").process;
+        let cfg = Cfg::from_env(rustup_home.path().to_owned(), true, false, &process).unwrap();
+
+        assert_eq!(complete(&cfg, &["rustup", "+cus"], 1), ["+custom"]);
+
+        for (args, index) in [
+            (&["rustup", "component", ""][..], 2),
+            (&["rustup", "+custom", "component", ""][..], 3),
+        ] {
+            let candidates = complete(&cfg, args, index);
+            let candidates = candidates
+                .iter()
+                .map(String::as_str)
+                .collect::<HashSet<_>>();
+            let expected = HashSet::from(["list", "add", "remove"]);
+            assert!(
+                candidates.is_superset(&expected),
+                "missing component subcommands for {args:?}: {:?}",
+                expected.difference(&candidates).collect::<Vec<_>>()
+            );
+            let root_only = HashSet::from(["install", "toolchain", "show"]);
+            assert!(
+                candidates.is_disjoint(&root_only),
+                "unexpected root-only subcommands for {args:?}: {:?}",
+                candidates.intersection(&root_only).collect::<Vec<_>>()
+            );
+        }
+    }
+}

tests/suite/cli_misc.rs

@@ -1245,10 +1245,28 @@ async fn nightly_backtrack_skips_missing() {
 #[tokio::test]
 async fn completion_rustup() {
     let cx = CliTestContext::new(Scenario::SimpleV2).await;
-    cx.config
+    let output = cx
+        .config
         .expect(["rustup", "completions", "bash", "rustup"])
-        .await
-        .is_ok();
+        .await;
+    output.is_ok();
+    assert!(output.output.stdout.contains("_rustup()"));
+    assert!(!output.output.stdout.contains("RUSTUP_COMPLETE"));
+}
+
+#[tokio::test]
+async fn dynamic_completion_rustup() {
+    let cx = CliTestContext::new(Scenario::SimpleV2).await;
+    let output = cx
+        .config
+        .cmd("rustup", std::iter::empty::<&str>())
+        .env("RUSTUP_COMPLETE", "bash")
+        .output()
+        .unwrap();
+    assert!(output.status.success());
+    let stdout = String::from_utf8(output.stdout).unwrap();
+    assert!(stdout.contains("RUSTUP_COMPLETE=\"bash\""));
+    assert!(stdout.contains("_clap_complete_rustup()"));
 }
 
 #[tokio::test]