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, StatementKind, TerminatorKind},
10  ty::TyCtxt,
11  util::Providers,
12};
13
14use crate::{BodyExt, block_timer, cache::Cache};
15
16/// MIR pass to remove instructions not important for Flowistry.
17///
18/// This pass helps reduce the number of intermediates during dataflow analysis, which
19/// reduces memory usage.
20pub fn simplify_mir(body: &mut Body<'_>) {
21  let return_blocks = body
22    .all_returns()
23    .filter_map(|loc| {
24      let bb = &body.basic_blocks[loc.block];
25      (bb.statements.len() == 0).then_some(loc.block)
26    })
27    .collect::<HashSet<_>>();
28
29  for block in body.basic_blocks_mut() {
30    block.statements.retain(|stmt| {
31      !matches!(
32        stmt.kind,
33        StatementKind::StorageLive(..) | StatementKind::StorageDead(..)
34      )
35    });
36
37    let terminator = block.terminator_mut();
38    terminator.kind = match terminator.kind {
39      TerminatorKind::FalseEdge { real_target, .. } => TerminatorKind::Goto {
40        target: real_target,
41      },
42      // Ensures that control dependencies can determine the independence of different
43      // return paths
44      TerminatorKind::Goto { target } if return_blocks.contains(&target) => {
45        TerminatorKind::Return
46      }
47      _ => continue,
48    }
49  }
50}
51
52static SIMPLIFY_MIR: AtomicBool = AtomicBool::new(false);
53
54pub fn enable_mir_simplification() {
55  SIMPLIFY_MIR.store(true, Ordering::SeqCst);
56}
57
58/// You must use this function in [`rustc_driver::Callbacks::config`] to call [`get_body_with_borrowck_facts`].
59///
60/// For why we need to do override mir_borrowck, see:
61/// <https://github.com/rust-lang/rust/blob/485ced56b8753ec86936903f2a8c95e9be8996a1/src/test/run-make-fulldeps/obtain-borrowck/driver.rs>
62pub fn override_queries(_session: &rustc_session::Session, local: &mut Providers) {
63  local.queries.mir_borrowck = mir_borrowck;
64}
65
66thread_local! {
67  static MIR_BODIES: Cache<LocalDefId, BodyWithBorrowckFacts<'static>> = Cache::default();
68}
69
70fn mir_borrowck(
71  tcx: TyCtxt<'_>,
72  def_id: LocalDefId,
73) -> rustc_middle::queries::mir_borrowck::ProvidedValue<'_> {
74  block_timer!(&format!(
75    "get_bodies_with_borrowck_facts for {}",
76    tcx.def_path_debug_str(def_id.to_def_id())
77  ));
78
79  let mut bodies_with_facts = rustc_borrowck::consumers::get_bodies_with_borrowck_facts(
80    tcx,
81    def_id,
82    ConsumerOptions::PoloniusInputFacts,
83  );
84
85  for (def_id, mut body_with_facts) in bodies_with_facts.drain() {
86    if SIMPLIFY_MIR.load(Ordering::SeqCst) {
87      simplify_mir(&mut body_with_facts.body);
88    }
89
90    // SAFETY: The reader casts the 'static lifetime to 'tcx before using it.
91    let body_with_facts: BodyWithBorrowckFacts<'static> =
92      unsafe { std::mem::transmute(body_with_facts) };
93    MIR_BODIES.with(|cache| {
94      cache.get(&def_id, |_| body_with_facts);
95    });
96  }
97
98  let mut providers = Providers::default();
99  rustc_borrowck::provide(&mut providers.queries);
100  let original_mir_borrowck = providers.queries.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      // Note: as of nightly-2026-05-01, mir_borrowck will panic if passed a child of an item (eg a closure),
125      // so we have to make sure we call for the child's parent instead.
126      let checkable_def_id = if tcx.is_typeck_child(def_id.to_def_id()) { tcx.local_parent(def_id) } else { def_id };
127      let _ = tcx.mir_borrowck(checkable_def_id);
128    }
129
130    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?"));
131    unsafe {
132      std::mem::transmute::<
133        &'_ BodyWithBorrowckFacts<'static>,
134        &'tcx BodyWithBorrowckFacts<'tcx>,
135      >(body)
136    }
137  })
138}
rustc_utils/mir/borrowck_facts.rs

rustc_utils/mir/
borrowck_facts.rs
