rustc_utils 0.15.2-nightly-2026-05-01

//! Polonius integration to extract borrowck facts from rustc.

use std::sync::atomic::{AtomicBool, Ordering};

use rustc_borrowck::consumers::{BodyWithBorrowckFacts, ConsumerOptions};
use rustc_data_structures::fx::FxHashSet as HashSet;
use rustc_hir::def_id::LocalDefId;
use rustc_middle::{
  mir::{Body, StatementKind, TerminatorKind},
  ty::TyCtxt,
  util::Providers,
};

use crate::{BodyExt, block_timer, cache::Cache};

/// MIR pass to remove instructions not important for Flowistry.
///
/// This pass helps reduce the number of intermediates during dataflow analysis, which
/// reduces memory usage.
pub fn simplify_mir(body: &mut Body<'_>) {
  let return_blocks = body
    .all_returns()
    .filter_map(|loc| {
      let bb = &body.basic_blocks[loc.block];
      (bb.statements.len() == 0).then_some(loc.block)
    })
    .collect::<HashSet<_>>();

  for block in body.basic_blocks_mut() {
    block.statements.retain(|stmt| {
      !matches!(
        stmt.kind,
        StatementKind::StorageLive(..) | StatementKind::StorageDead(..)
      )
    });

    let terminator = block.terminator_mut();
    terminator.kind = match terminator.kind {
      TerminatorKind::FalseEdge { real_target, .. } => TerminatorKind::Goto {
        target: real_target,
      },
      // Ensures that control dependencies can determine the independence of different
      // return paths
      TerminatorKind::Goto { target } if return_blocks.contains(&target) => {
        TerminatorKind::Return
      }
      _ => continue,
    }
  }
}

static SIMPLIFY_MIR: AtomicBool = AtomicBool::new(false);

pub fn enable_mir_simplification() {
  SIMPLIFY_MIR.store(true, Ordering::SeqCst);
}

/// You must use this function in [`rustc_driver::Callbacks::config`] to call [`get_body_with_borrowck_facts`].
///
/// For why we need to do override mir_borrowck, see:
/// <https://github.com/rust-lang/rust/blob/485ced56b8753ec86936903f2a8c95e9be8996a1/src/test/run-make-fulldeps/obtain-borrowck/driver.rs>
pub fn override_queries(_session: &rustc_session::Session, local: &mut Providers) {
  local.queries.mir_borrowck = mir_borrowck;
}

thread_local! {
  static MIR_BODIES: Cache<LocalDefId, BodyWithBorrowckFacts<'static>> = Cache::default();
}

fn mir_borrowck(
  tcx: TyCtxt<'_>,
  def_id: LocalDefId,
) -> rustc_middle::queries::mir_borrowck::ProvidedValue<'_> {
  block_timer!(&format!(
    "get_bodies_with_borrowck_facts for {}",
    tcx.def_path_debug_str(def_id.to_def_id())
  ));

  let mut bodies_with_facts = rustc_borrowck::consumers::get_bodies_with_borrowck_facts(
    tcx,
    def_id,
    ConsumerOptions::PoloniusInputFacts,
  );

  for (def_id, mut body_with_facts) in bodies_with_facts.drain() {
    if SIMPLIFY_MIR.load(Ordering::SeqCst) {
      simplify_mir(&mut body_with_facts.body);
    }

    // SAFETY: The reader casts the 'static lifetime to 'tcx before using it.
    let body_with_facts: BodyWithBorrowckFacts<'static> =
      unsafe { std::mem::transmute(body_with_facts) };
    MIR_BODIES.with(|cache| {
      cache.get(&def_id, |_| body_with_facts);
    });
  }

  let mut providers = Providers::default();
  rustc_borrowck::provide(&mut providers.queries);
  let original_mir_borrowck = providers.queries.mir_borrowck;
  original_mir_borrowck(tcx, def_id)
}

/// Gets the MIR body and [Polonius](https://github.com/rust-lang/polonius)-generated
/// [borrowck facts](https://doc.rust-lang.org/nightly/nightly-rustc/rustc_borrowck/struct.BodyWithBorrowckFacts.html)
/// for a given [`LocalDefId`].
///
/// For this function to work, you MUST add [`override_queries`] to the
/// [`rustc_interface::Config`](https://doc.rust-lang.org/nightly/nightly-rustc/rustc_interface/interface/struct.Config.html)
/// inside of your [`rustc_driver::Callbacks`]. For example, see
/// [example.rs](https://github.com/willcrichton/flowistry/tree/master/crates/flowistry/examples/example.rs).
///
/// Note that as of May 2022, Polonius can be *very* slow for large functions.
/// It may take up to 30 seconds to analyze a single body with a large CFG.
pub fn get_body_with_borrowck_facts<'tcx>(
  tcx: TyCtxt<'tcx>,
  def_id: LocalDefId,
) -> &'tcx BodyWithBorrowckFacts<'tcx> {
  MIR_BODIES.with(|cache| {
    // Note: as of nightly-2025-08-20, get_bodies_with_borrowck_facts also returns the bodies for children
    // of a checked body. We have to handle the case where the parent of the current `def_id` was checked,
    // or else rustc panics about a stolen body. Hence, we check for whether the cache contains the key already.
    if !cache.contains_key(&def_id) {
      // Note: as of nightly-2026-05-01, mir_borrowck will panic if passed a child of an item (eg a closure),
      // so we have to make sure we call for the child's parent instead.
      let checkable_def_id = if tcx.is_typeck_child(def_id.to_def_id()) { tcx.local_parent(def_id) } else { def_id };
      let _ = tcx.mir_borrowck(checkable_def_id);
    }

    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?"));
    unsafe {
      std::mem::transmute::<
        &'_ BodyWithBorrowckFacts<'static>,
        &'tcx BodyWithBorrowckFacts<'tcx>,
      >(body)
    }
  })
}