Lazy compute source file line index

crates/ruff/src/linter.rs

@@ -371,8 +371,12 @@ fn diagnostics_to_messages(
     directives: &Directives,
 ) -> Vec<Message> {
     let file = once_cell::unsync::Lazy::new(|| {
-        let mut builder = SourceFileBuilder::new(&path.to_string_lossy());
-        builder.set_source_code(&locator.to_source_code());
+        let mut builder =
+            SourceFileBuilder::new(path.to_string_lossy().as_ref(), locator.contents());
+
+        if let Some(line_index) = locator.line_index() {
+            builder.set_line_index(line_index.clone());
+        }
 
         builder.finish()
     });

crates/ruff/src/message/diff.rs

@@ -1,7 +1,7 @@
 use crate::message::Message;
 use colored::{Color, ColoredString, Colorize, Styles};
 use ruff_diagnostics::Fix;
-use ruff_python_ast::source_code::{OneIndexed, SourceCode};
+use ruff_python_ast::source_code::{OneIndexed, SourceFile};
 use ruff_text_size::{TextRange, TextSize};
 use similar::{ChangeTag, TextDiff};
 use std::fmt::{Display, Formatter};
@@ -17,7 +17,7 @@ use std::num::NonZeroUsize;
 /// * Compute the diff from the [`Edit`] because diff calculation is expensive.
 pub(super) struct Diff<'a> {
     fix: &'a Fix,
-    source_code: SourceCode<'a, 'a>,
+    source_code: &'a SourceFile,
 }
 
 impl<'a> Diff<'a> {
@@ -26,7 +26,7 @@ impl<'a> Diff<'a> {
             None
         } else {
             Some(Diff {
-                source_code: message.file.source_code(),
+                source_code: &message.file,
                 fix: &message.fix,
             })
         }
@@ -35,18 +35,22 @@ impl<'a> Diff<'a> {
 
 impl Display for Diff<'_> {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        let mut output = String::with_capacity(self.source_code.text().len());
+        let mut output = String::with_capacity(self.source_code.source_text().len());
         let mut last_end = TextSize::default();
 
         for edit in self.fix.edits() {
-            output.push_str(&self.source_code.text()[TextRange::new(last_end, edit.start())]);
+            output.push_str(
+                self
+                    .source_code
+                    .slice(TextRange::new(last_end, edit.start())),
+            );
             output.push_str(edit.content().unwrap_or_default());
             last_end = edit.end();
         }
 
-        output.push_str(&self.source_code.text()[usize::from(last_end)..]);
+        output.push_str(&self.source_code.source_text()[usize::from(last_end)..]);
 
-        let diff = TextDiff::from_lines(self.source_code.text(), &output);
+        let diff = TextDiff::from_lines(self.source_code.source_text(), &output);
 
         writeln!(f, "{}", "ℹ Suggested fix".blue())?;
 

crates/ruff/src/message/json.rs

@@ -1,5 +1,7 @@
 use crate::message::{Emitter, EmitterContext, Message};
 use crate::registry::AsRule;
+use ruff_diagnostics::Edit;
+use ruff_python_ast::source_code::SourceCode;
 use serde::ser::SerializeSeq;
 use serde::{Serialize, Serializer};
 use serde_json::json;
@@ -33,22 +35,20 @@ impl Serialize for ExpandedMessages<'_> {
         let mut s = serializer.serialize_seq(Some(self.messages.len()))?;
 
         for message in self.messages {
+            let source_code = message.file.to_source_code();
+
             let fix = if message.fix.is_empty() {
                 None
             } else {
                 Some(json!({
                     "message": message.kind.suggestion.as_deref(),
-                    "edits": &ExpandedEdits { message },
+                    "edits": &ExpandedEdits { edits: message.fix.edits(), source_code: &source_code },
                 }))
             };
 
-            let start_location = message.compute_start_location();
-            let end_location = message.compute_end_location();
-
-            let noqa_location = message
-                .file
-                .source_code()
-                .source_location(message.noqa_offset);
+            let start_location = source_code.source_location(message.start());
+            let end_location = source_code.source_location(message.end());
+            let noqa_location = source_code.source_location(message.noqa_offset);
 
             let value = json!({
                 "code": message.kind.rule().noqa_code().to_string(),
@@ -68,24 +68,20 @@ impl Serialize for ExpandedMessages<'_> {
 }
 
 struct ExpandedEdits<'a> {
-    message: &'a Message,
+    edits: &'a [Edit],
+    source_code: &'a SourceCode<'a, 'a>,
 }
 
 impl Serialize for ExpandedEdits<'_> {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where
         S: Serializer,
     {
-        let edits = self.message.fix.edits();
-        let mut s = serializer.serialize_seq(Some(edits.len()))?;
-
-        for edit in edits {
-            let start_location = self
-                .message
-                .file
-                .source_code()
-                .source_location(edit.start());
-            let end_location = self.message.file.source_code().source_location(edit.end());
+        let mut s = serializer.serialize_seq(Some(self.edits.len()))?;
+
+        for edit in self.edits {
+            let start_location = self.source_code.source_location(edit.start());
+            let end_location = self.source_code.source_location(edit.end());
             let value = json!({
                 "content": edit.content().unwrap_or_default(),
                 "location": start_location,

crates/ruff/src/message/mod.rs

@@ -57,11 +57,11 @@ impl Message {
     }
 
     pub fn compute_start_location(&self) -> SourceLocation {
-        self.file.source_code().source_location(self.start())
+        self.file.to_source_code().source_location(self.start())
     }
 
     pub fn compute_end_location(&self) -> SourceLocation {
-        self.file.source_code().source_location(self.end())
+        self.file.to_source_code().source_location(self.end())
     }
 
     pub const fn start(&self) -> TextSize {
@@ -178,7 +178,7 @@ def fibonacci(n):
             TextRange::new(TextSize::from(7), TextSize::from(9)),
         );
 
-        let fib_source = SourceFileBuilder::new("fib.py").source_text(fib).finish();
+        let fib_source = SourceFileBuilder::new("fib.py", fib).finish();
 
         let unused_variable = Diagnostic::new(
             UnusedVariable {
@@ -200,9 +200,7 @@ def fibonacci(n):
             TextRange::new(TextSize::from(3), TextSize::from(4)),
         );
 
-        let file_2_source = SourceFileBuilder::new("undef.py")
-            .source_text(file_2)
-            .finish();
+        let file_2_source = SourceFileBuilder::new("undef.py", file_2).finish();
 
         let unused_import_start = unused_import.start();
         let unused_variable_start = unused_variable.start();

crates/ruff/src/message/text.rs

@@ -160,8 +160,6 @@ impl Display for MessageCodeFrame<'_> {
             kind, file, range, ..
         } = self.message;
 
-        let source_code = file.source_code();
-
         let suggestion = kind.suggestion.as_deref();
         let footer = if suggestion.is_some() {
             vec![Annotation {
@@ -173,6 +171,7 @@ impl Display for MessageCodeFrame<'_> {
             Vec::new()
         };
 
+        let source_code = file.to_source_code();
         let content_start_index = source_code.line_index(range.start());
         let mut start_index = content_start_index.saturating_sub(2);
 
@@ -184,12 +183,10 @@ impl Display for MessageCodeFrame<'_> {
             start_index = start_index.saturating_add(1);
         }
 
-        let content_end_index = file.source_code().line_index(range.end());
+        let content_end_index = source_code.line_index(range.end());
         let mut end_index = content_end_index
             .saturating_add(2)
-            .min(OneIndexed::from_zero_indexed(
-                file.source_code().line_count(),
-            ));
+            .min(OneIndexed::from_zero_indexed(source_code.line_count()));
 
         // Trim trailing empty lines
         while end_index > content_end_index {
@@ -203,7 +200,7 @@ impl Display for MessageCodeFrame<'_> {
         let start_offset = source_code.line_start(start_index);
         let end_offset = source_code.line_end(end_index);
 
-        let source_text = &source_code.text()[TextRange::new(start_offset, end_offset)];
+        let source_text = source_code.slice(TextRange::new(start_offset, end_offset));
 
         let annotation_start_offset = range.start() - start_offset;
         let annotation_end_offset = range.end() - start_offset;

crates/ruff/src/rules/pycodestyle/rules/logical_lines/mod.rs

@@ -451,25 +451,22 @@ impl<'a> LogicalLineToken<'a> {
     /// Returns the token's start location
     #[inline]
     pub fn start(&self) -> TextSize {
-        #[allow(unsafe_code)]
-        unsafe {
-            *self.tokens.starts.get_unchecked(self.position)
-        }
+        self.range().start()
     }
 
     /// Returns the token's end location
     #[inline]
     pub fn end(&self) -> TextSize {
-        #[allow(unsafe_code)]
-        unsafe {
-            *self.tokens.ends.get_unchecked(self.position)
-        }
+        self.range().end()
     }
 
     /// Returns a tuple with the token's `(start, end)` locations
     #[inline]
     pub fn range(&self) -> TextRange {
-        TextRange::new(self.start(), self.end())
+        #[allow(unsafe_code)]
+        unsafe {
+            *self.tokens.ranges.get_unchecked(self.position)
+        }
     }
 }
 
@@ -653,20 +650,16 @@ struct Tokens {
     /// The token kinds
     kinds: Vec<TokenKind>,
 
-    /// The start locations
-    starts: Vec<TextSize>,
-
-    /// The end locations
-    ends: Vec<TextSize>,
+    /// The ranges
+    ranges: Vec<TextRange>,
 }
 
 impl Tokens {
     /// Creates new tokens with a reserved size of `capacity`
     fn with_capacity(capacity: usize) -> Self {
         Self {
             kinds: Vec::with_capacity(capacity),
-            starts: Vec::with_capacity(capacity),
-            ends: Vec::with_capacity(capacity),
+            ranges: Vec::with_capacity(capacity),
         }
     }
 
@@ -678,7 +671,6 @@ impl Tokens {
     /// Adds a new token with the given `kind` and `range`
     fn push(&mut self, kind: TokenKind, range: TextRange) {
         self.kinds.push(kind);
-        self.starts.push(range.start());
-        self.ends.push(range.end());
+        self.ranges.push(range);
     }
 }

crates/ruff/src/rules/pycodestyle/rules/logical_lines/whitespace_before_comment.rs

@@ -4,7 +4,7 @@ use ruff_diagnostics::Violation;
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::source_code::Locator;
 use ruff_python_ast::token_kind::TokenKind;
-use ruff_text_size::{TextRange, TextSize};
+use ruff_text_size::{TextLen, TextRange, TextSize};
 
 /// ## What it does
 /// Checks if inline comments are separated by at least two spaces.
@@ -157,12 +157,7 @@ pub(crate) fn whitespace_before_comment(
 
             let is_inline_comment = !line.trim().is_empty();
             if is_inline_comment {
-                if range.start() < prev_end + TextSize::from(2)
-                    && !locator.contains_line_break(TextRange::new(
-                        range.start(),
-                        prev_end + TextSize::from(2),
-                    ))
-                {
+                if range.start() - prev_end < "  ".text_len() {
                     diagnostics.push((
                         TextRange::new(prev_end, range.start()),
                         TooFewSpacesBeforeInlineComment.into(),

crates/ruff/src/test.rs

@@ -94,10 +94,11 @@ pub fn test_path(path: impl AsRef<Path>, settings: &Settings) -> Result<Vec<Mess
                     iterations += 1;
                     contents = fixed_contents.to_string();
                 } else {
-                    let source_code =
-                        SourceFileBuilder::new(&path.file_name().unwrap().to_string_lossy())
-                            .source_text_string(contents)
-                            .finish();
+                    let source_code = SourceFileBuilder::new(
+                        path.file_name().unwrap().to_string_lossy().as_ref(),
+                        contents,
+                    )
+                    .finish();
 
                     let messages: Vec<_> = diagnostics
                         .into_iter()
@@ -132,9 +133,11 @@ pub fn test_path(path: impl AsRef<Path>, settings: &Settings) -> Result<Vec<Mess
         }
     }
 
-    let source_code = SourceFileBuilder::new(&path.file_name().unwrap().to_string_lossy())
-        .source_text_string(contents)
-        .finish();
+    let source_code = SourceFileBuilder::new(
+        path.file_name().unwrap().to_string_lossy().as_ref(),
+        contents,
+    )
+    .finish();
 
     Ok(diagnostics
         .into_iter()

crates/ruff_cli/src/cache.rs

@@ -221,11 +221,7 @@ pub fn get(
 
             let source_files: Vec<_> = sources
                 .into_iter()
-                .map(|(filename, text)| {
-                    SourceFileBuilder::from_string(filename)
-                        .source_text_string(text)
-                        .finish()
-                })
+                .map(|(filename, text)| SourceFileBuilder::new(filename, text).finish())
                 .collect();
 
             for header in headers {

crates/ruff_cli/src/commands/run.rs

@@ -118,9 +118,8 @@ pub fn run(
                     );
                     let settings = resolver.resolve(path, pyproject_strategy);
                     if settings.rules.enabled(Rule::IOError) {
-                        let file = SourceFileBuilder::new(&path.to_string_lossy())
-                            .source_text("")
-                            .finish();
+                        let file =
+                            SourceFileBuilder::new(path.to_string_lossy().as_ref(), "").finish();
 
                         Diagnostics::new(
                             vec![Message::from_diagnostic(

crates/ruff_cli/src/diagnostics.rs

@@ -88,9 +88,7 @@ fn load_jupyter_notebook(path: &Path) -> Result<(String, JupyterIndex), Box<Diag
             return Err(Box::new(Diagnostics {
                 messages: vec![Message::from_diagnostic(
                     *diagnostic,
-                    SourceFileBuilder::new(&path.to_string_lossy())
-                        .source_text("")
-                        .finish(),
+                    SourceFileBuilder::new(path.to_string_lossy().as_ref(), "").finish(),
                     TextSize::default(),
                 )],
                 ..Diagnostics::default()

crates/ruff_python_ast/src/source_code/locator.rs

@@ -37,6 +37,10 @@ impl<'a> Locator<'a> {
             .get_or_init(|| LineIndex::from_source_text(self.contents))
     }
 
+    pub fn line_index(&self) -> Option<&LineIndex> {
+        self.index.get()
+    }
+
     pub fn to_source_code(&self) -> SourceCode {
         SourceCode {
             index: self.to_index(),