Ensure type checking each function is stealable by other threads

Author: Zoxc

src/librustc/ty/mod.rs

@@ -30,7 +30,7 @@ use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
-use rustc_data_structures::sync::{self, par_for_each, Lrc};
+use rustc_data_structures::sync::{self, balance_par_for_each, Lrc};
 use rustc_hir as hir;
 use rustc_hir::def::{CtorKind, CtorOf, DefKind, Res};
 use rustc_hir::def_id::{CrateNum, DefId, DefIdMap, LocalDefId, CRATE_DEF_INDEX, LOCAL_CRATE};
@@ -2642,7 +2642,7 @@ impl<'tcx> TyCtxt<'tcx> {
     }
 
     pub fn par_body_owners<F: Fn(DefId) + sync::Sync + sync::Send>(self, f: F) {
-        par_for_each(&self.hir().krate().body_ids, |&body_id| {
+        balance_par_for_each(&self.hir().krate().body_ids, |&body_id| {
             f(self.hir().body_owner_def_id(body_id))
         });
     }

src/librustc_data_structures/sync.rs

@@ -227,6 +227,13 @@ cfg_if! {
             resume(panic);
             }
 
+        pub fn balance_par_for_each<T: IntoIterator>(
+            t: T,
+            for_each: impl FnMut(<<T as IntoIterator>::IntoIter as Iterator>::Item),
+        ) {
+            par_for_each(t, for_each)
+        }
+
         pub fn par_map<T: IntoIterator, R, C: FromIterator<R>>(
             t: T,
             mut map: impl FnMut(<<T as IntoIterator>::IntoIter as Iterator>::Item) -> R,
@@ -405,7 +412,9 @@ cfg_if! {
 
         pub use rayon_core::WorkerLocal;
 
-        use rayon::iter::{ParallelIterator, FromParallelIterator, IntoParallelIterator};
+        use rayon::iter::{
+            ParallelIterator, IndexedParallelIterator, FromParallelIterator, IntoParallelIterator
+        };
 
         pub fn par_for_each<T: IntoParallelIterator>(
             t: T,
@@ -421,6 +430,25 @@ cfg_if! {
             resume(panic);
         }
 
+        /// This does the same as `par_for_each`,
+        /// but each loop iteration is stealable by other threads.
+        pub fn balance_par_for_each<T: IntoParallelIterator>(
+            t: T,
+            for_each: impl Fn(
+                <<T as IntoParallelIterator>::Iter as ParallelIterator>::Item
+            ) + Sync + Send
+        )
+        where
+            <T as IntoParallelIterator>::Iter: IndexedParallelIterator
+        {
+            // We catch panics here ensuring that all the loop iterations execute.
+            let panic = Lock::new(None);
+            t.into_par_iter().with_max_len(1).for_each(|i| {
+                catch(&panic, || for_each(i));
+            });
+            resume(panic);
+        }
+
         pub fn par_map<
             T: IntoParallelIterator,
             R: Send,