Harden the pre-tyctxt query system against accidental recomputation

Author: oli-obk

compiler/rustc_driver/src/lib.rs

@@ -310,7 +310,7 @@ fn run_compiler(
             if let Some(ppm) = &sess.opts.pretty {
                 if ppm.needs_ast_map() {
                     let expanded_crate = queries.expansion()?.peek().0.clone();
-                    queries.global_ctxt()?.peek_mut().enter(|tcx| {
+                    queries.global_ctxt()?.enter(|tcx| {
                         pretty::print_after_hir_lowering(
                             tcx,
                             compiler.input(),
@@ -371,7 +371,7 @@ fn run_compiler(
                 return early_exit();
             }
 
-            queries.global_ctxt()?.peek_mut().enter(|tcx| {
+            queries.global_ctxt()?.enter(|tcx| {
                 let result = tcx.analysis(());
                 if sess.opts.unstable_opts.save_analysis {
                     let crate_name = tcx.crate_name(LOCAL_CRATE);

compiler/rustc_interface/src/queries.rs

@@ -19,43 +19,68 @@ use rustc_session::{output::find_crate_name, Session};
 use rustc_span::symbol::sym;
 use rustc_span::Symbol;
 use std::any::Any;
-use std::cell::{Ref, RefCell, RefMut};
+use std::cell::{Ref, RefCell};
 use std::rc::Rc;
 use std::sync::Arc;
 
 /// Represent the result of a query.
 ///
-/// This result can be stolen with the [`take`] method and generated with the [`compute`] method.
+/// This result can be stolen once with the [`take`] method and generated with the [`compute`] method.
 ///
 /// [`take`]: Self::take
 /// [`compute`]: Self::compute
 pub struct Query<T> {
-    result: RefCell<Option<Result<T>>>,
+    result: RefCell<Option<Option<Result<T>>>>,
 }
 
 impl<T> Query<T> {
     fn compute<F: FnOnce() -> Result<T>>(&self, f: F) -> Result<&Query<T>> {
-        self.result.borrow_mut().get_or_insert_with(f).as_ref().map(|_| self).map_err(|&err| err)
+        self.result
+            .borrow_mut()
+            .get_or_insert_with(|| Some(f()))
+            .as_ref()
+            .expect("query already taken")
+            .as_ref()
+            .map(|_| self)
+            .map_err(|&err| err)
     }
 
     /// Takes ownership of the query result. Further attempts to take or peek the query
     /// result will panic unless it is generated by calling the `compute` method.
     pub fn take(&self) -> T {
-        self.result.borrow_mut().take().expect("missing query result").unwrap()
+        self.result
+            .borrow_mut()
+            .as_mut()
+            .expect("query never computed")
+            .take()
+            .expect("query already taken")
+            .unwrap()
     }
 
     /// Borrows the query result using the RefCell. Panics if the result is stolen.
     pub fn peek(&self) -> Ref<'_, T> {
         Ref::map(self.result.borrow(), |r| {
-            r.as_ref().unwrap().as_ref().expect("missing query result")
+            r.as_ref()
+                .unwrap()
+                .as_ref()
+                .expect("query never computed")
+                .as_ref()
+                .expect("query already taken")
         })
     }
+}
 
-    /// Mutably borrows the query result using the RefCell. Panics if the result is stolen.
-    pub fn peek_mut(&self) -> RefMut<'_, T> {
-        RefMut::map(self.result.borrow_mut(), |r| {
-            r.as_mut().unwrap().as_mut().expect("missing query result")
-        })
+impl<'tcx> Query<QueryContext<'tcx>> {
+    pub fn enter<T>(&self, f: impl FnOnce(TyCtxt<'tcx>) -> T) -> T {
+        self.result
+            .borrow_mut()
+            .as_mut()
+            .unwrap()
+            .as_mut()
+            .expect("query never computed")
+            .as_mut()
+            .expect("query already taken")
+            .enter(f)
     }
 }
 
@@ -237,7 +262,7 @@ impl<'tcx> Queries<'tcx> {
 
     pub fn ongoing_codegen(&'tcx self) -> Result<&Query<Box<dyn Any>>> {
         self.ongoing_codegen.compute(|| {
-            self.global_ctxt()?.peek_mut().enter(|tcx| {
+            self.global_ctxt()?.enter(|tcx| {
                 tcx.analysis(()).ok();
 
                 // Don't do code generation if there were any errors
@@ -296,7 +321,6 @@ impl<'tcx> Queries<'tcx> {
         let dep_graph = self.dep_graph()?.peek().clone();
         let (crate_hash, prepare_outputs) = self
             .global_ctxt()?
-            .peek_mut()
             .enter(|tcx| (tcx.crate_hash(LOCAL_CRATE), tcx.output_filenames(()).clone()));
         let ongoing_codegen = self.ongoing_codegen()?.take();
 
@@ -381,7 +405,7 @@ impl Compiler {
 
         // NOTE: intentionally does not compute the global context if it hasn't been built yet,
         // since that likely means there was a parse error.
-        if let Some(Ok(gcx)) = &mut *queries.global_ctxt.result.borrow_mut() {
+        if let Some(Some(Ok(gcx))) = &mut *queries.global_ctxt.result.borrow_mut() {
             // We assume that no queries are run past here. If there are new queries
             // after this point, they'll show up as "<unknown>" in self-profiling data.
             {

src/librustdoc/lib.rs

@@ -817,7 +817,7 @@ fn main_args(at_args: &[String]) -> MainResult {
                 sess.fatal("Compilation failed, aborting rustdoc");
             }
 
-            let mut global_ctxt = abort_on_err(queries.global_ctxt(), sess).peek_mut();
+            let global_ctxt = abort_on_err(queries.global_ctxt(), sess);
 
             global_ctxt.enter(|tcx| {
                 let (krate, render_opts, mut cache) = sess.time("run_global_ctxt", || {

src/test/run-make-fulldeps/obtain-borrowck/driver.rs

@@ -62,7 +62,7 @@ impl rustc_driver::Callbacks for CompilerCalls {
         queries: &'tcx Queries<'tcx>,
     ) -> Compilation {
         compiler.session().abort_if_errors();
-        queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
+        queries.global_ctxt().unwrap().enter(|tcx| {
             // Collect definition ids of MIR bodies.
             let hir = tcx.hir();
             let mut bodies = Vec::new();

src/tools/miri/src/bin/miri.rs

@@ -61,7 +61,7 @@ impl rustc_driver::Callbacks for MiriCompilerCalls {
     ) -> Compilation {
         compiler.session().abort_if_errors();
 
-        queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
+        queries.global_ctxt().unwrap().enter(|tcx| {
             init_late_loggers(tcx);
             if !tcx.sess.crate_types().contains(&CrateType::Executable) {
                 tcx.sess.fatal("miri only makes sense on bin crates");