From a1e14771370b96e726dedbec9e3ead816ba5c485 Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Tue, 6 Feb 2024 17:29:54 -0500 Subject: [PATCH] Try to improve trailing comma perf --- .../flake8_commas/rules/trailing_commas.rs | 105 +++++++++--------- 1 file changed, 52 insertions(+), 53 deletions(-) diff --git a/crates/ruff_linter/src/rules/flake8_commas/rules/trailing_commas.rs b/crates/ruff_linter/src/rules/flake8_commas/rules/trailing_commas.rs index f0dc92b6001bb2..61d9dae62da02c 100644 --- a/crates/ruff_linter/src/rules/flake8_commas/rules/trailing_commas.rs +++ b/crates/ruff_linter/src/rules/flake8_commas/rules/trailing_commas.rs @@ -1,10 +1,8 @@ -use itertools::Itertools; - use ruff_diagnostics::{AlwaysFixableViolation, Violation}; use ruff_diagnostics::{Diagnostic, Edit, Fix}; use ruff_macros::{derive_message_formats, violation}; use ruff_python_index::Indexer; -use ruff_python_parser::lexer::{LexResult, Spanned}; +use ruff_python_parser::lexer::LexResult; use ruff_python_parser::Tok; use ruff_source_file::Locator; use ruff_text_size::{Ranged, TextRange}; @@ -12,20 +10,20 @@ use ruff_text_size::{Ranged, TextRange}; /// Simplified token type. #[derive(Copy, Clone, PartialEq, Eq)] enum TokenType { - Irrelevant, - NonLogicalNewline, + Named, + String, Newline, - Comma, + NonLogicalNewline, OpeningBracket, + ClosingBracket, OpeningSquareBracket, + Colon, + Comma, OpeningCurlyBracket, - ClosingBracket, - For, - Named, Def, + For, Lambda, - Colon, - String, + Irrelevant, } /// Simplified token specialized for the task. @@ -54,30 +52,29 @@ impl Token { } } -impl From<&Spanned> for Token { - fn from(spanned: &Spanned) -> Self { - let r#type = match &spanned.0 { - Tok::NonLogicalNewline => TokenType::NonLogicalNewline, - Tok::Newline => TokenType::Newline, - Tok::For => TokenType::For, - Tok::Def => TokenType::Def, - Tok::Lambda => TokenType::Lambda, - // Import treated like a function. - Tok::Import => TokenType::Named, +impl From<(&Tok, TextRange)> for Token { + fn from((tok, range): (&Tok, TextRange)) -> Self { + let r#type = match tok { Tok::Name { .. } => TokenType::Named, Tok::String { .. } => TokenType::String, - Tok::Comma => TokenType::Comma, + Tok::Newline => TokenType::Newline, + Tok::NonLogicalNewline => TokenType::NonLogicalNewline, Tok::Lpar => TokenType::OpeningBracket, + Tok::Rpar => TokenType::ClosingBracket, Tok::Lsqb => TokenType::OpeningSquareBracket, - Tok::Lbrace => TokenType::OpeningCurlyBracket, - Tok::Rpar | Tok::Rsqb | Tok::Rbrace => TokenType::ClosingBracket, + Tok::Rsqb => TokenType::ClosingBracket, Tok::Colon => TokenType::Colon, + Tok::Comma => TokenType::Comma, + Tok::Lbrace => TokenType::OpeningCurlyBracket, + Tok::Rbrace => TokenType::ClosingBracket, + Tok::Def => TokenType::Def, + Tok::For => TokenType::For, + Tok::Lambda => TokenType::Lambda, + // Import treated like a function. + Tok::Import => TokenType::Named, _ => TokenType::Irrelevant, }; - Self { - range: spanned.1, - r#type, - } + Self { range, r#type } } } @@ -237,10 +234,12 @@ pub(crate) fn trailing_commas( indexer: &Indexer, ) { let mut fstrings = 0u32; - let tokens = tokens - .iter() - .flatten() - .filter_map(|spanned @ (tok, tok_range)| match tok { + let mut tokens = tokens.iter().filter_map(|result| { + let Ok((tok, tok_range)) = result else { + return None; + }; + + match tok { // Completely ignore comments -- they just interfere with the logic. Tok::Comment(_) => None, // F-strings are handled as `String` token type with the complete range @@ -263,31 +262,28 @@ pub(crate) fn trailing_commas( } _ => { if fstrings == 0 { - Some(Token::from(spanned)) + Some(Token::from((tok, *tok_range))) } else { None } } - }); - let tokens = [Token::irrelevant(), Token::irrelevant()] - .into_iter() - .chain(tokens); - // Collapse consecutive newlines to the first one -- trailing commas are - // added before the first newline. - let tokens = tokens.coalesce(|previous, current| { - if previous.r#type == TokenType::NonLogicalNewline - && current.r#type == TokenType::NonLogicalNewline - { - Ok(previous) - } else { - Err((previous, current)) } }); - // The current nesting of the comma contexts. + let mut prev = Token::irrelevant(); + let mut prev_prev = Token::irrelevant(); + let mut stack = vec![Context::new(ContextType::No)]; - for (prev_prev, prev, token) in tokens.tuple_windows() { + while let Some(token) = tokens.next() { + if prev.r#type == TokenType::NonLogicalNewline + && token.r#type == TokenType::NonLogicalNewline + { + // Collapse consecutive newlines to the first one -- trailing commas are + // added before the first newline. + continue; + } + // Update the comma context stack. match token.r#type { TokenType::OpeningBracket => match (prev.r#type, prev_prev.r#type) { @@ -316,16 +312,16 @@ pub(crate) fn trailing_commas( stack.push(Context::new(ContextType::LambdaParameters)); } TokenType::For => { - let len = stack.len(); - stack[len - 1] = Context::new(ContextType::No); + let last = stack.last_mut().expect("Stack to never be empty"); + *last = Context::new(ContextType::No); } TokenType::Comma => { - let len = stack.len(); - stack[len - 1].inc(); + let last = stack.last_mut().expect("Stack to never be empty"); + last.inc(); } _ => {} } - let context = &stack[stack.len() - 1]; + let context = stack.last().expect("Stack to never be empty"); // Is it allowed to have a trailing comma before this token? let comma_allowed = token.r#type == TokenType::ClosingBracket @@ -412,5 +408,8 @@ pub(crate) fn trailing_commas( if pop_context && stack.len() > 1 { stack.pop(); } + + prev_prev = prev; + prev = token; } }