Use `Freeze` for `SourceFile.external_src`

8 files changed, 91 insertions, 46 deletions

diff --git a/compiler/rustc_data_structures/src/sync/freeze.rs b/compiler/rustc_data_structures/src/sync/freeze.rs
index d9f1d72d851..88c710e5236 100644
--- a/compiler/rustc_data_structures/src/sync/freeze.rs
+++ b/compiler/rustc_data_structures/src/sync/freeze.rs
@@ -27,7 +27,26 @@ unsafe impl<T: DynSync + DynSend> DynSync for FreezeLock<T> {}
 impl<T> FreezeLock<T> {
     #[inline]
     pub fn new(value: T) -> Self {
-        Self { data: UnsafeCell::new(value), frozen: AtomicBool::new(false), lock: RwLock::new(()) }
+        Self::with(value, false)
+    }
+
+    #[inline]
+    pub fn frozen(value: T) -> Self {
+        Self::with(value, true)
+    }
+
+    #[inline]
+    pub fn with(value: T, frozen: bool) -> Self {
+        Self {
+            data: UnsafeCell::new(value),
+            frozen: AtomicBool::new(frozen),
+            lock: RwLock::new(()),
+        }
+    }
+
+    #[inline]
+    pub fn is_frozen(&self) -> bool {
+        self.frozen.load(Ordering::Acquire)
     }
 
     #[inline]
@@ -43,12 +62,25 @@ impl<T> FreezeLock<T> {
     }
 
     #[inline]
+    pub fn borrow(&self) -> FreezeReadGuard<'_, T> {
+        self.read()
+    }
+
+    #[inline]
     #[track_caller]
     pub fn write(&self) -> FreezeWriteGuard<'_, T> {
+        self.try_write().expect("still mutable")
+    }
+
+    #[inline]
+    pub fn try_write(&self) -> Option<FreezeWriteGuard<'_, T>> {
         let _lock_guard = self.lock.write();
         // Use relaxed ordering since we're in the write lock.
-        assert!(!self.frozen.load(Ordering::Relaxed), "still mutable");
-        FreezeWriteGuard { _lock_guard, lock: self, marker: PhantomData }
+        if self.frozen.load(Ordering::Relaxed) {
+            None
+        } else {
+            Some(FreezeWriteGuard { _lock_guard, lock: self, marker: PhantomData })
+        }
     }
 
     #[inline]
@@ -90,6 +122,15 @@ pub struct FreezeWriteGuard<'a, T> {
     marker: PhantomData<&'a mut T>,
 }
 
+impl<'a, T> FreezeWriteGuard<'a, T> {
+    pub fn freeze(self) -> &'a T {
+        self.lock.frozen.store(true, Ordering::Release);
+
+        // SAFETY: This is frozen so the data cannot be modified and shared access is sound.
+        unsafe { &*self.lock.data.get() }
+    }
+}
+
 impl<'a, T: 'a> Deref for FreezeWriteGuard<'a, T> {
     type Target = T;
     #[inline]
diff --git a/compiler/rustc_errors/src/annotate_snippet_emitter_writer.rs b/compiler/rustc_errors/src/annotate_snippet_emitter_writer.rs
index d7a008f9a57..5d3b2f45166 100644
--- a/compiler/rustc_errors/src/annotate_snippet_emitter_writer.rs
+++ b/compiler/rustc_errors/src/annotate_snippet_emitter_writer.rs
@@ -169,7 +169,7 @@ impl AnnotateSnippetEmitterWriter {
                         .map(|line| {
                             // Ensure the source file is present before we try
                             // to load a string from it.
-                            source_map.ensure_source_file_source_present(file.clone());
+                            source_map.ensure_source_file_source_present(&file);
                             (
                                 format!("{}", source_map.filename_for_diagnostics(&file.name)),
                                 source_string(file.clone(), &line),
diff --git a/compiler/rustc_errors/src/emitter.rs b/compiler/rustc_errors/src/emitter.rs
index da108327ae7..58be74f887b 100644
--- a/compiler/rustc_errors/src/emitter.rs
+++ b/compiler/rustc_errors/src/emitter.rs
@@ -1193,7 +1193,7 @@ impl EmitterWriter {
         let will_be_emitted = |span: Span| {
             !span.is_dummy() && {
                 let file = sm.lookup_source_file(span.hi());
-                sm.ensure_source_file_source_present(file)
+                sm.ensure_source_file_source_present(&file)
             }
         };
 
@@ -1388,7 +1388,7 @@ impl EmitterWriter {
         // Print out the annotate source lines that correspond with the error
         for annotated_file in annotated_files {
             // we can't annotate anything if the source is unavailable.
-            if !sm.ensure_source_file_source_present(annotated_file.file.clone()) {
+            if !sm.ensure_source_file_source_present(&annotated_file.file) {
                 if !self.short_message {
                     // We'll just print an unannotated message.
                     for (annotation_id, line) in annotated_file.lines.iter().enumerate() {
diff --git a/compiler/rustc_errors/src/json.rs b/compiler/rustc_errors/src/json.rs
index 390bf28df09..38667c5ff81 100644
--- a/compiler/rustc_errors/src/json.rs
+++ b/compiler/rustc_errors/src/json.rs
@@ -558,7 +558,7 @@ impl DiagnosticSpanLine {
             .span_to_lines(span)
             .map(|lines| {
                 // We can't get any lines if the source is unavailable.
-                if !je.sm.ensure_source_file_source_present(lines.file.clone()) {
+                if !je.sm.ensure_source_file_source_present(&lines.file) {
                     return vec![];
                 }
 
diff --git a/compiler/rustc_errors/src/lib.rs b/compiler/rustc_errors/src/lib.rs
index 55c4ec66cd9..6cb6549a00b 100644
--- a/compiler/rustc_errors/src/lib.rs
+++ b/compiler/rustc_errors/src/lib.rs
@@ -273,7 +273,7 @@ impl CodeSuggestion {
                 assert!(!lines.lines.is_empty() || bounding_span.is_dummy());
 
                 // We can't splice anything if the source is unavailable.
-                if !sm.ensure_source_file_source_present(lines.file.clone()) {
+                if !sm.ensure_source_file_source_present(&lines.file) {
                     return None;
                 }
 
diff --git a/compiler/rustc_metadata/src/rmeta/encoder.rs b/compiler/rustc_metadata/src/rmeta/encoder.rs
index 4f4351633a2..bace48cec44 100644
--- a/compiler/rustc_metadata/src/rmeta/encoder.rs
+++ b/compiler/rustc_metadata/src/rmeta/encoder.rs
@@ -280,8 +280,8 @@ impl<'a, 'tcx> Encodable<EncodeContext<'a, 'tcx>> for SpanData {
             // All of this logic ensures that the final result of deserialization is a 'normal'
             // Span that can be used without any additional trouble.
             let metadata_index = {
-                // Introduce a new scope so that we drop the 'lock()' temporary
-                match &*source_file.external_src.lock() {
+                // Introduce a new scope so that we drop the 'read()' temporary
+                match &*source_file.external_src.read() {
                     ExternalSource::Foreign { metadata_index, .. } => *metadata_index,
                     src => panic!("Unexpected external source {src:?}"),
                 }
diff --git a/compiler/rustc_span/src/lib.rs b/compiler/rustc_span/src/lib.rs
index f1a6e9059d7..8ef17ed671b 100644
--- a/compiler/rustc_span/src/lib.rs
+++ b/compiler/rustc_span/src/lib.rs
@@ -64,7 +64,7 @@ pub mod fatal_error;
 pub mod profiling;
 
 use rustc_data_structures::stable_hasher::{Hash128, Hash64, HashStable, StableHasher};
-use rustc_data_structures::sync::{Lock, Lrc};
+use rustc_data_structures::sync::{FreezeLock, FreezeWriteGuard, Lock, Lrc};
 
 use std::borrow::Cow;
 use std::cmp::{self, Ordering};
@@ -1206,7 +1206,6 @@ pub enum ExternalSourceKind {
     AbsentOk,
     /// A failed attempt has been made to load the external source.
     AbsentErr,
-    Unneeded,
 }
 
 impl ExternalSource {
@@ -1343,7 +1342,7 @@ pub struct SourceFile {
     pub src_hash: SourceFileHash,
     /// The external source code (used for external crates, which will have a `None`
     /// value as `self.src`.
-    pub external_src: Lock<ExternalSource>,
+    pub external_src: FreezeLock<ExternalSource>,
     /// The start position of this source in the `SourceMap`.
     pub start_pos: BytePos,
     /// The byte length of this source.
@@ -1368,7 +1367,10 @@ impl Clone for SourceFile {
             name: self.name.clone(),
             src: self.src.clone(),
             src_hash: self.src_hash,
-            external_src: Lock::new(self.external_src.borrow().clone()),
+            external_src: {
+                let lock = self.external_src.read();
+                FreezeLock::with(lock.clone(), self.external_src.is_frozen())
+            },
             start_pos: self.start_pos,
             source_len: self.source_len,
             lines: Lock::new(self.lines.borrow().clone()),
@@ -1488,7 +1490,7 @@ impl<D: Decoder> Decodable<D> for SourceFile {
             src_hash,
             // Unused - the metadata decoder will construct
             // a new SourceFile, filling in `external_src` properly
-            external_src: Lock::new(ExternalSource::Unneeded),
+            external_src: FreezeLock::frozen(ExternalSource::Unneeded),
             lines: Lock::new(lines),
             multibyte_chars,
             non_narrow_chars,
@@ -1530,7 +1532,7 @@ impl SourceFile {
             name,
             src: Some(Lrc::new(src)),
             src_hash,
-            external_src: Lock::new(ExternalSource::Unneeded),
+            external_src: FreezeLock::frozen(ExternalSource::Unneeded),
             start_pos: BytePos::from_u32(0),
             source_len: RelativeBytePos::from_u32(source_len),
             lines: Lock::new(SourceFileLines::Lines(lines)),
@@ -1612,35 +1614,37 @@ impl SourceFile {
     where
         F: FnOnce() -> Option<String>,
     {
-        if matches!(
-            *self.external_src.borrow(),
-            ExternalSource::Foreign { kind: ExternalSourceKind::AbsentOk, .. }
-        ) {
+        if !self.external_src.is_frozen() {
             let src = get_src();
-            let mut external_src = self.external_src.borrow_mut();
-            // Check that no-one else have provided the source while we were getting it
-            if let ExternalSource::Foreign {
-                kind: src_kind @ ExternalSourceKind::AbsentOk, ..
-            } = &mut *external_src
-            {
-                if let Some(mut src) = src {
-                    // The src_hash needs to be computed on the pre-normalized src.
-                    if self.src_hash.matches(&src) {
-                        normalize_src(&mut src);
-                        *src_kind = ExternalSourceKind::Present(Lrc::new(src));
-                        return true;
-                    }
+            let src = src.and_then(|mut src| {
+                // The src_hash needs to be computed on the pre-normalized src.
+                self.src_hash.matches(&src).then(|| {
+                    normalize_src(&mut src);
+                    src
+                })
+            });
+
+            self.external_src.try_write().map(|mut external_src| {
+                if let ExternalSource::Foreign {
+                    kind: src_kind @ ExternalSourceKind::AbsentOk,
+                    ..
+                } = &mut *external_src
+                {
+                    *src_kind = if let Some(src) = src {
+                        ExternalSourceKind::Present(Lrc::new(src))
+                    } else {
+                        ExternalSourceKind::AbsentErr
+                    };
                 } else {
-                    *src_kind = ExternalSourceKind::AbsentErr;
+                    panic!("unexpected state {:?}", *external_src)
                 }
 
-                false
-            } else {
-                self.src.is_some() || external_src.get_source().is_some()
-            }
-        } else {
-            self.src.is_some() || self.external_src.borrow().get_source().is_some()
+                // Freeze this so we don't try to load the source again.
+                FreezeWriteGuard::freeze(external_src)
+            });
         }
+
+        self.src.is_some() || self.external_src.read().get_source().is_some()
     }
 
     /// Gets a line from the list of pre-computed line-beginnings.
diff --git a/compiler/rustc_span/src/source_map.rs b/compiler/rustc_span/src/source_map.rs
index 81730f2f608..20bd57cae6d 100644
--- a/compiler/rustc_span/src/source_map.rs
+++ b/compiler/rustc_span/src/source_map.rs
@@ -340,7 +340,7 @@ impl SourceMap {
             name: filename,
             src: None,
             src_hash,
-            external_src: Lock::new(ExternalSource::Foreign {
+            external_src: FreezeLock::new(ExternalSource::Foreign {
                 kind: ExternalSourceKind::AbsentOk,
                 metadata_index,
             }),
@@ -564,7 +564,7 @@ impl SourceMap {
                 end: (local_end.sf.name.clone(), local_end.sf.start_pos),
             })))
         } else {
-            self.ensure_source_file_source_present(local_begin.sf.clone());
+            self.ensure_source_file_source_present(&local_begin.sf);
 
             let start_index = local_begin.pos.to_usize();
             let end_index = local_end.pos.to_usize();
@@ -581,7 +581,7 @@ impl SourceMap {
 
             if let Some(ref src) = local_begin.sf.src {
                 extract_source(src, start_index, end_index)
-            } else if let Some(src) = local_begin.sf.external_src.borrow().get_source() {
+            } else if let Some(src) = local_begin.sf.external_src.read().get_source() {
                 extract_source(src, start_index, end_index)
             } else {
                 Err(SpanSnippetError::SourceNotAvailable { filename: local_begin.sf.name.clone() })
@@ -873,7 +873,7 @@ impl SourceMap {
             let sp = sp.data();
             let local_begin = self.lookup_byte_offset(sp.lo);
             let start_index = local_begin.pos.to_usize();
-            let src = local_begin.sf.external_src.borrow();
+            let src = local_begin.sf.external_src.read();
 
             let snippet = if let Some(ref src) = local_begin.sf.src {
                 Some(&src[start_index..])
@@ -983,7 +983,7 @@ impl SourceMap {
             return 1;
         }
 
-        let src = local_begin.sf.external_src.borrow();
+        let src = local_begin.sf.external_src.read();
 
         let snippet = if let Some(src) = &local_begin.sf.src {
             src
@@ -1030,7 +1030,7 @@ impl SourceMap {
         self.files().iter().fold(0, |a, f| a + f.count_lines())
     }
 
-    pub fn ensure_source_file_source_present(&self, source_file: Lrc<SourceFile>) -> bool {
+    pub fn ensure_source_file_source_present(&self, source_file: &SourceFile) -> bool {
         source_file.add_external_src(|| {
             let FileName::Real(ref name) = source_file.name else {
                 return None;