From 1aaae04e76dd53d009d8a9cd7979a4cba63b7789 Mon Sep 17 00:00:00 2001
From: shannmu <shanmu1901@gmail.com>
Date: Tue, 17 Sep 2024 16:14:44 +0800
Subject: [PATCH] feat: Add custom completer for completing `cargo build
 --packge <TAB>`

---
 src/cargo/util/command_prelude.rs | 30 +++++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/src/cargo/util/command_prelude.rs b/src/cargo/util/command_prelude.rs
index 4c8656bae182..4c975b5442e8 100644
--- a/src/cargo/util/command_prelude.rs
+++ b/src/cargo/util/command_prelude.rs
@@ -1,6 +1,6 @@
 use crate::core::compiler::{BuildConfig, MessageFormat, TimingOutput};
 use crate::core::resolver::CliFeatures;
-use crate::core::{shell, Edition, Target, TargetKind, Workspace};
+use crate::core::{shell, Edition, Package, Target, TargetKind, Workspace};
 use crate::ops::lockfile::LOCKFILE_NAME;
 use crate::ops::registry::RegistryOrIndex;
 use crate::ops::{CompileFilter, CompileOptions, NewOptions, Packages, VersionControl};
@@ -86,7 +86,10 @@ pub trait CommandExt: Sized {
         self._arg(
             optional_multi_opt("package", "SPEC", package)
                 .short('p')
-                .help_heading(heading::PACKAGE_SELECTION),
+                .help_heading(heading::PACKAGE_SELECTION)
+                .add(clap_complete::ArgValueCandidates::new(
+                    get_package_candidates,
+                )),
         )
     }
 
@@ -95,7 +98,10 @@ pub trait CommandExt: Sized {
             optional_opt("package", package)
                 .short('p')
                 .value_name("SPEC")
-                .help_heading(heading::PACKAGE_SELECTION),
+                .help_heading(heading::PACKAGE_SELECTION)
+                .add(clap_complete::ArgValueCandidates::new(
+                    get_package_candidates,
+                )),
         )
     }
 
@@ -1067,6 +1073,24 @@ fn get_targets_from_metadata() -> CargoResult<Vec<Target>> {
     Ok(targets)
 }
 
+fn get_package_candidates() -> Vec<clap_complete::CompletionCandidate> {
+    get_packages_from_metadata()
+        .unwrap_or_default()
+        .into_iter()
+        .map(|pkg| clap_complete::CompletionCandidate::new(pkg.name().as_str()))
+        .collect::<Vec<_>>()
+}
+
+fn get_packages_from_metadata() -> CargoResult<Vec<Package>> {
+    let cwd = std::env::current_dir()?;
+    let gctx = GlobalContext::new(shell::Shell::new(), cwd.clone(), cargo_home_with_cwd(&cwd)?);
+    let ws = Workspace::new(&find_root_manifest_for_wd(&cwd)?, &gctx)?;
+
+    let packages = ws.members().map(|pkg| pkg.clone()).collect::<Vec<_>>();
+
+    Ok(packages)
+}
+
 #[track_caller]
 pub fn ignore_unknown<T: Default>(r: Result<T, clap::parser::MatchesError>) -> T {
     match r {