flowistry_ide/
plugin.rs

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        // Note: this command must share certain parameters with rustc_plugin so Cargo will not recompute
101        // dependencies when actually running the driver, e.g. RUSTFLAGS.
102        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
298// Implement FlowistryAnalysis for all functions with a type signature that matches
299// FlowistryAnalysis::analyze
300impl<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}