wip

crates/ruff_benchmark/benches/ty.rs

@@ -141,7 +141,7 @@ fn benchmark_incremental(criterion: &mut Criterion) {
                 &case.file_path,
                 format!(
                     "{}\n# A comment\n",
-                    source_text(&case.db, case.file).load().as_str()
+                    source_text(&case.db, case.file).load(&case.db).as_str()
                 ),
             )
             .unwrap();

crates/ruff_db/src/diagnostic/render.rs

@@ -645,7 +645,7 @@ impl FileResolver for &dyn Db {
 
     fn input(&self, file: File) -> Input {
         Input {
-            text: source_text(*self, file).load(),
+            text: source_text(*self, file).load(*self),
             line_index: line_index(*self, file),
         }
     }
@@ -2159,7 +2159,7 @@ watermelon
             let span = self.path(path);
 
             let file = span.expect_ty_file();
-            let text = source_text(&self.db, file).load();
+            let text = source_text(&self.db, file).load(&self.db);
             let line_index = line_index(&self.db, file);
             let source = SourceCode::new(text.as_str(), &line_index);
 

crates/ruff_db/src/parsed.rs

@@ -30,7 +30,7 @@ pub fn parsed_module(db: &dyn Db, file: File) -> ParsedModule {
 }
 
 pub fn parsed_module_impl(db: &dyn Db, file: File) -> Parsed<ModModule> {
-    let source = source_text(db, file).load();
+    let source = source_text(db, file).load(db);
     let ty = file.source_type(db);
 
     let target_version = db.python_version();

crates/ruff_db/src/source.rs

@@ -14,6 +14,15 @@ use crate::files::{File, FilePath};
 /// Reads the source text of a python text file (must be valid UTF8) or notebook.
 #[salsa::tracked(no_eq)]
 pub fn source_text(db: &dyn Db, file: File) -> SourceText {
+    let source_text = source_text_impl(db, file);
+
+    SourceText {
+        file,
+        inner: Arc::new(ArcSwapOption::new(Some(Arc::new(source_text)))),
+    }
+}
+
+fn source_text_impl(db: &dyn Db, file: File) -> SourceTextInner {
     let path = file.path(db);
     let _span = tracing::trace_span!("source_text", file = %path).entered();
     let mut read_error = None;
@@ -38,12 +47,10 @@ pub fn source_text(db: &dyn Db, file: File) -> SourceText {
             .into()
     };
 
-    SourceText {
-        inner: Arc::new(ArcSwapOption::new(Some(Arc::new(SourceTextInner {
-            kind,
-            read_error,
-            _count: Count::new(),
-        })))),
+    SourceTextInner {
+        kind,
+        read_error,
+        _count: Count::new(),
     }
 }
 
@@ -68,6 +75,7 @@ fn is_notebook(path: &FilePath) -> bool {
 /// Cheap cloneable in `O(1)`.
 #[derive(Clone)]
 pub struct SourceText {
+    file: File,
     inner: Arc<ArcSwapOption<SourceTextInner>>,
 }
 
@@ -80,9 +88,12 @@ impl PartialEq for SourceText {
 impl Eq for SourceText {}
 
 impl SourceText {
-    pub fn load(&self) -> SourceTextRef {
+    pub fn load(&self, db: &dyn Db) -> SourceTextRef {
         SourceTextRef {
-            inner: self.inner.load_full().unwrap(),
+            inner: self
+                .inner
+                .load_full()
+                .unwrap_or_else(|| Arc::new(source_text_impl(db, self.file))),
         }
     }
 
@@ -194,7 +205,7 @@ pub enum SourceTextError {
 pub fn line_index(db: &dyn Db, file: File) -> LineIndex {
     let _span = tracing::trace_span!("line_index", ?file).entered();
 
-    let source = source_text(db, file).load();
+    let source = source_text(db, file).load(db);
 
     LineIndex::from_source_text(&source)
 }
@@ -221,11 +232,11 @@ mod tests {
 
         let file = system_path_to_file(&db, path).unwrap();
 
-        assert_eq!(source_text(&db, file).load().as_str(), "x = 10");
+        assert_eq!(source_text(&db, file).load(&db).as_str(), "x = 10");
 
         db.write_file(path, "x = 20").unwrap();
 
-        assert_eq!(source_text(&db, file).load().as_str(), "x = 20");
+        assert_eq!(source_text(&db, file).load(&db).as_str(), "x = 20");
 
         Ok(())
     }
@@ -239,13 +250,13 @@ mod tests {
 
         let file = system_path_to_file(&db, path).unwrap();
 
-        assert_eq!(source_text(&db, file).load().as_str(), "x = 10");
+        assert_eq!(source_text(&db, file).load(&db).as_str(), "x = 10");
 
         // Change the file permission only
         file.set_permissions(&mut db).to(Some(0o777));
 
         db.clear_salsa_events();
-        assert_eq!(source_text(&db, file).load().as_str(), "x = 10");
+        assert_eq!(source_text(&db, file).load(&db).as_str(), "x = 10");
 
         let events = db.take_salsa_events();
 
@@ -267,7 +278,7 @@ mod tests {
 
         let file = system_path_to_file(&db, path).unwrap();
         let index = line_index(&db, file);
-        let source = source_text(&db, file).load();
+        let source = source_text(&db, file).load(&db);
 
         assert_eq!(index.line_count(), 2);
         assert_eq!(
@@ -309,7 +320,7 @@ mod tests {
         )?;
 
         let file = system_path_to_file(&db, path).unwrap();
-        let source = source_text(&db, file).load();
+        let source = source_text(&db, file).load(&db);
 
         assert!(source.is_notebook());
         assert_eq!(source.as_str(), "x = 10\n");

crates/ruff_python_formatter/src/lib.rs

@@ -172,7 +172,7 @@ pub fn formatted_file(db: &dyn Db, file: File) -> Result<Option<String>, FormatM
     }
 
     let comment_ranges = CommentRanges::from(parsed.tokens());
-    let source = source_text(db.upcast(), file).load();
+    let source = source_text(db.upcast(), file).load(db.upcast());
 
     let formatted = format_node(&parsed, &comment_ranges, &source, options)?;
     let printed = formatted.print()?;

crates/ty/tests/file_watching.rs

@@ -643,7 +643,7 @@ fn changed_file() -> anyhow::Result<()> {
     let foo_path = case.project_path("foo.py");
 
     let foo = case.system_file(&foo_path)?;
-    assert_eq!(source_text(case.db(), foo).load().as_str(), foo_source);
+    assert_eq!(source_text(case.db(), foo).load(case.db()).as_str(), foo_source);
     case.assert_indexed_project_files([foo]);
 
     update_file(&foo_path, "print('Version 2')")?;
@@ -655,7 +655,7 @@ fn changed_file() -> anyhow::Result<()> {
     case.apply_changes(changes, None);
 
     assert_eq!(
-        source_text(case.db(), foo).load().as_str(),
+        source_text(case.db(), foo).load(case.db()).as_str(),
         "print('Version 2')"
     );
     case.assert_indexed_project_files([foo]);
@@ -1266,11 +1266,11 @@ fn hard_links_in_project() -> anyhow::Result<()> {
     let bar = case.system_file(&bar_path).unwrap();
 
     assert_eq!(
-        source_text(case.db(), foo).load().as_str(),
+        source_text(case.db(), foo).load(case.db()).as_str(),
         "print('Version 1')"
     );
     assert_eq!(
-        source_text(case.db(), bar).load().as_str(),
+        source_text(case.db(), bar).load(case.db()).as_str(),
         "print('Version 1')"
     );
     case.assert_indexed_project_files([bar, foo]);
@@ -1283,14 +1283,14 @@ fn hard_links_in_project() -> anyhow::Result<()> {
     case.apply_changes(changes, None);
 
     assert_eq!(
-        source_text(case.db(), foo).load().as_str(),
+        source_text(case.db(), foo).load(case.db()).as_str(),
         "print('Version 2')"
     );
 
     // macOS is the only platform that emits events for every hardlink.
     if cfg!(target_os = "macos") {
         assert_eq!(
-            source_text(case.db(), bar).load().as_str(),
+            source_text(case.db(), bar).load(case.db()).as_str(),
             "print('Version 2')"
         );
     }
@@ -1350,11 +1350,11 @@ fn hard_links_to_target_outside_project() -> anyhow::Result<()> {
     let bar = case.system_file(&bar_path).unwrap();
 
     assert_eq!(
-        source_text(case.db(), foo).load().as_str(),
+        source_text(case.db(), foo).load(case.db()).as_str(),
         "print('Version 1')"
     );
     assert_eq!(
-        source_text(case.db(), bar).load().as_str(),
+        source_text(case.db(), bar).load(case.db()).as_str(),
         "print('Version 1')"
     );
 
@@ -1366,7 +1366,7 @@ fn hard_links_to_target_outside_project() -> anyhow::Result<()> {
     case.apply_changes(changes, None);
 
     assert_eq!(
-        source_text(case.db(), bar).load().as_str(),
+        source_text(case.db(), bar).load(case.db()).as_str(),
         "print('Version 2')"
     );
 
@@ -1475,7 +1475,7 @@ mod unix {
         let baz_file = baz.file().unwrap();
 
         assert_eq!(
-            source_text(case.db(), baz_file).load().as_str(),
+            source_text(case.db(), baz_file).load(case.db()).as_str(),
             "def baz(): ..."
         );
         assert_eq!(
@@ -1494,7 +1494,7 @@ mod unix {
         case.apply_changes(changes, None);
 
         assert_eq!(
-            source_text(case.db(), baz_file).load().as_str(),
+            source_text(case.db(), baz_file).load(case.db()).as_str(),
             "def baz(): print('Version 2')"
         );
 
@@ -1507,7 +1507,7 @@ mod unix {
         case.apply_changes(changes, None);
 
         assert_eq!(
-            source_text(case.db(), baz_file).load().as_str(),
+            source_text(case.db(), baz_file).load(case.db()).as_str(),
             "def baz(): print('Version 3')"
         );
 
@@ -1560,12 +1560,12 @@ mod unix {
         let patched_bar_baz_file = case.system_file(&patched_bar_baz).unwrap();
 
         assert_eq!(
-            source_text(case.db(), patched_bar_baz_file).load().as_str(),
+            source_text(case.db(), patched_bar_baz_file).load(case.db()).as_str(),
             "def baz(): ..."
         );
 
         assert_eq!(
-            source_text(case.db(), baz_file).load().as_str(),
+            source_text(case.db(), baz_file).load(case.db()).as_str(),
             "def baz(): ..."
         );
         assert_eq!(baz_file.path(case.db()).as_system_path(), Some(&*bar_baz));
@@ -1594,10 +1594,10 @@ mod unix {
         //
         // That's why I think it's fine to not support this case for now.
 
-        let patched_baz_text = source_text(case.db(), patched_bar_baz_file).load();
+        let patched_baz_text = source_text(case.db(), patched_bar_baz_file).load(case.db());
         let did_update_patched_baz = patched_baz_text.as_str() == "def baz(): print('Version 2')";
 
-        let bar_baz_text = source_text(case.db(), baz_file).load();
+        let bar_baz_text = source_text(case.db(), baz_file).load(case.db());
         let did_update_bar_baz = bar_baz_text.as_str() == "def baz(): print('Version 2')";
 
         assert!(
@@ -1671,12 +1671,12 @@ mod unix {
         let baz_original_file = case.system_file(&baz_original).unwrap();
 
         assert_eq!(
-            source_text(case.db(), baz_original_file).load().as_str(),
+            source_text(case.db(), baz_original_file).load(case.db()).as_str(),
             "def baz(): ..."
         );
 
         assert_eq!(
-            source_text(case.db(), baz_site_packages).load().as_str(),
+            source_text(case.db(), baz_site_packages).load(case.db()).as_str(),
             "def baz(): ..."
         );
         assert_eq!(
@@ -1695,7 +1695,7 @@ mod unix {
         case.apply_changes(changes, None);
 
         assert_eq!(
-            source_text(case.db(), baz_original_file).load().as_str(),
+            source_text(case.db(), baz_original_file).load(case.db()).as_str(),
             "def baz(): print('Version 2')"
         );
 
@@ -1707,7 +1707,7 @@ mod unix {
         // it doesn't seem worth doing considering that as prominent tools like PyCharm don't support it.
         // Pyright does support it, thanks to chokidar.
         assert_ne!(
-            source_text(case.db(), baz_site_packages).load().as_str(),
+            source_text(case.db(), baz_site_packages).load(case.db()).as_str(),
             "def baz(): print('Version 2')"
         );
 

crates/ty_ide/src/inlay_hints.rs

@@ -216,7 +216,7 @@ mod tests {
         fn inlay_hints(&self) -> String {
             let hints = inlay_hints(&self.db, self.file, self.range);
 
-            let mut buf = source_text(&self.db, self.file).load().as_str().to_string();
+            let mut buf = source_text(&self.db, self.file).load(&self.db).as_str().to_string();
 
             let mut offset = 0;
 

crates/ty_project/src/lib.rs

@@ -479,7 +479,7 @@ impl Project {
         // Abort checking if there are IO errors.
         let source = source_text(db.upcast(), file);
 
-        if let Some(read_error) = source.load().read_error() {
+        if let Some(read_error) = source.load(db.upcast()).read_error() {
             diagnostics.push(
                 IOErrorDiagnostic {
                     file: Some(file),
@@ -749,7 +749,7 @@ mod tests {
         db.memory_file_system().remove_file(path)?;
         file.sync(&mut db);
 
-        assert_eq!(source_text(&db, file).load().as_str(), "");
+        assert_eq!(source_text(&db, file).load(&db).as_str(), "");
         assert_eq!(
             db.project()
                 .check_file_impl(&db, file)
@@ -766,7 +766,7 @@ mod tests {
         // content returned by `source_text` remains unchanged, but the diagnostics should get updated.
         db.write_file(path, "").unwrap();
 
-        assert_eq!(source_text(&db, file).load().as_str(), "");
+        assert_eq!(source_text(&db, file).load(&db).as_str(), "");
         assert_eq!(
             db.project()
                 .check_file_impl(&db, file)

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -2276,7 +2276,7 @@ impl SemanticSyntaxContext for SemanticIndexBuilder<'_, '_> {
     }
 
     fn source(&self) -> Cow<'_, str> {
-        let source_text = self.source_text().load();
+        let source_text = self.source_text().load(self.db.upcast());
         Cow::Owned(source_text.as_str().to_string())
     }
 
@@ -2366,7 +2366,7 @@ impl SemanticSyntaxContext for SemanticIndexBuilder<'_, '_> {
     }
 
     fn in_notebook(&self) -> bool {
-        self.source_text().load().is_notebook()
+        self.source_text().load(self.db.upcast()).is_notebook()
     }
 
     fn report_semantic_error(&self, error: SemanticSyntaxError) {

crates/ty_python_semantic/src/suppression.rs

@@ -89,7 +89,7 @@ declare_lint! {
 #[salsa::tracked(returns(ref))]
 pub(crate) fn suppressions(db: &dyn Db, file: File) -> Suppressions {
     let parsed = parsed_module(db.upcast(), file).load(db.upcast());
-    let source = source_text(db.upcast(), file).load();
+    let source = source_text(db.upcast(), file).load(db.upcast());
 
     let mut builder = SuppressionsBuilder::new(&source, db.lint_registry());
     let mut line_start = TextSize::default();

crates/ty_python_semantic/src/types/definition.rs

@@ -32,7 +32,7 @@ impl TypeDefinition<'_> {
         match self {
             Self::Module(module) => {
                 let file = module.file()?;
-                let source = source_text(db.upcast(), file).load();
+                let source = source_text(db.upcast(), file).load(db.upcast());
                 Some(FileRange::new(file, TextRange::up_to(source.text_len())))
             }
             Self::Class(definition)

crates/ty_python_semantic/src/types/string_annotation.rs

@@ -137,7 +137,7 @@ pub(crate) fn parse_string_annotation(
     let _span = tracing::trace_span!("parse_string_annotation", string=?string_expr.range(), ?file)
         .entered();
 
-    let source = source_text(db.upcast(), file).load();
+    let source = source_text(db.upcast(), file).load(db.upcast());
 
     if let Some(string_literal) = string_expr.as_single_part_string() {
         let prefix = string_literal.flags.prefix();

crates/ty_server/src/document/location.rs

@@ -31,14 +31,14 @@ impl ToLink for NavigationTarget {
     ) -> Option<lsp_types::LocationLink> {
         let file = self.file();
         let uri = file_to_url(db.upcast(), file)?;
-        let source = source_text(db.upcast(), file).load();
+        let source = source_text(db.upcast(), file).load(db.upcast());
         let index = line_index(db.upcast(), file);
 
         let target_range = self.full_range().to_lsp_range(&source, &index, encoding);
         let selection_range = self.focus_range().to_lsp_range(&source, &index, encoding);
 
         let src = src.map(|src| {
-            let source = source_text(db.upcast(), src.file()).load();
+            let source = source_text(db.upcast(), src.file()).load(db.upcast());
             let index = line_index(db.upcast(), src.file());
 
             src.range().to_lsp_range(&source, &index, encoding)

crates/ty_server/src/document/range.rs

@@ -163,7 +163,7 @@ impl FileRangeExt for FileRange {
     fn to_location(&self, db: &dyn Db, encoding: PositionEncoding) -> Option<Location> {
         let file = self.file();
         let uri = file_to_url(db, file)?;
-        let source = source_text(db.upcast(), file).load();
+        let source = source_text(db.upcast(), file).load(db.upcast());
         let line_index = line_index(db.upcast(), file);
 
         let range = self.range().to_lsp_range(&source, &line_index, encoding);

crates/ty_server/src/server/api/diagnostics.rs

@@ -180,7 +180,7 @@ fn to_lsp_diagnostic(
     let range = if let Some(span) = diagnostic.primary_span() {
         let file = span.expect_ty_file();
         let index = line_index(db.upcast(), file);
-        let source = source_text(db.upcast(), file).load();
+        let source = source_text(db.upcast(), file).load(db.upcast());
 
         span.range()
             .map(|range| range.to_lsp_range(&source, &index, encoding))

crates/ty_server/src/server/api/requests/completion.rs

@@ -39,7 +39,7 @@ impl BackgroundDocumentRequestHandler for CompletionRequestHandler {
             return Ok(None);
         };
 
-        let source = source_text(db, file).load();
+        let source = source_text(db, file).load(db);
         let line_index = line_index(db, file);
         let offset = params.text_document_position.position.to_text_size(
             &source,

crates/ty_server/src/server/api/requests/goto_type_definition.rs

@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for GotoTypeDefinitionRequestHandler {
             return Ok(None);
         };
 
-        let source = source_text(db, file).load();
+        let source = source_text(db, file).load(db);
         let line_index = line_index(db, file);
         let offset = params.text_document_position_params.position.to_text_size(
             &source,

crates/ty_server/src/server/api/requests/hover.rs

@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for HoverRequestHandler {
             return Ok(None);
         };
 
-        let source = source_text(db, file).load();
+        let source = source_text(db, file).load(db);
         let line_index = line_index(db, file);
         let offset = params.text_document_position_params.position.to_text_size(
             &source,

crates/ty_server/src/server/api/requests/inlay_hints.rs

@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for InlayHintRequestHandler {
         };
 
         let index = line_index(db, file);
-        let source = source_text(db, file).load();
+        let source = source_text(db, file).load(db);
 
         let range = params
             .range

crates/ty_test/src/assertion.rs

@@ -55,7 +55,7 @@ pub(crate) struct InlineFileAssertions {
 
 impl InlineFileAssertions {
     pub(crate) fn from_file(db: &Db, file: File) -> Self {
-        let source = source_text(db, file).load();
+        let source = source_text(db, file).load(db);
         let lines = line_index(db, file);
         let parsed = parsed_module(db, file).load(db);
         let comment_ranges = CommentRanges::from(parsed.tokens());

crates/ty_test/src/matcher.rs

@@ -211,7 +211,7 @@ impl Matcher {
     fn from_file(db: &Db, file: File) -> Self {
         Self {
             line_index: line_index(db, file),
-            source: source_text(db, file).load(),
+            source: source_text(db, file).load(db),
         }
     }
 

crates/ty_wasm/src/lib.rs

@@ -219,7 +219,7 @@ impl Workspace {
 
     #[wasm_bindgen(js_name = "sourceText")]
     pub fn source_text(&self, file_id: &FileHandle) -> Result<String, Error> {
-        let source_text = ruff_db::source::source_text(&self.db, file_id.file).load();
+        let source_text = ruff_db::source::source_text(&self.db, file_id.file).load(&self.db);
 
         Ok(source_text.to_string())
     }
@@ -230,7 +230,7 @@ impl Workspace {
         file_id: &FileHandle,
         position: Position,
     ) -> Result<Vec<LocationLink>, Error> {
-        let source = source_text(&self.db, file_id.file).load();
+        let source = source_text(&self.db, file_id.file).load(&self.db);
         let index = line_index(&self.db, file_id.file);
 
         let offset = position.to_text_size(&source, &index, self.position_encoding)?;
@@ -269,7 +269,7 @@ impl Workspace {
 
     #[wasm_bindgen]
     pub fn hover(&self, file_id: &FileHandle, position: Position) -> Result<Option<Hover>, Error> {
-        let source = source_text(&self.db, file_id.file).load();
+        let source = source_text(&self.db, file_id.file).load(&self.db);
         let index = line_index(&self.db, file_id.file);
 
         let offset = position.to_text_size(&source, &index, self.position_encoding)?;
@@ -299,7 +299,7 @@ impl Workspace {
         file_id: &FileHandle,
         position: Position,
     ) -> Result<Vec<Completion>, Error> {
-        let source = source_text(&self.db, file_id.file).load();
+        let source = source_text(&self.db, file_id.file).load(&self.db);
         let index = line_index(&self.db, file_id.file);
 
         let offset = position.to_text_size(&source, &index, self.position_encoding)?;
@@ -317,7 +317,7 @@ impl Workspace {
     #[wasm_bindgen(js_name = "inlayHints")]
     pub fn inlay_hints(&self, file_id: &FileHandle, range: Range) -> Result<Vec<InlayHint>, Error> {
         let index = line_index(&self.db, file_id.file);
-        let source = source_text(&self.db, file_id.file).load();
+        let source = source_text(&self.db, file_id.file).load(&self.db);
 
         let result = inlay_hints(
             &self.db,
@@ -440,7 +440,7 @@ impl Range {
         position_encoding: PositionEncoding,
     ) -> Self {
         let index = line_index(db.upcast(), file_range.file());
-        let source = source_text(db.upcast(), file_range.file()).load();
+        let source = source_text(db.upcast(), file_range.file()).load(db.upcast());
 
         Self::from_text_range(file_range.range(), &index, &source, position_encoding)
     }