rustc_plugin/
driver.rs

1use std::{
2  env,
3  ops::Deref,
4  path::{Path, PathBuf},
5  process::{Command, exit},
6};
7
8use rustc_session::{EarlyDiagCtxt, config::ErrorOutputType};
9use rustc_tools_util::VersionInfo;
10
11use super::plugin::{PLUGIN_ARGS, RustcPlugin};
12use crate::cli::{RUN_ON_ALL_CRATES, SPECIFIC_CRATE, SPECIFIC_TARGET};
13
14/// If a command-line option matches `find_arg`, then apply the predicate `pred` on its value. If
15/// true, then return it. The parameter is assumed to be either `--arg=value` or `--arg value`.
16fn arg_value<'a, T: Deref<Target = str>>(
17  args: &'a [T],
18  find_arg: &str,
19  pred: impl Fn(&str) -> bool,
20) -> Option<&'a str> {
21  let mut args = args.iter().map(Deref::deref);
22  while let Some(arg) = args.next() {
23    let mut arg = arg.splitn(2, '=');
24    if arg.next() != Some(find_arg) {
25      continue;
26    }
27
28    match arg.next().or_else(|| args.next()) {
29      Some(v) if pred(v) => return Some(v),
30      _ => {}
31    }
32  }
33  None
34}
35
36fn toolchain_path(home: Option<String>, toolchain: Option<String>) -> Option<PathBuf> {
37  home.and_then(|home| {
38    toolchain.map(|toolchain| {
39      let mut path = PathBuf::from(home);
40      path.push("toolchains");
41      path.push(toolchain);
42      path
43    })
44  })
45}
46
47fn get_sysroot(orig_args: &[String]) -> (bool, String) {
48  // Get the sysroot, looking from most specific to this invocation to the least:
49  // - command line
50  // - runtime environment
51  //    - SYSROOT
52  //    - RUSTUP_HOME, MULTIRUST_HOME, RUSTUP_TOOLCHAIN, MULTIRUST_TOOLCHAIN
53  // - sysroot from rustc in the path
54  // - compile-time environment
55  //    - SYSROOT
56  //    - RUSTUP_HOME, MULTIRUST_HOME, RUSTUP_TOOLCHAIN, MULTIRUST_TOOLCHAIN
57  let sys_root_arg = arg_value(orig_args, "--sysroot", |_| true);
58  let have_sys_root_arg = sys_root_arg.is_some();
59  let sys_root = sys_root_arg
60    .map(PathBuf::from)
61    .or_else(|| std::env::var("MIRI_SYSROOT").ok().map(PathBuf::from))
62    .or_else(|| std::env::var("SYSROOT").ok().map(PathBuf::from))
63    .or_else(|| {
64      let home = std::env::var("RUSTUP_HOME")
65        .or_else(|_| std::env::var("MULTIRUST_HOME"))
66        .ok();
67      let toolchain = std::env::var("RUSTUP_TOOLCHAIN")
68        .or_else(|_| std::env::var("MULTIRUST_TOOLCHAIN"))
69        .ok();
70      toolchain_path(home, toolchain)
71    })
72    .or_else(|| {
73      Command::new("rustc")
74        .arg("--print")
75        .arg("sysroot")
76        .output()
77        .ok()
78        .and_then(|out| String::from_utf8(out.stdout).ok())
79        .map(|s| PathBuf::from(s.trim()))
80    })
81    .or_else(|| option_env!("SYSROOT").map(PathBuf::from))
82    .or_else(|| {
83      let home = option_env!("RUSTUP_HOME")
84        .or(option_env!("MULTIRUST_HOME"))
85        .map(ToString::to_string);
86      let toolchain = option_env!("RUSTUP_TOOLCHAIN")
87        .or(option_env!("MULTIRUST_TOOLCHAIN"))
88        .map(ToString::to_string);
89      toolchain_path(home, toolchain)
90    })
91    .map(|pb| pb.to_string_lossy().to_string())
92    .expect(
93      "need to specify SYSROOT env var during clippy compilation, or use rustup or multirust",
94    );
95  (have_sys_root_arg, sys_root)
96}
97
98struct DefaultCallbacks;
99impl rustc_driver::Callbacks for DefaultCallbacks {}
100
101/// The top-level function that should be called by your internal driver binary.
102pub fn driver_main<T: RustcPlugin>(plugin: T) {
103  let early_dcx = EarlyDiagCtxt::new(ErrorOutputType::default());
104  rustc_driver::init_rustc_env_logger(&early_dcx);
105
106  exit(rustc_driver::catch_with_exit_code(move || {
107    let mut orig_args: Vec<String> = env::args().collect();
108
109    let (have_sys_root_arg, sys_root) = get_sysroot(&orig_args);
110
111    if orig_args.iter().any(|a| a == "--version" || a == "-V") {
112      let version_info = rustc_tools_util::get_version_info!();
113      println!("{version_info}");
114      exit(0);
115    }
116
117    // Setting RUSTC_WRAPPER causes Cargo to pass 'rustc' as the first argument.
118    // We're invoking the compiler programmatically, so we ignore this
119    let wrapper_mode =
120      orig_args.get(1).map(Path::new).and_then(Path::file_stem) == Some("rustc".as_ref());
121
122    if wrapper_mode {
123      // we still want to be able to invoke it normally though
124      orig_args.remove(1);
125    }
126
127    // this conditional check for the --sysroot flag is there so users can call
128    // the driver directly without having to pass --sysroot or anything
129    let mut args: Vec<String> = orig_args.clone();
130    if !have_sys_root_arg {
131      args.extend(["--sysroot".into(), sys_root]);
132    };
133
134    // On a given invocation of rustc, we have to decide whether to act as rustc,
135    // or actually execute the plugin. There are two conditions for executing the plugin:
136    // 1. Either we're supposed to run on all crates, or CARGO_PRIMARY_PACKAGE is set.
137    // 2. --print is NOT passed, since Cargo does that to get info about rustc.
138    let primary_package = env::var("CARGO_PRIMARY_PACKAGE").is_ok();
139    let run_on_all_crates = env::var(RUN_ON_ALL_CRATES).is_ok();
140    let normal_rustc = arg_value(&args, "--print", |_| true).is_some();
141    let is_target_crate = is_target_crate(&args);
142    let run_plugin =
143      !normal_rustc && (run_on_all_crates || primary_package) && is_target_crate;
144
145    if run_plugin {
146      log::debug!("Running plugin...");
147      let plugin_args: T::Args =
148        serde_json::from_str(&env::var(PLUGIN_ARGS).unwrap()).unwrap();
149      plugin.run(args, plugin_args).unwrap();
150    } else {
151      log::debug!(
152        "Running normal Rust. Relevant variables:\
153normal_rustc={normal_rustc}, \
154run_on_all_crates={run_on_all_crates}, \
155primary_package={primary_package}, \
156is_target_crate={is_target_crate}"
157      );
158      rustc_driver::run_compiler(&args, &mut DefaultCallbacks);
159    }
160  }))
161}
162
163fn is_target_crate(args: &[String]) -> bool {
164  match (env::var(SPECIFIC_CRATE), env::var(SPECIFIC_TARGET)) {
165    (Ok(krate), Ok(target)) => {
166      arg_value(args, "--crate-name", |name| name == krate).is_some()
167        && (arg_value(args, "--crate-type", |_| true).is_none()  // integration test crate
168          || arg_value(args, "--crate-type", |name| name == target).is_some())
169    }
170    _ => true,
171  }
172}