3 files changed, 74 insertions, 102 deletions

diff --git a/xtask/src/flags.rs b/xtask/src/flags.rs
index 5710fbdb5..56eda5b1e 100644
--- a/xtask/src/flags.rs
+++ b/xtask/src/flags.rs
@@ -1,5 +1,7 @@
 #![allow(unreachable_pub)]
 
+use crate::install::{ClientOpt, Malloc, ServerOpt};
+
 xflags::args_parser! {
     /// Run custom build command.
     cmd xtask {
@@ -137,3 +139,25 @@ impl Xtask {
     }
 }
 // generated end
+
+impl Install {
+    pub(crate) fn server(&self) -> Option<ServerOpt> {
+        if self.client && !self.server {
+            return None;
+        }
+        let malloc = if self.mimalloc {
+            Malloc::Mimalloc
+        } else if self.jemalloc {
+            Malloc::Jemalloc
+        } else {
+            Malloc::System
+        };
+        Some(ServerOpt { malloc })
+    }
+    pub(crate) fn client(&self) -> Option<ClientOpt> {
+        if !self.client && self.server {
+            return None;
+        }
+        Some(ClientOpt { code_bin: self.code_bin.clone() })
+    }
+}
diff --git a/xtask/src/install.rs b/xtask/src/install.rs
index ea2194248..177028b08 100644
--- a/xtask/src/install.rs
+++ b/xtask/src/install.rs
@@ -5,60 +5,32 @@ use std::{env, path::PathBuf, str};
 use anyhow::{bail, format_err, Context, Result};
 use xshell::{cmd, pushd};
 
+use crate::flags;
+
 // Latest stable, feel free to send a PR if this lags behind.
 const REQUIRED_RUST_VERSION: u32 = 50;
 
-pub(crate) struct InstallCmd {
-    pub(crate) client: Option<ClientOpt>,
-    pub(crate) server: Option<ServerOpt>,
-}
-
-#[derive(Clone, Copy)]
-pub(crate) enum ClientOpt {
-    VsCode,
-    VsCodeExploration,
-    VsCodeInsiders,
-    VsCodium,
-    VsCodeOss,
-    Any,
-}
-
-impl ClientOpt {
-    pub(crate) const fn as_cmds(&self) -> &'static [&'static str] {
-        match self {
-            ClientOpt::VsCode => &["code"],
-            ClientOpt::VsCodeExploration => &["code-exploration"],
-            ClientOpt::VsCodeInsiders => &["code-insiders"],
-            ClientOpt::VsCodium => &["codium"],
-            ClientOpt::VsCodeOss => &["code-oss"],
-            ClientOpt::Any => &["code", "code-exploration", "code-insiders", "codium", "code-oss"],
+impl flags::Install {
+    pub(crate) fn run(self) -> Result<()> {
+        if cfg!(target_os = "macos") {
+            fix_path_for_mac().context("Fix path for mac")?
+        }
+        if let Some(server) = self.server() {
+            install_server(server).context("install server")?;
+        }
+        if let Some(client) = self.client() {
+            install_client(client).context("install client")?;
         }
+        Ok(())
     }
 }
 
-impl Default for ClientOpt {
-    fn default() -> Self {
-        ClientOpt::Any
-    }
+#[derive(Clone)]
+pub(crate) struct ClientOpt {
+    pub(crate) code_bin: Option<String>,
 }
 
-impl std::str::FromStr for ClientOpt {
-    type Err = anyhow::Error;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        [
-            ClientOpt::VsCode,
-            ClientOpt::VsCodeExploration,
-            ClientOpt::VsCodeInsiders,
-            ClientOpt::VsCodium,
-            ClientOpt::VsCodeOss,
-        ]
-        .iter()
-        .copied()
-        .find(|c| [s] == c.as_cmds())
-        .ok_or_else(|| anyhow::format_err!("no such client"))
-    }
-}
+const VS_CODES: &[&str] = &["code", "code-exploration", "code-insiders", "codium", "code-oss"];
 
 pub(crate) struct ServerOpt {
     pub(crate) malloc: Malloc,
@@ -70,21 +42,6 @@ pub(crate) enum Malloc {
     Jemalloc,
 }
 
-impl InstallCmd {
-    pub(crate) fn run(self) -> Result<()> {
-        if cfg!(target_os = "macos") {
-            fix_path_for_mac().context("Fix path for mac")?
-        }
-        if let Some(server) = self.server {
-            install_server(server).context("install server")?;
-        }
-        if let Some(client) = self.client {
-            install_client(client).context("install client")?;
-        }
-        Ok(())
-    }
-}
-
 fn fix_path_for_mac() -> Result<()> {
     let mut vscode_path: Vec<PathBuf> = {
         const COMMON_APP_PATH: &str =
@@ -121,21 +78,12 @@ fn fix_path_for_mac() -> Result<()> {
 fn install_client(client_opt: ClientOpt) -> Result<()> {
     let _dir = pushd("./editors/code");
 
-    let find_code = |f: fn(&str) -> bool| -> Result<&'static str> {
-        client_opt.as_cmds().iter().copied().find(|bin| f(bin)).ok_or_else(|| {
-            format_err!("Can't execute `code --version`. Perhaps it is not in $PATH?")
-        })
-    };
-
-    let installed_extensions = if cfg!(unix) {
+    // Package extension.
+    if cfg!(unix) {
         cmd!("npm --version").run().context("`npm` is required to build the VS Code plugin")?;
         cmd!("npm ci").run()?;
 
         cmd!("npm run package --scripts-prepend-node-path").run()?;
-
-        let code = find_code(|bin| cmd!("{bin} --version").read().is_ok())?;
-        cmd!("{code} --install-extension rust-analyzer.vsix --force").run()?;
-        cmd!("{code} --list-extensions").read()?
     } else {
         cmd!("cmd.exe /c npm --version")
             .run()
@@ -143,8 +91,36 @@ fn install_client(client_opt: ClientOpt) -> Result<()> {
         cmd!("cmd.exe /c npm ci").run()?;
 
         cmd!("cmd.exe /c npm run package").run()?;
+    };
+
+    // Find the appropriate VS Code binary.
+    let lifetime_extender;
+    let candidates: &[&str] = match client_opt.code_bin.as_deref() {
+        Some(it) => {
+            lifetime_extender = [it];
+            &lifetime_extender[..]
+        }
+        None => VS_CODES,
+    };
+    let code = candidates
+        .iter()
+        .copied()
+        .find(|&bin| {
+            if cfg!(unix) {
+                cmd!("{bin} --version").read().is_ok()
+            } else {
+                cmd!("cmd.exe /c {bin}.cmd --version").read().is_ok()
+            }
+        })
+        .ok_or_else(|| {
+            format_err!("Can't execute `{} --version`. Perhaps it is not in $PATH?", candidates[0])
+        })?;
 
-        let code = find_code(|bin| cmd!("cmd.exe /c {bin}.cmd --version").read().is_ok())?;
+    // Install & verify.
+    let installed_extensions = if cfg!(unix) {
+        cmd!("{code} --install-extension rust-analyzer.vsix --force").run()?;
+        cmd!("{code} --list-extensions").read()?
+    } else {
         cmd!("cmd.exe /c {code}.cmd --install-extension rust-analyzer.vsix --force").run()?;
         cmd!("cmd.exe /c {code}.cmd --list-extensions").read()?
     };
diff --git a/xtask/src/main.rs b/xtask/src/main.rs
index e419db7a7..3c4332f75 100644
--- a/xtask/src/main.rs
+++ b/xtask/src/main.rs
@@ -28,11 +28,7 @@ use std::{
 use walkdir::{DirEntry, WalkDir};
 use xshell::{cmd, cp, pushd, pushenv};
 
-use crate::{
-    codegen::Mode,
-    dist::DistCmd,
-    install::{InstallCmd, Malloc, ServerOpt},
-};
+use crate::{codegen::Mode, dist::DistCmd};
 
 fn main() -> Result<()> {
     let _d = pushd(project_root())?;
@@ -43,31 +39,7 @@ fn main() -> Result<()> {
             println!("{}", flags::Xtask::HELP);
             return Ok(());
         }
-        flags::XtaskCmd::Install(flags) => {
-            if flags.server && flags.client {
-                eprintln!(
-                    "error: The argument `--server` cannot be used with `--client`\n\n\
-                     For more information try --help"
-                );
-                return Ok(());
-            }
-
-            let malloc = if flags.mimalloc {
-                Malloc::Mimalloc
-            } else if flags.jemalloc {
-                Malloc::Jemalloc
-            } else {
-                Malloc::System
-            };
-
-            let client_bin = flags.code_bin.map(|it| it.parse()).transpose()?;
-
-            InstallCmd {
-                client: if flags.server { None } else { client_bin },
-                server: if flags.client { None } else { Some(ServerOpt { malloc }) },
-            }
-            .run()
-        }
+        flags::XtaskCmd::Install(cmd) => cmd.run(),
         flags::XtaskCmd::Codegen(cmd) => cmd.run(),
         flags::XtaskCmd::Lint(_) => run_clippy(),
         flags::XtaskCmd::FuzzTests(_) => run_fuzzer(),