refactor(ruff): Implement doc_lines_from_tokens as iterator (#3124)

This is a nit refactor... It implements the extraction of document lines as an iterator instead of a Vector to avoid the extra allocation.

crates/ruff/src/doc_lines.rs

@@ -3,31 +3,57 @@
 
 use rustpython_parser::ast::{Constant, ExprKind, Stmt, StmtKind, Suite};
 use rustpython_parser::lexer::{LexResult, Tok};
+use std::iter::FusedIterator;
 
 use crate::ast::visitor;
 use crate::ast::visitor::Visitor;
 
 /// Extract doc lines (standalone comments) from a token sequence.
-pub fn doc_lines_from_tokens(lxr: &[LexResult]) -> Vec<usize> {
+pub fn doc_lines_from_tokens(lxr: &[LexResult]) -> DocLines {
-    let mut doc_lines: Vec<usize> = Vec::default();
+    DocLines::new(lxr)
-    let mut prev: Option<usize> = None;
-    for (start, tok, end) in lxr.iter().flatten() {
-        if matches!(tok, Tok::Indent | Tok::Dedent | Tok::Newline) {
-            continue;
 }
-        if matches!(tok, Tok::Comment(..)) {
+
-            if let Some(prev) = prev {
+pub struct DocLines<'a> {
+    inner: std::iter::Flatten<core::slice::Iter<'a, LexResult>>,
+    prev: Option<usize>,
+}
+
+impl<'a> DocLines<'a> {
+    fn new(lxr: &'a [LexResult]) -> Self {
+        Self {
+            inner: lxr.iter().flatten(),
+            prev: None,
+        }
+    }
+}
+
+impl Iterator for DocLines<'_> {
+    type Item = usize;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            let (start, tok, end) = self.inner.next()?;
+
+            match tok {
+                Tok::Indent | Tok::Dedent | Tok::Newline => continue,
+                Tok::Comment(..) => {
+                    if let Some(prev) = self.prev {
                         if start.row() > prev {
-                    doc_lines.push(start.row());
+                            break Some(start.row());
                         }
                     } else {
-                doc_lines.push(start.row());
+                        break Some(start.row());
                     }
                 }
-        prev = Some(end.row());
+                _ => {}
             }
-    doc_lines
+
+            self.prev = Some(end.row());
         }
+    }
+}
+
+impl FusedIterator for DocLines<'_> {}
 
 #[derive(Default)]
 struct StringLinesVisitor {

crates/ruff/src/fs.rs

@@ -1,5 +1,3 @@
-use std::fs::File;
-use std::io::{BufReader, Read};
 use std::ops::Deref;
 use std::path::{Path, PathBuf};
 
@@ -88,12 +86,3 @@ pub fn relativize_path(path: impl AsRef<Path>) -> String {
     }
     format!("{}", path.display())
 }
-
-/// Read a file's contents from disk.
-pub fn read_file<P: AsRef<Path>>(path: P) -> Result<String> {
-    let file = File::open(path)?;
-    let mut buf_reader = BufReader::new(file);
-    let mut contents = String::new();
-    buf_reader.read_to_string(&mut contents)?;
-    Ok(contents)
-}

crates/ruff/src/linter.rs

@@ -223,7 +223,7 @@ const MAX_ITERATIONS: usize = 100;
 /// Add any missing `# noqa` pragmas to the source code at the given `Path`.
 pub fn add_noqa_to_path(path: &Path, package: Option<&Path>, settings: &Settings) -> Result<usize> {
     // Read the file from disk.
-    let contents = fs::read_file(path)?;
+    let contents = std::fs::read_to_string(path)?;
 
     // Tokenize once.
     let tokens: Vec<LexResult> = rustpython_helpers::tokenize(&contents);

crates/ruff/src/settings/pyproject.rs

@@ -5,7 +5,6 @@ use std::path::{Path, PathBuf};
 use anyhow::{anyhow, Result};
 use serde::{Deserialize, Serialize};
 
-use crate::fs;
 use crate::settings::options::Options;
 
 #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
@@ -30,13 +29,13 @@ impl Pyproject {
 
 /// Parse a `ruff.toml` file.
 fn parse_ruff_toml<P: AsRef<Path>>(path: P) -> Result<Options> {
-    let contents = fs::read_file(path)?;
+    let contents = std::fs::read_to_string(path)?;
     toml::from_str(&contents).map_err(Into::into)
 }
 
 /// Parse a `pyproject.toml` file.
 fn parse_pyproject_toml<P: AsRef<Path>>(path: P) -> Result<Pyproject> {
-    let contents = fs::read_file(path)?;
+    let contents = std::fs::read_to_string(path)?;
     toml::from_str(&contents).map_err(Into::into)
 }
 

crates/ruff/src/test.rs

@@ -12,7 +12,7 @@ use crate::packaging::detect_package_root;
 use crate::registry::Diagnostic;
 use crate::settings::{flags, Settings};
 use crate::source_code::{Indexer, Locator, Stylist};
-use crate::{directives, fs, rustpython_helpers};
+use crate::{directives, rustpython_helpers};
 
 pub fn test_resource_path(path: impl AsRef<Path>) -> std::path::PathBuf {
     Path::new("./resources/test/").join(path)
@@ -22,7 +22,7 @@ pub fn test_resource_path(path: impl AsRef<Path>) -> std::path::PathBuf {
 /// asserts that autofixes converge after 10 iterations.
 pub fn test_path(path: &Path, settings: &Settings) -> Result<Vec<Diagnostic>> {
     let path = test_resource_path("fixtures").join(path);
-    let contents = fs::read_file(&path)?;
+    let contents = std::fs::read_to_string(&path)?;
     let tokens: Vec<LexResult> = rustpython_helpers::tokenize(&contents);
     let locator = Locator::new(&contents);
     let stylist = Stylist::from_contents(&contents, &locator);

crates/ruff_cli/src/diagnostics.rs

@@ -81,7 +81,7 @@ pub fn lint_path(
     };
 
     // Read the file from disk.
-    let contents = fs::read_file(path)?;
+    let contents = std::fs::read_to_string(path)?;
 
     // Lint the file.
     let (