Rollup merge of rust-lang#152417 - Zalathar:arena-alloc, r=nnethercote

Move the needs-drop check for `arena_cache` queries out of macro code

This is slightly simpler than before, because now the macro only needs to call a single function, and can just unconditionally supply `tcx` and a typed arena.

There should be no actual change to compiler behaviour.

Author: JonathanBrouwer

compiler/rustc_middle/src/queries.rs

@@ -63,7 +63,6 @@
 #![allow(unused_parens)]
 
 use std::ffi::OsStr;
-use std::mem;
 use std::path::PathBuf;
 use std::sync::Arc;
 

compiler/rustc_middle/src/query/arena_cached.rs

@@ -1,3 +1,9 @@
+use std::mem;
+
+use rustc_arena::TypedArena;
+
+use crate::ty::TyCtxt;
+
 /// Helper trait that allows `arena_cache` queries to return `Option<&T>`
 /// instead of `&Option<T>`, and avoid allocating `None` in the arena.
 ///
@@ -11,10 +17,11 @@ pub trait ArenaCached<'tcx>: Sized {
     /// Type that is stored in the arena.
     type Allocated: 'tcx;
 
-    /// Takes a provided value, and allocates it in the arena (if appropriate)
-    /// with the help of the given `arena_alloc` closure.
+    /// Takes a provided value, and allocates it in an appropriate arena,
+    /// unless the particular value doesn't need allocation (e.g. `None`).
     fn alloc_in_arena(
-        arena_alloc: impl Fn(Self::Allocated) -> &'tcx Self::Allocated,
+        tcx: TyCtxt<'tcx>,
+        typed_arena: &'tcx TypedArena<Self::Allocated>,
         value: Self::Provided,
     ) -> Self;
 }
@@ -23,12 +30,9 @@ impl<'tcx, T> ArenaCached<'tcx> for &'tcx T {
     type Provided = T;
     type Allocated = T;
 
-    fn alloc_in_arena(
-        arena_alloc: impl Fn(Self::Allocated) -> &'tcx Self::Allocated,
-        value: Self::Provided,
-    ) -> Self {
+    fn alloc_in_arena(tcx: TyCtxt<'tcx>, typed_arena: &'tcx TypedArena<T>, value: T) -> Self {
         // Just allocate in the arena normally.
-        arena_alloc(value)
+        do_alloc(tcx, typed_arena, value)
     }
 }
 
@@ -38,10 +42,17 @@ impl<'tcx, T> ArenaCached<'tcx> for Option<&'tcx T> {
     type Allocated = T;
 
     fn alloc_in_arena(
-        arena_alloc: impl Fn(Self::Allocated) -> &'tcx Self::Allocated,
-        value: Self::Provided,
+        tcx: TyCtxt<'tcx>,
+        typed_arena: &'tcx TypedArena<T>,
+        value: Option<T>,
     ) -> Self {
         // Don't store None in the arena, and wrap the allocated reference in Some.
-        value.map(arena_alloc)
+        try { do_alloc(tcx, typed_arena, value?) }
     }
 }
+
+/// Allocates a value in either its dedicated arena, or in the common dropless
+/// arena, depending on whether it needs to be dropped.
+fn do_alloc<'tcx, T>(tcx: TyCtxt<'tcx>, typed_arena: &'tcx TypedArena<T>, value: T) -> &'tcx T {
+    if mem::needs_drop::<T>() { typed_arena.alloc(value) } else { tcx.arena.dropless.alloc(value) }
+}

compiler/rustc_middle/src/query/plumbing.rs

@@ -295,36 +295,28 @@ macro_rules! define_callbacks {
                 ($V)
             );
 
-            /// This function takes `ProvidedValue` and converts it to an erased `Value` by
-            /// allocating it on an arena if the query has the `arena_cache` modifier. The
-            /// value is then erased and returned. This will happen when computing the query
-            /// using a provider or decoding a stored result.
+            /// This helper function takes a value returned by the query provider
+            /// (or loaded from disk, or supplied by query feeding), allocates
+            /// it in an arena if requested by the `arena_cache` modifier, and
+            /// then returns an erased copy of it.
             #[inline(always)]
             pub fn provided_to_erased<'tcx>(
-                _tcx: TyCtxt<'tcx>,
+                tcx: TyCtxt<'tcx>,
                 provided_value: ProvidedValue<'tcx>,
             ) -> Erased<Value<'tcx>> {
-                // Store the provided value in an arena and get a reference
-                // to it, for queries with `arena_cache`.
-                let value: Value<'tcx> = query_if_arena!([$($modifiers)*]
-                    {
-                        use $crate::query::arena_cached::ArenaCached;
-
-                        if mem::needs_drop::<<$V as ArenaCached<'tcx>>::Allocated>() {
-                            <$V as ArenaCached>::alloc_in_arena(
-                                |v| _tcx.query_system.arenas.$name.alloc(v),
-                                provided_value,
-                            )
-                        } else {
-                            <$V as ArenaCached>::alloc_in_arena(
-                                |v| _tcx.arena.dropless.alloc(v),
-                                provided_value,
-                            )
-                        }
-                    }
-                    // Otherwise, the provided value is the value.
-                    (provided_value)
-                );
+                // For queries with the `arena_cache` modifier, store the
+                // provided value in an arena and get a reference to it.
+                let value: Value<'tcx> = query_if_arena!([$($modifiers)*] {
+                    <$V as $crate::query::arena_cached::ArenaCached>::alloc_in_arena(
+                        tcx,
+                        &tcx.query_system.arenas.$name,
+                        provided_value,
+                    )
+                } {
+                    // Otherwise, the provided value is the value (and `tcx` is unused).
+                    let _ = tcx;
+                    provided_value
+                });
                 erase::erase_val(value)
             }