Auto merge of #86977 - vakaras:body_with_borrowck_facts, r=nikomatsakis

Enable compiler consumers to obtain mir::Body with Polonius facts.

This PR adds a function (``get_body_with_borrowck_facts``) that can be used by compiler consumers to obtain ``mir::Body`` with accompanying borrow checker information.

The most important borrow checker information that [our verifier called Prusti](https://github.com/viperproject/prusti-dev) needs is lifetime constraints. I have not found a reasonable way to compute the lifetime constraints on the Prusti side. In the compiler, the constraints are computed during the borrow checking phase and then dropped. This PR adds an additional parameter to the `do_mir_borrowck` function that tells it to return the computed information instead of dropping it.

The additionally returned information by `do_mir_borrowck` contains a ``mir::Body`` with non-erased lifetime regions and Polonius facts. I have decided to reuse the Polonius facts because this way I needed fewer changes to the compiler and Polonius facts contains other useful information that we otherwise would need to recompute.

Just FYI: up to now, Prusti was obtaining this information by [parsing the compiler logs](https://github.com/viperproject/prusti-dev/blob/b58ced8dfd14ef30582b503d517167ccd771eaff/prusti-interface/src/environment/borrowck/regions.rs#L25-L39). This is not only a hacky approach, but we also reached its limits.

r? `@nikomatsakis`

4 files changed, 237 insertions, 0 deletions

diff --git a/src/test/run-make-fulldeps/obtain-borrowck/Makefile b/src/test/run-make-fulldeps/obtain-borrowck/Makefile
new file mode 100644
index 00000000000..223993125c8
--- /dev/null
+++ b/src/test/run-make-fulldeps/obtain-borrowck/Makefile
@@ -0,0 +1,26 @@
+include ../tools.mk
+
+# This example shows how to implement a rustc driver that retrieves MIR bodies
+# together with the borrow checker information.
+
+# How to run this
+# $ ./x.py test src/test/run-make-fulldeps/obtain-borrowck
+
+DRIVER_BINARY := "$(TMPDIR)"/driver
+SYSROOT := $(shell $(RUSTC) --print sysroot)
+
+ifdef IS_WINDOWS
+LIBSTD := -L "$(SYSROOT)\\lib\\rustlib\\$(TARGET)\\lib"
+else
+LIBSTD :=
+endif
+
+all:
+	$(RUSTC) driver.rs -o "$(DRIVER_BINARY)"
+	$(TARGET_RPATH_ENV) "$(DRIVER_BINARY)" --sysroot $(SYSROOT) $(LIBSTD) test.rs -o "$(TMPDIR)/driver_test" > "$(TMPDIR)"/output.stdout
+
+ifdef RUSTC_BLESS_TEST
+	cp "$(TMPDIR)"/output.stdout output.stdout
+else
+	$(DIFF) output.stdout "$(TMPDIR)"/output.stdout
+endif
diff --git a/src/test/run-make-fulldeps/obtain-borrowck/driver.rs b/src/test/run-make-fulldeps/obtain-borrowck/driver.rs
new file mode 100644
index 00000000000..308df0b030c
--- /dev/null
+++ b/src/test/run-make-fulldeps/obtain-borrowck/driver.rs
@@ -0,0 +1,171 @@
+#![feature(rustc_private)]
+
+//! This program implements a rustc driver that retrieves MIR bodies with
+//! borrowck information. This cannot be done in a straightforward way because
+//! `get_body_with_borrowck_facts`–the function for retrieving a MIR body with
+//! borrowck facts–can panic if the body is stolen before it is invoked.
+//! Therefore, the driver overrides `mir_borrowck` query (this is done in the
+//! `config` callback), which retrieves the body that is about to be borrow
+//! checked and stores it in a thread local `MIR_BODIES`. Then, `after_analysis`
+//! callback triggers borrow checking of all MIR bodies by retrieving
+//! `optimized_mir` and pulls out the MIR bodies with the borrowck information
+//! from the thread local storage.
+
+extern crate rustc_driver;
+extern crate rustc_hir;
+extern crate rustc_interface;
+extern crate rustc_middle;
+extern crate rustc_mir;
+extern crate rustc_session;
+
+use rustc_driver::Compilation;
+use rustc_hir::def_id::LocalDefId;
+use rustc_hir::itemlikevisit::ItemLikeVisitor;
+use rustc_interface::interface::Compiler;
+use rustc_interface::{Config, Queries};
+use rustc_middle::ty::query::query_values::mir_borrowck;
+use rustc_middle::ty::query::Providers;
+use rustc_middle::ty::{self, TyCtxt};
+use rustc_mir::consumers::BodyWithBorrowckFacts;
+use rustc_session::Session;
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::thread_local;
+
+fn main() {
+    let exit_code = rustc_driver::catch_with_exit_code(move || {
+        let mut rustc_args: Vec<_> = std::env::args().collect();
+        // We must pass -Zpolonius so that the borrowck information is computed.
+        rustc_args.push("-Zpolonius".to_owned());
+        let mut callbacks = CompilerCalls::default();
+        // Call the Rust compiler with our callbacks.
+        rustc_driver::RunCompiler::new(&rustc_args, &mut callbacks).run()
+    });
+    std::process::exit(exit_code);
+}
+
+#[derive(Default)]
+pub struct CompilerCalls;
+
+impl rustc_driver::Callbacks for CompilerCalls {
+
+    // In this callback we override the mir_borrowck query.
+    fn config(&mut self, config: &mut Config) {
+        assert!(config.override_queries.is_none());
+        config.override_queries = Some(override_queries);
+    }
+
+    // In this callback we trigger borrow checking of all functions and obtain
+    // the result.
+    fn after_analysis<'tcx>(
+        &mut self,
+        compiler: &Compiler,
+        queries: &'tcx Queries<'tcx>,
+    ) -> Compilation {
+        compiler.session().abort_if_errors();
+        queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
+
+            // Collect definition ids of MIR bodies.
+            let hir = tcx.hir();
+            let krate = hir.krate();
+            let mut visitor = HirVisitor { bodies: Vec::new() };
+            krate.visit_all_item_likes(&mut visitor);
+
+            // Trigger borrow checking of all bodies.
+            for def_id in visitor.bodies {
+                let _ = tcx.optimized_mir(def_id);
+            }
+
+            // See what bodies were borrow checked.
+            let mut bodies = get_bodies(tcx);
+            bodies.sort_by(|(def_id1, _), (def_id2, _)| def_id1.cmp(def_id2));
+            println!("Bodies retrieved for:");
+            for (def_id, body) in bodies {
+                println!("{}", def_id);
+                assert!(body.input_facts.cfg_edge.len() > 0);
+            }
+        });
+
+        Compilation::Continue
+    }
+}
+
+fn override_queries(_session: &Session, local: &mut Providers, external: &mut Providers) {
+    local.mir_borrowck = mir_borrowck;
+    external.mir_borrowck = mir_borrowck;
+}
+
+// Since mir_borrowck does not have access to any other state, we need to use a
+// thread-local for storing the obtained MIR bodies.
+//
+// Note: We are using 'static lifetime here, which is in general unsound.
+// Unfortunately, that is the only lifetime allowed here. Our use is safe
+// because we cast it back to `'tcx` before using.
+thread_local! {
+    pub static MIR_BODIES:
+        RefCell<HashMap<LocalDefId, BodyWithBorrowckFacts<'static>>> =
+        RefCell::new(HashMap::new());
+}
+
+fn mir_borrowck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> mir_borrowck<'tcx> {
+    let body_with_facts = rustc_mir::consumers::get_body_with_borrowck_facts(
+        tcx,
+        ty::WithOptConstParam::unknown(def_id),
+    );
+    // 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(|state| {
+        let mut map = state.borrow_mut();
+        assert!(map.insert(def_id, body_with_facts).is_none());
+    });
+    let mut providers = Providers::default();
+    rustc_mir::provide(&mut providers);
+    let original_mir_borrowck = providers.mir_borrowck;
+    original_mir_borrowck(tcx, def_id)
+}
+
+/// Visitor that collects all body definition ids mentioned in the program.
+struct HirVisitor {
+    bodies: Vec<LocalDefId>,
+}
+
+impl<'tcx> ItemLikeVisitor<'tcx> for HirVisitor {
+    fn visit_item(&mut self, item: &rustc_hir::Item) {
+        if let rustc_hir::ItemKind::Fn(..) = item.kind {
+            self.bodies.push(item.def_id);
+        }
+    }
+
+    fn visit_trait_item(&mut self, trait_item: &rustc_hir::TraitItem) {
+        if let rustc_hir::TraitItemKind::Fn(_, trait_fn) = &trait_item.kind {
+            if let rustc_hir::TraitFn::Provided(_) = trait_fn {
+                self.bodies.push(trait_item.def_id);
+            }
+        }
+    }
+
+    fn visit_impl_item(&mut self, impl_item: &rustc_hir::ImplItem) {
+        if let rustc_hir::ImplItemKind::Fn(..) = impl_item.kind {
+            self.bodies.push(impl_item.def_id);
+        }
+    }
+
+    fn visit_foreign_item(&mut self, _foreign_item: &rustc_hir::ForeignItem) {}
+}
+
+/// Pull MIR bodies stored in the thread-local.
+fn get_bodies<'tcx>(tcx: TyCtxt<'tcx>) -> Vec<(String, BodyWithBorrowckFacts<'tcx>)> {
+    MIR_BODIES.with(|state| {
+        let mut map = state.borrow_mut();
+        map.drain()
+            .map(|(def_id, body)| {
+                let def_path = tcx.def_path(def_id.to_def_id());
+                // SAFETY: For soundness we need to ensure that the bodies have
+                // the same lifetime (`'tcx`), which they had before they were
+                // stored in the thread local.
+                (def_path.to_string_no_crate_verbose(), body)
+            })
+            .collect()
+    })
+}
diff --git a/src/test/run-make-fulldeps/obtain-borrowck/output.stdout b/src/test/run-make-fulldeps/obtain-borrowck/output.stdout
new file mode 100644
index 00000000000..e011622e6b2
--- /dev/null
+++ b/src/test/run-make-fulldeps/obtain-borrowck/output.stdout
@@ -0,0 +1,8 @@
+Bodies retrieved for:
+::X::provided
+::foo
+::main
+::main::{constant#0}
+::{impl#0}::new
+::{impl#1}::provided
+::{impl#1}::required
diff --git a/src/test/run-make-fulldeps/obtain-borrowck/test.rs b/src/test/run-make-fulldeps/obtain-borrowck/test.rs
new file mode 100644
index 00000000000..f7b4b41feaf
--- /dev/null
+++ b/src/test/run-make-fulldeps/obtain-borrowck/test.rs
@@ -0,0 +1,32 @@
+trait X {
+    fn provided(&self) -> usize {
+        5
+    }
+    fn required(&self) -> u32;
+}
+
+struct Bar;
+
+impl Bar {
+    fn new() -> Self {
+        Self
+    }
+}
+
+impl X for Bar {
+    fn provided(&self) -> usize {
+        1
+    }
+    fn required(&self) -> u32 {
+        7
+    }
+}
+
+const fn foo() -> usize {
+    1
+}
+
+fn main() {
+    let bar: [Bar; foo()] = [Bar::new()];
+    assert_eq!(bar[0].provided(), foo());
+}