1use std::{
2 borrow::Cow,
3 env,
4 path::PathBuf,
5 process::{Command, exit},
6 time::Instant,
7};
8
9use anyhow::Context;
10use base64::Engine;
11use clap::{Parser, Subcommand};
12use flowistry::extensions::{
13 ContextMode, EVAL_MODE, EvalMode, MutabilityMode, PointerMode,
14};
15use fluid_let::fluid_set;
16use log::{debug, info};
17use rustc_hir::BodyId;
18use rustc_interface::interface::Result as RustcResult;
19use rustc_middle::ty::TyCtxt;
20use rustc_plugin::{CrateFilter, RustcPlugin, RustcPluginArgs, Utf8Path};
21use rustc_span::ErrorGuaranteed;
22use rustc_utils::{
23 mir::borrowck_facts,
24 source_map::{
25 filename::Filename,
26 find_bodies::find_enclosing_bodies,
27 range::{CharPos, CharRange, FunctionIdentifier, ToSpan},
28 },
29 timer::elapsed,
30};
31use serde::{Deserialize, Serialize};
32
33#[derive(Parser, Serialize, Deserialize)]
34pub struct FlowistryPluginArgs {
35 #[clap(long)]
36 bench: Option<bool>,
37
38 #[clap(long)]
39 context_mode: Option<ContextMode>,
40 #[clap(long)]
41 mutability_mode: Option<MutabilityMode>,
42 #[clap(long)]
43 pointer_mode: Option<PointerMode>,
44
45 #[clap(subcommand)]
46 command: FlowistryCommand,
47}
48
49#[derive(Subcommand, Serialize, Deserialize)]
50enum FlowistryCommand {
51 Spans {
52 file: String,
53 },
54
55 Focus {
56 file: String,
57 pos_line: usize,
58 pos_column: usize,
59 },
60
61 Decompose {
62 file: String,
63 pos: usize,
64 },
65
66 Playground {
67 file: String,
68 start_line: usize,
69 start_column: usize,
70 end_line: usize,
71 end_column: usize,
72 },
73
74 Preload,
75
76 RustcVersion,
77}
78
79pub struct FlowistryPlugin;
80impl RustcPlugin for FlowistryPlugin {
81 type Args = FlowistryPluginArgs;
82
83 fn driver_name(&self) -> Cow<'static, str> {
84 "flowistry-driver".into()
85 }
86
87 fn version(&self) -> Cow<'static, str> {
88 env!("CARGO_PKG_VERSION").into()
89 }
90
91 fn args(&self, target_dir: &Utf8Path) -> RustcPluginArgs<FlowistryPluginArgs> {
92 let args = FlowistryPluginArgs::parse_from(env::args().skip(1));
93
94 let cargo_path = env::var("CARGO_PATH").unwrap_or_else(|_| "cargo".to_string());
95
96 use FlowistryCommand::*;
97 match &args.command {
98 Preload => {
99 let mut cmd = Command::new(cargo_path);
100 cmd
103 .args(["check", "--all", "--all-features", "--target-dir"])
104 .arg(target_dir);
105 let exit_status = cmd.status().expect("could not run cargo");
106 exit(exit_status.code().unwrap_or(-1));
107 }
108 RustcVersion => {
109 let version_str = rustc_interface::util::rustc_version_str().unwrap_or("unknown");
110 println!("{version_str}");
111 exit(0);
112 }
113 _ => {}
114 };
115
116 let file = match &args.command {
117 Spans { file, .. } => file,
118 Focus { file, .. } => file,
119 Decompose { file, .. } => file,
120 Playground { file, .. } => file,
121 _ => unreachable!(),
122 };
123
124 RustcPluginArgs {
125 filter: CrateFilter::CrateContainingFile(PathBuf::from(file)),
126 args,
127 }
128 }
129
130 fn run(
131 self,
132 compiler_args: Vec<String>,
133 plugin_args: FlowistryPluginArgs,
134 ) -> RustcResult<()> {
135 let eval_mode = EvalMode {
136 context_mode: plugin_args.context_mode.unwrap_or(ContextMode::SigOnly),
137 mutability_mode: plugin_args
138 .mutability_mode
139 .unwrap_or(MutabilityMode::DistinguishMut),
140 pointer_mode: plugin_args.pointer_mode.unwrap_or(PointerMode::Precise),
141 };
142 fluid_set!(EVAL_MODE, eval_mode);
143
144 use FlowistryCommand::*;
145 match plugin_args.command {
146 Spans { file, .. } => postprocess(crate::spans::spans(&compiler_args, file)),
147 Playground {
148 file,
149 start_line,
150 start_column,
151 end_line,
152 end_column,
153 ..
154 } => {
155 let compute_target = || CharRange {
156 start: CharPos {
157 line: start_line,
158 column: start_column,
159 },
160 end: CharPos {
161 line: end_line,
162 column: end_column,
163 },
164 filename: Filename::intern(&file),
165 };
166 postprocess(run(
167 crate::playground::playground,
168 compute_target,
169 &compiler_args,
170 ))
171 }
172 Focus {
173 file,
174 pos_line,
175 pos_column,
176 ..
177 } => {
178 let compute_target = || {
179 let cpos = CharPos {
180 line: pos_line,
181 column: pos_column,
182 };
183 let range = CharRange {
184 start: cpos,
185 end: cpos,
186 filename: Filename::intern(&file),
187 };
188 debug!("eyo WTF {range:?} {file}");
189 FunctionIdentifier::Range(range)
190 };
191 postprocess(run(crate::focus::focus, compute_target, &compiler_args))
192 }
193 Decompose {
194 file: _file,
195 pos: _pos,
196 ..
197 } => {
198 cfg_if::cfg_if! {
199 if #[cfg(feature = "decompose")] {
200 let indices = GraphemeIndices::from_path(&_file).unwrap();
201 let id =
202 FunctionIdentifier::Range(ByteRange::from_char_range(_pos, _pos, &_file, &indices));
203 postprocess(run(
204 crate::decompose::decompose,
205 id,
206 &compiler_args,
207 ))
208 } else {
209 panic!("Flowistry must be built with the decompose feature")
210 }
211 }
212 }
213 _ => unreachable!(),
214 }
215 }
216}
217
218fn postprocess<T: Serialize>(result: FlowistryResult<T>) -> RustcResult<()> {
219 let result = match result {
220 Ok(output) => Ok(output),
221 Err(e) => match e {
222 FlowistryError::BuildError => {
223 #[allow(deprecated)]
224 return Err(ErrorGuaranteed::unchecked_error_guaranteed());
225 }
226 e => Err(e),
227 },
228 };
229
230 let mut encoder =
231 flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::best());
232 serde_json::to_writer(&mut encoder, &result).unwrap();
233 let buffer = encoder.finish().unwrap();
234 print!(
235 "{}",
236 base64::engine::general_purpose::STANDARD.encode(buffer)
237 );
238
239 Ok(())
240}
241
242pub fn run_with_callbacks(
243 args: &[String],
244 callbacks: &mut (dyn rustc_driver::Callbacks + Send),
245) -> FlowistryResult<()> {
246 let mut args = args.to_vec();
247 args.extend(
248 "-Z identify-regions -Z mir-opt-level=0 -A warnings -Z maximal-hir-to-mir-coverage"
249 .split(' ')
250 .map(|s| s.to_owned()),
251 );
252
253 rustc_driver::catch_fatal_errors(move || rustc_driver::run_compiler(&args, callbacks))
254 .map_err(|_| FlowistryError::BuildError)
255}
256
257fn run<A: FlowistryAnalysis, T: ToSpan>(
258 analysis: A,
259 compute_target: impl FnOnce() -> T + Send,
260 args: &[String],
261) -> FlowistryResult<A::Output> {
262 let mut callbacks = FlowistryCallbacks {
263 analysis: Some(analysis),
264 compute_target: Some(compute_target),
265 output: None,
266 rustc_start: Instant::now(),
267 eval_mode: EVAL_MODE.copied(),
268 };
269
270 info!("Starting rustc analysis...");
271 debug!("Eval mode: {:?}", callbacks.eval_mode);
272
273 run_with_callbacks(args, &mut callbacks)?;
274
275 callbacks
276 .output
277 .unwrap()
278 .map_err(|e| FlowistryError::AnalysisError {
279 error: e.to_string(),
280 })
281}
282
283#[derive(Debug, Serialize)]
284#[serde(tag = "type")]
285pub enum FlowistryError {
286 BuildError,
287 AnalysisError { error: String },
288 FileNotFound,
289}
290
291pub type FlowistryResult<T> = Result<T, FlowistryError>;
292
293pub trait FlowistryAnalysis: Sized + Send + Sync {
294 type Output: Serialize + Send + Sync;
295 fn analyze(&mut self, tcx: TyCtxt, id: BodyId) -> anyhow::Result<Self::Output>;
296}
297
298impl<F, O> FlowistryAnalysis for F
301where
302 F: for<'tcx> Fn<(TyCtxt<'tcx>, BodyId), Output = anyhow::Result<O>> + Send + Sync,
303 O: Serialize + Send + Sync,
304{
305 type Output = O;
306 fn analyze(&mut self, tcx: TyCtxt, id: BodyId) -> anyhow::Result<Self::Output> {
307 (self)(tcx, id)
308 }
309}
310
311struct FlowistryCallbacks<A: FlowistryAnalysis, T: ToSpan, F: FnOnce() -> T> {
312 analysis: Option<A>,
313 compute_target: Option<F>,
314 output: Option<anyhow::Result<A::Output>>,
315 rustc_start: Instant,
316 eval_mode: Option<EvalMode>,
317}
318
319impl<A: FlowistryAnalysis, T: ToSpan, F: FnOnce() -> T> rustc_driver::Callbacks
320 for FlowistryCallbacks<A, T, F>
321{
322 fn config(&mut self, config: &mut rustc_interface::Config) {
323 borrowck_facts::enable_mir_simplification();
324 config.override_queries = Some(borrowck_facts::override_queries);
325 }
326
327 fn after_analysis<'tcx>(
328 &mut self,
329 _compiler: &rustc_interface::interface::Compiler,
330 tcx: TyCtxt<'tcx>,
331 ) -> rustc_driver::Compilation {
332 elapsed("rustc", self.rustc_start);
333 fluid_set!(EVAL_MODE, self.eval_mode.unwrap_or_default());
334
335 let mut analysis = self.analysis.take().unwrap();
336 self.output = Some((|| {
337 let target = (self.compute_target.take().unwrap())().to_span(tcx)?;
338 debug!("target span: {target:?}");
339 let mut bodies = find_enclosing_bodies(tcx, target);
340 let body = bodies.next().context("Selection did not map to a body")?;
341 analysis.analyze(tcx, body)
342 })());
343
344 rustc_driver::Compilation::Stop
345 }
346}