Register the file while computing its start position.

4 files changed, 43 insertions, 45 deletions

diff --git a/compiler/rustc_span/src/lib.rs b/compiler/rustc_span/src/lib.rs
index b1cde3093ef..34714957aef 100644
--- a/compiler/rustc_span/src/lib.rs
+++ b/compiler/rustc_span/src/lib.rs
@@ -1519,7 +1519,11 @@ impl fmt::Debug for SourceFile {
 }
 
 impl SourceFile {
-    pub fn new(name: FileName, mut src: String, hash_kind: SourceFileHashAlgorithm) -> Self {
+    pub fn new(
+        name: FileName,
+        mut src: String,
+        hash_kind: SourceFileHashAlgorithm,
+    ) -> Result<Self, OffsetOverflowError> {
         // Compute the file hash before any normalization.
         let src_hash = SourceFileHash::new(hash_kind, &src);
         let normalized_pos = normalize_src(&mut src);
@@ -1530,25 +1534,25 @@ impl SourceFile {
             hasher.finish()
         };
         let source_len = src.len();
-        assert!(source_len <= u32::MAX as usize);
+        let source_len = u32::try_from(source_len).map_err(|_| OffsetOverflowError)?;
 
         let (lines, multibyte_chars, non_narrow_chars) =
             analyze_source_file::analyze_source_file(&src);
 
-        SourceFile {
+        Ok(SourceFile {
             name,
             src: Some(Lrc::new(src)),
             src_hash,
             external_src: Lock::new(ExternalSource::Unneeded),
             start_pos: BytePos::from_u32(0),
-            source_len: RelativeBytePos::from_usize(source_len),
+            source_len: RelativeBytePos::from_u32(source_len),
             lines: Lock::new(SourceFileLines::Lines(lines)),
             multibyte_chars,
             non_narrow_chars,
             normalized_pos,
             name_hash,
             cnum: LOCAL_CRATE,
-        }
+        })
     }
 
     pub fn lines<F, R>(&self, f: F) -> R
diff --git a/compiler/rustc_span/src/source_map.rs b/compiler/rustc_span/src/source_map.rs
index 07483280f41..2290f9c7c6f 100644
--- a/compiler/rustc_span/src/source_map.rs
+++ b/compiler/rustc_span/src/source_map.rs
@@ -267,10 +267,13 @@ impl SourceMap {
         self.files.borrow().stable_id_to_source_file.get(&stable_id).cloned()
     }
 
-    fn allocate_address_space(&self, size: usize) -> Result<usize, OffsetOverflowError> {
-        let size = u32::try_from(size).map_err(|_| OffsetOverflowError)?;
+    fn register_source_file(
+        &self,
+        mut file: SourceFile,
+    ) -> Result<Lrc<SourceFile>, OffsetOverflowError> {
+        let size = file.source_len.to_u32();
 
-        loop {
+        let start_pos = loop {
             let current = self.used_address_space.load(Ordering::Relaxed);
             let next = current
                 .checked_add(size)
@@ -284,9 +287,20 @@ impl SourceMap {
                 .compare_exchange(current, next, Ordering::Relaxed, Ordering::Relaxed)
                 .is_ok()
             {
-                return Ok(usize::try_from(current).unwrap());
+                break usize::try_from(current).unwrap();
             }
-        }
+        };
+
+        file.start_pos = BytePos::from_usize(start_pos);
+        let file_id = StableSourceFileId::new(&file);
+
+        let mut files = self.files.borrow_mut();
+
+        let file = Lrc::new(file);
+        files.source_files.push(file.clone());
+        files.stable_id_to_source_file.insert(file_id, file.clone());
+
+        Ok(file)
     }
 
     /// Creates a new `SourceFile`.
@@ -310,29 +324,18 @@ impl SourceMap {
         let (filename, _) = self.path_mapping.map_filename_prefix(&filename);
 
         let file_id = StableSourceFileId::new_from_name(&filename, LOCAL_CRATE);
-
-        let lrc_sf = match self.source_file_by_stable_id(file_id) {
-            Some(lrc_sf) => lrc_sf,
+        match self.source_file_by_stable_id(file_id) {
+            Some(lrc_sf) => Ok(lrc_sf),
             None => {
-                let mut source_file = SourceFile::new(filename, src, self.hash_kind);
+                let source_file = SourceFile::new(filename, src, self.hash_kind)?;
 
                 // Let's make sure the file_id we generated above actually matches
                 // the ID we generate for the SourceFile we just created.
                 debug_assert_eq!(StableSourceFileId::new(&source_file), file_id);
 
-                let start_pos = self.allocate_address_space(source_file.source_len.to_usize())?;
-                source_file.start_pos = BytePos::from_usize(start_pos);
-
-                let mut files = self.files.borrow_mut();
-
-                let source_file = Lrc::new(source_file);
-                files.source_files.push(source_file.clone());
-                files.stable_id_to_source_file.insert(file_id, source_file.clone());
-
-                source_file
+                self.register_source_file(source_file)
             }
-        };
-        Ok(lrc_sf)
+        }
     }
 
     /// Allocates a new `SourceFile` representing a source file from an external
@@ -344,7 +347,7 @@ impl SourceMap {
         filename: FileName,
         src_hash: SourceFileHash,
         name_hash: Hash128,
-        source_len: usize,
+        source_len: u32,
         cnum: CrateNum,
         file_local_lines: Lock<SourceFileLines>,
         multibyte_chars: Vec<MultiByteChar>,
@@ -352,13 +355,9 @@ impl SourceMap {
         normalized_pos: Vec<NormalizedPos>,
         metadata_index: u32,
     ) -> Lrc<SourceFile> {
-        let start_pos = self
-            .allocate_address_space(source_len)
-            .expect("not enough address space for imported source file");
+        let source_len = RelativeBytePos::from_u32(source_len);
 
-        let source_len = RelativeBytePos::from_usize(source_len);
-
-        let source_file = Lrc::new(SourceFile {
+        let source_file = SourceFile {
             name: filename,
             src: None,
             src_hash,
@@ -366,7 +365,7 @@ impl SourceMap {
                 kind: ExternalSourceKind::AbsentOk,
                 metadata_index,
             }),
-            start_pos: BytePos::from_usize(start_pos),
+            start_pos: BytePos(0),
             source_len,
             lines: file_local_lines,
             multibyte_chars,
@@ -374,16 +373,10 @@ impl SourceMap {
             normalized_pos,
             name_hash,
             cnum,
-        });
-
-        let mut files = self.files.borrow_mut();
-
-        files.source_files.push(source_file.clone());
-        files
-            .stable_id_to_source_file
-            .insert(StableSourceFileId::new(&source_file), source_file.clone());
+        };
 
-        source_file
+        self.register_source_file(source_file)
+            .expect("not enough address space for imported source file")
     }
 
     /// If there is a doctest offset, applies it to the line.
diff --git a/compiler/rustc_span/src/source_map/tests.rs b/compiler/rustc_span/src/source_map/tests.rs
index d5cb3dac55c..7689e6afac5 100644
--- a/compiler/rustc_span/src/source_map/tests.rs
+++ b/compiler/rustc_span/src/source_map/tests.rs
@@ -244,7 +244,7 @@ fn t10() {
         name,
         src_hash,
         name_hash,
-        source_len.to_usize(),
+        source_len.to_u32(),
         CrateNum::new(0),
         lines,
         multibyte_chars,
diff --git a/compiler/rustc_span/src/tests.rs b/compiler/rustc_span/src/tests.rs
index 3b69295bb93..a980ee8d9e0 100644
--- a/compiler/rustc_span/src/tests.rs
+++ b/compiler/rustc_span/src/tests.rs
@@ -4,7 +4,8 @@ use super::*;
 fn test_lookup_line() {
     let source = "abcdefghijklm\nabcdefghij\n...".to_owned();
     let mut sf =
-        SourceFile::new(FileName::Anon(Hash64::ZERO), source, SourceFileHashAlgorithm::Sha256);
+        SourceFile::new(FileName::Anon(Hash64::ZERO), source, SourceFileHashAlgorithm::Sha256)
+            .unwrap();
     sf.start_pos = BytePos(3);
     sf.lines(|lines| {
         assert_eq!(lines, &[RelativeBytePos(0), RelativeBytePos(14), RelativeBytePos(25)])