perf(logical-lines): Various small perf improvements (#4022)

crates/ruff/src/checkers/logical_lines.rs

@@ -1,7 +1,7 @@
 use ruff_text_size::TextRange;
 use rustpython_parser::lexer::LexResult;
 
-use ruff_diagnostics::{Diagnostic, Fix};
+use ruff_diagnostics::{Diagnostic, DiagnosticKind, Fix};
 use ruff_python_ast::source_code::{Locator, Stylist};
 use ruff_python_ast::token_kind::TokenKind;
 
@@ -37,7 +37,7 @@ pub fn check_logical_lines(
     settings: &Settings,
     autofix: flags::Autofix,
 ) -> Vec<Diagnostic> {
-    let mut diagnostics = vec![];
+    let mut context = LogicalLinesContext::new(settings);
 
     #[cfg(feature = "logical_lines")]
     let should_fix_missing_whitespace =
@@ -59,106 +59,33 @@ pub fn check_logical_lines(
 
     for line in &LogicalLines::from_tokens(tokens, locator) {
         if line.flags().contains(TokenFlags::OPERATOR) {
-            for (location, kind) in space_around_operator(&line) {
-                if settings.rules.enabled(kind.rule()) {
-                    diagnostics.push(Diagnostic {
-                        kind,
-                        range: TextRange::empty(location),
-                        fix: Fix::empty(),
-                        parent: None,
-                    });
-                }
-            }
-
-            for (location, kind) in whitespace_around_named_parameter_equals(&line.tokens()) {
-                if settings.rules.enabled(kind.rule()) {
-                    diagnostics.push(Diagnostic {
-                        kind,
-                        range: TextRange::empty(location),
-                        fix: Fix::empty(),
-                        parent: None,
-                    });
-                }
-            }
-            for (location, kind) in missing_whitespace_around_operator(&line.tokens()) {
-                if settings.rules.enabled(kind.rule()) {
-                    diagnostics.push(Diagnostic {
-                        kind,
-                        range: TextRange::empty(location),
-                        fix: Fix::empty(),
-                        parent: None,
-                    });
-                }
-            }
-
-            for diagnostic in missing_whitespace(&line, should_fix_missing_whitespace) {
-                if settings.rules.enabled(diagnostic.kind.rule()) {
-                    diagnostics.push(diagnostic);
-                }
-            }
+            space_around_operator(&line, &mut context);
+            whitespace_around_named_parameter_equals(&line, &mut context);
+            missing_whitespace_around_operator(&line, &mut context);
+            missing_whitespace(&line, should_fix_missing_whitespace, &mut context);
         }
+
         if line
             .flags()
             .contains(TokenFlags::OPERATOR | TokenFlags::PUNCTUATION)
         {
-            for (location, kind) in extraneous_whitespace(&line) {
-                if settings.rules.enabled(kind.rule()) {
-                    diagnostics.push(Diagnostic {
-                        kind,
-                        range: TextRange::empty(location),
-                        fix: Fix::empty(),
-                        parent: None,
-                    });
-                }
-            }
+            extraneous_whitespace(&line, &mut context);
         }
         if line.flags().contains(TokenFlags::KEYWORD) {
-            for (location, kind) in whitespace_around_keywords(&line) {
-                if settings.rules.enabled(kind.rule()) {
-                    diagnostics.push(Diagnostic {
-                        kind,
-                        range: TextRange::empty(location),
-                        fix: Fix::empty(),
-                        parent: None,
-                    });
-                }
-            }
-
-            for (location, kind) in missing_whitespace_after_keyword(&line.tokens()) {
-                if settings.rules.enabled(kind.rule()) {
-                    diagnostics.push(Diagnostic {
-                        kind,
-                        range: TextRange::empty(location),
-                        fix: Fix::empty(),
-                        parent: None,
-                    });
-                }
-            }
+            whitespace_around_keywords(&line, &mut context);
+            missing_whitespace_after_keyword(&line, &mut context);
         }
+
         if line.flags().contains(TokenFlags::COMMENT) {
-            for (range, kind) in
-                whitespace_before_comment(&line.tokens(), locator, prev_line.is_none())
-            {
-                if settings.rules.enabled(kind.rule()) {
-                    diagnostics.push(Diagnostic {
-                        kind,
-                        range,
-                        fix: Fix::empty(),
-                        parent: None,
-                    });
-                }
-            }
+            whitespace_before_comment(&line, locator, prev_line.is_none(), &mut context);
         }
 
         if line.flags().contains(TokenFlags::BRACKET) {
-            for diagnostic in whitespace_before_parameters(
-                &line.tokens(),
+            whitespace_before_parameters(
+                &line,
                 should_fix_whitespace_before_parameters,
-            ) {
-                if settings.rules.enabled(diagnostic.kind.rule()) {
-                    diagnostics.push(diagnostic);
-                }
-            }
+                &mut context,
+            );
         }
 
         // Extract the indentation level.
@@ -185,12 +112,7 @@ pub fn check_logical_lines(
             indent_size,
         ) {
             if settings.rules.enabled(kind.rule()) {
-                diagnostics.push(Diagnostic {
-                    kind,
-                    range,
-                    fix: Fix::empty(),
-                    parent: None,
-                });
+                context.push(kind, range);
             }
         }
 
@@ -199,7 +121,40 @@ pub fn check_logical_lines(
             prev_indent_level = Some(indent_level);
         }
     }
-    diagnostics
+    context.diagnostics
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct LogicalLinesContext<'a> {
+    settings: &'a Settings,
+    diagnostics: Vec<Diagnostic>,
+}
+
+impl<'a> LogicalLinesContext<'a> {
+    fn new(settings: &'a Settings) -> Self {
+        Self {
+            settings,
+            diagnostics: Vec::new(),
+        }
+    }
+
+    pub fn push<K: Into<DiagnosticKind>>(&mut self, kind: K, range: TextRange) {
+        let kind = kind.into();
+        if self.settings.rules.enabled(kind.rule()) {
+            self.diagnostics.push(Diagnostic {
+                kind,
+                range,
+                fix: Fix::empty(),
+                parent: None,
+            });
+        }
+    }
+
+    pub fn push_diagnostic(&mut self, diagnostic: Diagnostic) {
+        if self.settings.rules.enabled(diagnostic.kind.rule()) {
+            self.diagnostics.push(diagnostic);
+        }
+    }
 }
 
 #[cfg(test)]

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

@@ -1,6 +1,7 @@
-use ruff_text_size::TextSize;
+use ruff_text_size::TextRange;
 
 use super::{LogicalLine, Whitespace};
+use crate::checkers::logical_lines::LogicalLinesContext;
 use ruff_diagnostics::DiagnosticKind;
 use ruff_diagnostics::Violation;
 use ruff_macros::{derive_message_formats, violation};
@@ -101,17 +102,15 @@ impl Violation for WhitespaceBeforePunctuation {
 }
 
 /// E201, E202, E203
-pub(crate) fn extraneous_whitespace(line: &LogicalLine) -> Vec<(TextSize, DiagnosticKind)> {
-    let mut diagnostics = vec![];
-    let mut last_token: Option<TokenKind> = None;
+pub(crate) fn extraneous_whitespace(line: &LogicalLine, context: &mut LogicalLinesContext) {
+    let mut last_token = TokenKind::EndOfFile;
 
     for token in line.tokens() {
         let kind = token.kind();
         match kind {
             TokenKind::Lbrace | TokenKind::Lpar | TokenKind::Lsqb => {
-                if !matches!(line.trailing_whitespace(&token), Whitespace::None) {
-                    let end = token.end();
-                    diagnostics.push((end, WhitespaceAfterOpenBracket.into()));
+                if !matches!(line.trailing_whitespace(token), Whitespace::None) {
+                    context.push(WhitespaceAfterOpenBracket, TextRange::empty(token.end()));
                 }
             }
             TokenKind::Rbrace
@@ -120,28 +119,27 @@ pub(crate) fn extraneous_whitespace(line: &LogicalLine) -> Vec<(TextSize, Diagno
             | TokenKind::Comma
             | TokenKind::Semi
             | TokenKind::Colon => {
-                let diagnostic_kind =
-                    if matches!(kind, TokenKind::Comma | TokenKind::Semi | TokenKind::Colon) {
-                        DiagnosticKind::from(WhitespaceBeforePunctuation)
-                    } else {
-                        DiagnosticKind::from(WhitespaceBeforeCloseBracket)
-                    };
-
                 if let (Whitespace::Single | Whitespace::Many | Whitespace::Tab, offset) =
-                    line.leading_whitespace(&token)
+                    line.leading_whitespace(token)
                 {
-                    if !matches!(last_token, Some(TokenKind::Comma)) {
-                        let start = token.start();
-                        diagnostics.push((start - offset, diagnostic_kind));
+                    if !matches!(last_token, TokenKind::Comma | TokenKind::EndOfFile) {
+                        let diagnostic_kind = if matches!(
+                            kind,
+                            TokenKind::Comma | TokenKind::Semi | TokenKind::Colon
+                        ) {
+                            DiagnosticKind::from(WhitespaceBeforePunctuation)
+                        } else {
+                            DiagnosticKind::from(WhitespaceBeforeCloseBracket)
+                        };
+
+                        context.push(diagnostic_kind, TextRange::empty(token.start() - offset));
                     }
                 }
             }
 
             _ => {}
         }
 
-        last_token = Some(kind);
+        last_token = kind;
     }
-
-    diagnostics
 }

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

@@ -1,9 +1,10 @@
 use super::LogicalLine;
+use crate::checkers::logical_lines::LogicalLinesContext;
 use ruff_diagnostics::Edit;
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::token_kind::TokenKind;
-use ruff_text_size::TextRange;
+use ruff_text_size::{TextRange, TextSize};
 
 #[violation]
 pub struct MissingWhitespace {
@@ -35,30 +36,32 @@ impl AlwaysAutofixableViolation for MissingWhitespace {
 }
 
 /// E231
-pub(crate) fn missing_whitespace(line: &LogicalLine, autofix: bool) -> Vec<Diagnostic> {
-    let mut diagnostics = vec![];
-
+pub(crate) fn missing_whitespace(
+    line: &LogicalLine,
+    autofix: bool,
+    context: &mut LogicalLinesContext,
+) {
     let mut open_parentheses = 0u32;
-    let mut prev_lsqb = None;
-    let mut prev_lbrace = None;
+    let mut prev_lsqb = TextSize::default();
+    let mut prev_lbrace = TextSize::default();
     let mut iter = line.tokens().iter().peekable();
 
     while let Some(token) = iter.next() {
         let kind = token.kind();
         match kind {
             TokenKind::Lsqb => {
                 open_parentheses += 1;
-                prev_lsqb = Some(token.start());
+                prev_lsqb = token.start();
             }
             TokenKind::Rsqb => {
                 open_parentheses += 1;
             }
             TokenKind::Lbrace => {
-                prev_lbrace = Some(token.start());
+                prev_lbrace = token.start();
             }
 
             TokenKind::Comma | TokenKind::Semi | TokenKind::Colon => {
-                let after = line.text_after(&token);
+                let after = line.text_after(token);
 
                 if !after.chars().next().map_or(false, char::is_whitespace) {
                     if let Some(next_token) = iter.peek() {
@@ -85,11 +88,10 @@ pub(crate) fn missing_whitespace(line: &LogicalLine, autofix: bool) -> Vec<Diagn
                     if autofix {
                         diagnostic.set_fix(Edit::insertion(" ".to_string(), token.end()));
                     }
-                    diagnostics.push(diagnostic);
+                    context.push_diagnostic(diagnostic);
                 }
             }
             _ => {}
         }
     }
-    diagnostics
 }

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

@@ -1,13 +1,12 @@
 use itertools::Itertools;
-use ruff_text_size::TextSize;
+use ruff_text_size::TextRange;
 
-use ruff_diagnostics::DiagnosticKind;
+use crate::checkers::logical_lines::LogicalLinesContext;
+use crate::rules::pycodestyle::rules::logical_lines::LogicalLine;
 use ruff_diagnostics::Violation;
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::token_kind::TokenKind;
 
-use super::LogicalLineTokens;
-
 #[violation]
 pub struct MissingWhitespaceAfterKeyword;
 
@@ -20,11 +19,10 @@ impl Violation for MissingWhitespaceAfterKeyword {
 
 /// E275
 pub(crate) fn missing_whitespace_after_keyword(
-    tokens: &LogicalLineTokens,
-) -> Vec<(TextSize, DiagnosticKind)> {
-    let mut diagnostics = vec![];
-
-    for (tok0, tok1) in tokens.iter().tuple_windows() {
+    line: &LogicalLine,
+    context: &mut LogicalLinesContext,
+) {
+    for (tok0, tok1) in line.tokens().iter().tuple_windows() {
         let tok0_kind = tok0.kind();
         let tok1_kind = tok1.kind();
 
@@ -36,8 +34,7 @@ pub(crate) fn missing_whitespace_after_keyword(
                 || matches!(tok1_kind, TokenKind::Colon | TokenKind::Newline))
             && tok0.end() == tok1.start()
         {
-            diagnostics.push((tok0.end(), MissingWhitespaceAfterKeyword.into()));
+            context.push(MissingWhitespaceAfterKeyword, TextRange::empty(tok0.end()));
         }
     }
-    diagnostics
 }