rustc_utils/
test_utils.rs

1//! Running rustc and Flowistry in tests.
2
3use std::{
4  fmt::Debug,
5  fs,
6  hash::Hash,
7  io, panic,
8  path::Path,
9  process::Command,
10  sync::{Arc, LazyLock},
11};
12
13use anyhow::{Context, Result, anyhow, ensure};
14use log::debug;
15use rustc_abi::{FieldIdx, VariantIdx};
16use rustc_borrowck::consumers::BodyWithBorrowckFacts;
17use rustc_data_structures::fx::{FxHashMap as HashMap, FxHashSet as HashSet};
18use rustc_driver::run_compiler;
19use rustc_hir::BodyId;
20use rustc_middle::{
21  mir::{Body, HasLocalDecls, Local, Place, PlaceTy},
22  ty::TyCtxt,
23};
24use rustc_span::source_map::FileLoader;
25
26use crate::{
27  BodyExt, PlaceExt,
28  mir::borrowck_facts,
29  source_map::{
30    filename::{Filename, FilenameIndex},
31    find_bodies::{find_bodies, find_enclosing_bodies},
32    range::{BytePos, ByteRange, CharPos, CharRange, ToSpan},
33  },
34};
35
36pub struct StringLoader(pub String);
37impl FileLoader for StringLoader {
38  fn file_exists(&self, _: &Path) -> bool {
39    true
40  }
41
42  fn read_file(&self, _: &Path) -> io::Result<String> {
43    Ok(self.0.clone())
44  }
45
46  fn read_binary_file(&self, path: &Path) -> io::Result<Arc<[u8]>> {
47    Ok(fs::read(path)?.into())
48  }
49
50  fn current_directory(&self) -> io::Result<std::path::PathBuf> {
51    Ok(std::env::current_dir().unwrap())
52  }
53}
54
55static SYSROOT: LazyLock<String> = LazyLock::new(|| {
56  let rustc_output = Command::new("rustc")
57    .args(["--print", "sysroot"])
58    .output()
59    .unwrap()
60    .stdout;
61  String::from_utf8(rustc_output).unwrap().trim().to_owned()
62});
63
64pub const DUMMY_FILE_NAME: &str = "dummy.rs";
65
66thread_local! {
67  pub static DUMMY_FILE: FilenameIndex = Filename::intern(DUMMY_FILE_NAME);
68  pub static DUMMY_BYTE_RANGE: ByteRange = DUMMY_FILE.with(|filename| ByteRange {
69    start: BytePos(0),
70    end: BytePos(0),
71    filename: *filename,
72  });
73  pub static DUMMY_CHAR_RANGE: CharRange = DUMMY_FILE.with(|filename| CharRange {
74    start: CharPos { line: 0, column: 0 },
75    end: CharPos { line: 0, column: 0 },
76    filename: *filename,
77  });
78}
79
80/// Programmatically build a rustc compilation session
81pub struct CompileBuilder {
82  input: String,
83  arguments: Vec<String>,
84}
85
86impl CompileBuilder {
87  /// Initialize a compilation from this string of source code.
88  pub fn new(input: impl Into<String>) -> Self {
89    Self {
90      input: input.into(),
91      arguments: vec![],
92    }
93  }
94
95  /// Append additional rustc arguments
96  pub fn with_args(&mut self, args: impl IntoIterator<Item = String>) -> &mut Self {
97    self.arguments.extend(args);
98    self
99  }
100
101  /// Perform the compilation, providing access to it's intermediates state to
102  /// the provided closure
103  pub fn compile(&self, f: impl for<'tcx> FnOnce(CompileResult<'tcx>) + Send) {
104    let mut callbacks = TestCallbacks {
105      input: self.input.clone(),
106      callback: Some(move |tcx: TyCtxt<'_>| f(CompileResult { tcx })),
107    };
108    let args = [
109      "rustc",
110      DUMMY_FILE_NAME,
111      "--crate-type",
112      "lib",
113      "--edition=2024",
114      "-Zidentify-regions",
115      "-Zmir-opt-level=0",
116      "-Zmaximal-hir-to-mir-coverage",
117      "--allow",
118      "warnings",
119      "--sysroot",
120      &*SYSROOT,
121    ]
122    .into_iter()
123    .map(str::to_owned)
124    .chain(self.arguments.clone())
125    .collect::<Box<_>>();
126
127    rustc_driver::catch_fatal_errors(|| {
128      run_compiler(&args, &mut callbacks);
129    })
130    .unwrap();
131  }
132}
133
134/// Convenience alias for `CompileBuilder::new(input).compile(...)` if the
135/// callback is going to use [`CompileResult::as_body`].
136pub fn compile_body(
137  input: impl Into<String>,
138  callback: impl for<'tcx> FnOnce(TyCtxt<'tcx>, BodyId, &'tcx BodyWithBorrowckFacts<'tcx>)
139  + Send,
140) {
141  CompileBuilder::new(input).compile(|result| {
142    let (body_id, body_with_facts) = result.as_body();
143    callback(result.tcx, body_id, body_with_facts);
144  });
145}
146
147/// State during the rust compilation. Most of the time you only care about
148/// `self.tcx`, but this wrapper provides additional convenience methods for
149/// getting, e.g. the body of the configured entrypoint.
150pub struct CompileResult<'tcx> {
151  pub tcx: TyCtxt<'tcx>,
152}
153
154impl<'tcx> CompileResult<'tcx> {
155  /// Assume that we compiled only one function and return that function's id and body.
156  pub fn as_body(&self) -> (BodyId, &'tcx BodyWithBorrowckFacts<'tcx>) {
157    let tcx = self.tcx;
158    let (_, body_id) = find_bodies(tcx).remove(0);
159    let def_id = tcx.hir_body_owner_def_id(body_id);
160    let body_with_facts = borrowck_facts::get_body_with_borrowck_facts(tcx, def_id);
161    debug!("{}", body_with_facts.body.to_string(tcx).unwrap());
162    (body_id, body_with_facts)
163  }
164
165  /// Find a body in the target byte range.
166  pub fn as_body_with_range(
167    &self,
168    target: ByteRange,
169  ) -> (BodyId, &'tcx BodyWithBorrowckFacts<'tcx>) {
170    let tcx = self.tcx;
171    let body_id = find_enclosing_bodies(tcx, target.to_span(tcx).unwrap())
172      .next()
173      .unwrap();
174    let def_id = tcx.hir_body_owner_def_id(body_id);
175    let body_with_facts = borrowck_facts::get_body_with_borrowck_facts(tcx, def_id);
176    debug!("{}", body_with_facts.body.to_string(tcx).unwrap());
177
178    (body_id, body_with_facts)
179  }
180}
181
182struct TestCallbacks<Cb> {
183  input: String,
184  callback: Option<Cb>,
185}
186
187impl<Cb> rustc_driver::Callbacks for TestCallbacks<Cb>
188where
189  Cb: FnOnce(TyCtxt<'_>),
190{
191  fn config(&mut self, config: &mut rustc_interface::Config) {
192    config.override_queries = Some(borrowck_facts::override_queries);
193    config.file_loader = Some(Box::new(StringLoader(self.input.clone())));
194  }
195
196  fn after_expansion(
197    &mut self,
198    _compiler: &rustc_interface::interface::Compiler,
199    tcx: TyCtxt<'_>,
200  ) -> rustc_driver::Compilation {
201    let callback = self.callback.take().unwrap();
202    callback(tcx);
203    rustc_driver::Compilation::Stop
204  }
205}
206
207pub type RangeMap = HashMap<&'static str, Vec<ByteRange>>;
208
209pub fn parse_ranges(
210  src: impl AsRef<str>,
211  delimiters: impl AsRef<[(&'static str, &'static str)]>,
212) -> Result<(String, RangeMap)> {
213  let src = src.as_ref();
214  let delimiters = delimiters.as_ref();
215
216  let mut in_idx = 0;
217  let mut out_idx = 0;
218  let mut buf = Vec::new();
219  let bytes = src.bytes().collect::<Vec<_>>();
220  let mut stack = vec![];
221
222  let (opens, closes): (Vec<_>, Vec<_>) = delimiters.iter().copied().unzip();
223  let mut ranges: HashMap<_, Vec<_>> = HashMap::default();
224
225  macro_rules! check_token {
226    ($tokens:expr) => {
227      $tokens
228        .iter()
229        .find(|t| {
230          in_idx + t.len() <= bytes.len()
231            && t.as_bytes() == &bytes[in_idx .. in_idx + t.len()]
232        })
233        .map(|t| *t)
234    };
235  }
236
237  while in_idx < bytes.len() {
238    if let Some(open) = check_token!(opens) {
239      stack.push((out_idx, open));
240      in_idx += open.len();
241      continue;
242    }
243
244    if let Some(close) = check_token!(closes) {
245      let (start, delim) = stack
246        .pop()
247        .with_context(|| anyhow!("Missing open delimiter for \"{close}\""))?;
248      let range = DUMMY_FILE.with(|filename| ByteRange {
249        start: BytePos(start),
250        end: BytePos(out_idx),
251        filename: *filename,
252      });
253      ranges.entry(delim).or_default().push(range);
254      in_idx += close.len();
255      continue;
256    }
257
258    buf.push(bytes[in_idx]);
259    in_idx += 1;
260    out_idx += 1;
261  }
262
263  ensure!(stack.is_empty(), "Unclosed delimiters: {stack:?}");
264
265  let prog_clean = String::from_utf8(buf)?;
266  Ok((prog_clean, ranges))
267}
268
269pub fn color_ranges(prog: &str, all_ranges: &[(&str, &HashSet<ByteRange>)]) -> String {
270  let mut new_tokens = all_ranges
271    .iter()
272    .flat_map(|(_, ranges)| {
273      ranges.iter().flat_map(|range| {
274        let contained = all_ranges.iter().any(|(_, ranges)| {
275          ranges.iter().any(|other| {
276            range != other && other.start.0 <= range.end.0 && range.end.0 < other.end.0
277          })
278        });
279        let end_marker = if contained { "]" } else { "\x1B[0m]" };
280        [("[\x1B[31m", range.start), (end_marker, range.end)]
281      })
282    })
283    .collect::<Vec<_>>();
284  new_tokens.sort_by_key(|(_, i)| -(isize::try_from(i.0).unwrap()));
285
286  let mut output = prog.to_owned();
287  for (s, i) in new_tokens {
288    output.insert_str(i.0, s);
289  }
290
291  output
292}
293
294pub fn fmt_ranges(prog: &str, s: &HashSet<ByteRange>) -> String {
295  textwrap::indent(&color_ranges(prog, &[("", s)]), "  ")
296}
297
298pub fn compare_ranges(
299  expected: &HashSet<ByteRange>,
300  actual: &HashSet<ByteRange>,
301  prog: &str,
302) {
303  let missing = expected - actual;
304  let extra = actual - expected;
305
306  let check = |s: HashSet<ByteRange>, message: &str| {
307    if s.len() > 0 {
308      println!("Expected ranges:\n{}", fmt_ranges(prog, expected));
309      println!("Actual ranges:\n{}", fmt_ranges(prog, actual));
310      panic!("{message} ranges:\n{}", fmt_ranges(prog, &s));
311    }
312  };
313
314  check(missing, "Analysis did NOT have EXPECTED");
315  check(extra, "Actual DID have UNEXPECTED");
316}
317
318pub struct Placer<'a, 'tcx> {
319  tcx: TyCtxt<'tcx>,
320  body: &'a Body<'tcx>,
321  local_map: HashMap<String, Local>,
322}
323
324impl<'a, 'tcx> Placer<'a, 'tcx> {
325  pub fn new(tcx: TyCtxt<'tcx>, body: &'a Body<'tcx>) -> Self {
326    let local_map = body.debug_info_name_map();
327    Placer {
328      tcx,
329      body,
330      local_map,
331    }
332  }
333
334  pub fn local(&self, name: &str) -> PlaceBuilder<'a, 'tcx> {
335    PlaceBuilder {
336      place: Place::from_local(self.local_map[name], self.tcx),
337      body: self.body,
338      tcx: self.tcx,
339    }
340  }
341}
342
343#[derive(Copy, Clone)]
344pub struct PlaceBuilder<'a, 'tcx> {
345  tcx: TyCtxt<'tcx>,
346  body: &'a Body<'tcx>,
347  place: Place<'tcx>,
348}
349
350impl<'tcx> PlaceBuilder<'_, 'tcx> {
351  pub fn field(mut self, i: usize) -> Self {
352    let f = FieldIdx::from_usize(i);
353    let place_ty = self.place.ty(self.body.local_decls(), self.tcx);
354    let ty = PlaceTy::field_ty(self.tcx, place_ty.ty, place_ty.variant_index, f);
355    self.place = self.tcx.mk_place_field(self.place, f, ty);
356    self
357  }
358
359  pub fn deref(mut self) -> Self {
360    self.place = self.tcx.mk_place_deref(self.place);
361    self
362  }
363
364  pub fn downcast(mut self, i: usize) -> Self {
365    let ty = self.place.ty(self.body.local_decls(), self.tcx).ty;
366    let adt_def = ty.ty_adt_def().unwrap();
367    let v = VariantIdx::from_usize(i);
368    self.place = self.tcx.mk_place_downcast(self.place, adt_def, v);
369    self
370  }
371
372  pub fn index(mut self, i: usize) -> Self {
373    self.place = self.tcx.mk_place_index(self.place, Local::from_usize(i));
374    self
375  }
376
377  pub fn mk(self) -> Place<'tcx> {
378    self.place
379  }
380}
381
382pub fn compare_sets<T: PartialEq + Eq + Clone + Hash + Debug>(
383  actual: impl IntoIterator<Item = T>,
384  expected: impl IntoIterator<Item = T>,
385) {
386  let expected = expected.into_iter().collect::<HashSet<_>>();
387  let actual = actual.into_iter().collect::<HashSet<_>>();
388
389  let missing = &expected - &actual;
390  let extra = &actual - &expected;
391
392  let check = |s: HashSet<T>, message: &str| {
393    if s.len() > 0 {
394      println!(
395        "Expected:\n{}",
396        textwrap::indent(&format!("{expected:#?}"), "  ")
397      );
398      println!(
399        "Actual:\n{}",
400        textwrap::indent(&format!("{actual:#?}"), "  ")
401      );
402      panic!(
403        "{message} ranges:\n{}",
404        textwrap::indent(&format!("{s:#?}"), "  ")
405      );
406    }
407  };
408
409  check(missing, "Result did NOT have EXPECTED");
410  check(extra, "Result DID have UNEXPECTED");
411}
412
413#[cfg(test)]
414mod test {
415  use super::*;
416
417  #[test]
418  fn test_parse_ranges() {
419    DUMMY_FILE.with(|filename| {
420      let s = "`[`[f]`oo]`";
421      let (clean, ranges) = parse_ranges(s, vec![("`[", "]`")]).unwrap();
422      assert_eq!(clean, "foo");
423      assert_eq!(ranges["`["], vec![
424        ByteRange {
425          start: BytePos(0),
426          end: BytePos(1),
427          filename: *filename,
428        },
429        ByteRange {
430          start: BytePos(0),
431          end: BytePos(3),
432          filename: *filename,
433        },
434      ]);
435    });
436  }
437}
rustc_utils/test_utils.rs

rustc_utils/
test_utils.rs
