Use an Arc to share caches between clones.

Author: cjgillot

compiler/rustc_middle/src/mir/basic_blocks.rs

@@ -1,4 +1,4 @@
-use std::sync::OnceLock;
+use std::sync::{Arc, OnceLock};
 
 use rustc_data_structures::graph;
 use rustc_data_structures::graph::dominators::{Dominators, dominators};
@@ -14,7 +14,8 @@ use crate::mir::{BasicBlock, BasicBlockData, START_BLOCK};
 #[derive(Clone, TyEncodable, TyDecodable, Debug, HashStable, TypeFoldable, TypeVisitable)]
 pub struct BasicBlocks<'tcx> {
     basic_blocks: IndexVec<BasicBlock, BasicBlockData<'tcx>>,
-    cache: Cache,
+    /// Use an `Arc` so we can share the cache when we clone the MIR body, as borrowck does.
+    cache: Arc<Cache>,
 }
 
 // Typically 95%+ of basic blocks have 4 or fewer predecessors.
@@ -38,9 +39,10 @@ struct Cache {
 impl<'tcx> BasicBlocks<'tcx> {
     #[inline]
     pub fn new(basic_blocks: IndexVec<BasicBlock, BasicBlockData<'tcx>>) -> Self {
-        BasicBlocks { basic_blocks, cache: Cache::default() }
+        BasicBlocks { basic_blocks, cache: Arc::new(Cache::default()) }
     }
 
+    #[inline]
     pub fn dominators(&self) -> &Dominators<BasicBlock> {
         self.cache.dominators.get_or_init(|| dominators(self))
     }
@@ -104,7 +106,14 @@ impl<'tcx> BasicBlocks<'tcx> {
     /// All other methods that allow you to mutate the basic blocks also call this method
     /// themselves, thereby avoiding any risk of accidentally cache invalidation.
     pub fn invalidate_cfg_cache(&mut self) {
-        self.cache = Cache::default();
+        if let Some(cache) = Arc::get_mut(&mut self.cache) {
+            // If we only have a single reference to this cache, clear it.
+            *cache = Cache::default();
+        } else {
+            // If we have several references to this cache, overwrite the pointer itself so other
+            // users can continue to use their (valid) cache.
+            self.cache = Arc::new(Cache::default());
+        }
     }
 }