rustc_utils/mir/
borrowck_facts.rs

1//! Polonius integration to extract borrowck facts from rustc.
2
3use std::sync::atomic::{AtomicBool, Ordering};
4
5use rustc_borrowck::consumers::{BodyWithBorrowckFacts, ConsumerOptions};
6use rustc_data_structures::fx::FxHashSet as HashSet;
7use rustc_hir::def_id::LocalDefId;
8use rustc_middle::{
9  mir::{Body, ConcreteOpaqueTypes, StatementKind, TerminatorKind},
10  ty::TyCtxt,
11  util::Providers,
12};
13use rustc_span::ErrorGuaranteed;
14
15use crate::{BodyExt, block_timer, cache::Cache};
16
17/// MIR pass to remove instructions not important for Flowistry.
18///
19/// This pass helps reduce the number of intermediates during dataflow analysis, which
20/// reduces memory usage.
21pub fn simplify_mir(body: &mut Body<'_>) {
22  let return_blocks = body
23    .all_returns()
24    .filter_map(|loc| {
25      let bb = &body.basic_blocks[loc.block];
26      (bb.statements.len() == 0).then_some(loc.block)
27    })
28    .collect::<HashSet<_>>();
29
30  for block in body.basic_blocks_mut() {
31    block.statements.retain(|stmt| {
32      !matches!(
33        stmt.kind,
34        StatementKind::StorageLive(..) | StatementKind::StorageDead(..)
35      )
36    });
37
38    let terminator = block.terminator_mut();
39    terminator.kind = match terminator.kind {
40      TerminatorKind::FalseEdge { real_target, .. } => TerminatorKind::Goto {
41        target: real_target,
42      },
43      // Ensures that control dependencies can determine the independence of different
44      // return paths
45      TerminatorKind::Goto { target } if return_blocks.contains(&target) => {
46        TerminatorKind::Return
47      }
48      _ => continue,
49    }
50  }
51}
52
53static SIMPLIFY_MIR: AtomicBool = AtomicBool::new(false);
54
55pub fn enable_mir_simplification() {
56  SIMPLIFY_MIR.store(true, Ordering::SeqCst);
57}
58
59/// You must use this function in [`rustc_driver::Callbacks::config`] to call [`get_body_with_borrowck_facts`].
60///
61/// For why we need to do override mir_borrowck, see:
62/// <https://github.com/rust-lang/rust/blob/485ced56b8753ec86936903f2a8c95e9be8996a1/src/test/run-make-fulldeps/obtain-borrowck/driver.rs>
63pub fn override_queries(_session: &rustc_session::Session, local: &mut Providers) {
64  local.mir_borrowck = mir_borrowck;
65}
66
67thread_local! {
68  static MIR_BODIES: Cache<LocalDefId, BodyWithBorrowckFacts<'static>> = Cache::default();
69}
70
71fn mir_borrowck(
72  tcx: TyCtxt<'_>,
73  def_id: LocalDefId,
74) -> Result<&ConcreteOpaqueTypes<'_>, ErrorGuaranteed> {
75  block_timer!(&format!(
76    "get_bodies_with_borrowck_facts for {}",
77    tcx.def_path_debug_str(def_id.to_def_id())
78  ));
79
80  let mut body_with_facts = rustc_borrowck::consumers::get_bodies_with_borrowck_facts(
81    tcx,
82    def_id,
83    ConsumerOptions::PoloniusInputFacts,
84  );
85
86  for (def_id, mut body_with_facts) in body_with_facts.drain() {
87    if SIMPLIFY_MIR.load(Ordering::SeqCst) {
88      simplify_mir(&mut body_with_facts.body);
89    }
90
91    // SAFETY: The reader casts the 'static lifetime to 'tcx before using it.
92    let body_with_facts: BodyWithBorrowckFacts<'static> =
93      unsafe { std::mem::transmute(body_with_facts) };
94    MIR_BODIES.with(|cache| {
95      cache.get(&def_id, |_| body_with_facts);
96    });
97  }
98  let mut providers = Providers::default();
99  rustc_borrowck::provide(&mut providers);
100  let original_mir_borrowck = providers.mir_borrowck;
101  original_mir_borrowck(tcx, def_id)
102}
103
104/// Gets the MIR body and [Polonius](https://github.com/rust-lang/polonius)-generated
105/// [borrowck facts](https://doc.rust-lang.org/nightly/nightly-rustc/rustc_borrowck/struct.BodyWithBorrowckFacts.html)
106/// for a given [`LocalDefId`].
107///
108/// For this function to work, you MUST add [`override_queries`] to the
109/// [`rustc_interface::Config`](https://doc.rust-lang.org/nightly/nightly-rustc/rustc_interface/interface/struct.Config.html)
110/// inside of your [`rustc_driver::Callbacks`]. For example, see
111/// [example.rs](https://github.com/willcrichton/flowistry/tree/master/crates/flowistry/examples/example.rs).
112///
113/// Note that as of May 2022, Polonius can be *very* slow for large functions.
114/// It may take up to 30 seconds to analyze a single body with a large CFG.
115pub fn get_body_with_borrowck_facts<'tcx>(
116  tcx: TyCtxt<'tcx>,
117  def_id: LocalDefId,
118) -> &'tcx BodyWithBorrowckFacts<'tcx> {
119  MIR_BODIES.with(|cache| {
120    // Note: as of nightly-2025-08-20, get_bodies_with_borrowck_facts also returns the bodies for children
121    // of a checked body. We have to handle the case where the parent of the current `def_id` was checked,
122    // or else rustc panics about a stolen body. Hence, we check for whether the cache contains the key already.
123    if !cache.contains_key(&def_id) {
124      let _ = tcx.mir_borrowck(def_id);
125    }
126
127    let body = cache.get(&def_id, |_| panic!("mir_borrowck override should have stored body for item: {def_id:?}. Are you sure you registered borrowck_facts::override_queries?"));
128    unsafe {
129      std::mem::transmute::<
130        &'_ BodyWithBorrowckFacts<'static>,
131        &'tcx BodyWithBorrowckFacts<'tcx>,
132      >(body)
133    }
134  })
135}
rustc_utils/mir/borrowck_facts.rs

rustc_utils/mir/
borrowck_facts.rs
