Refactor and rename skip_trailing_trivia (#6312)

Based on feedback here:
https://github.com/astral-sh/ruff/pull/6274#discussion_r1282747964.

crates/ruff_python_formatter/src/comments/format.rs

@@ -3,7 +3,7 @@ use ruff_text_size::{TextLen, TextRange, TextSize};
 
 use ruff_formatter::{format_args, write, FormatError, SourceCode};
 use ruff_python_ast::node::{AnyNodeRef, AstNode};
-use ruff_python_trivia::{lines_after, lines_before, skip_trailing_trivia};
+use ruff_python_trivia::{lines_after, lines_after_ignoring_trivia, lines_before};
 
 use crate::comments::SourceComment;
 use crate::context::NodeLevel;
@@ -90,10 +90,9 @@ impl Format<PyFormatContext<'_>> for FormatLeadingAlternateBranchComments<'_> {
 
             write!(f, [leading_comments(self.comments)])?;
         } else if let Some(last_preceding) = self.last_node {
-            let full_end = skip_trailing_trivia(last_preceding.end(), f.context().source());
             // The leading comments formatting ensures that it preserves the right amount of lines after
             // We need to take care of this ourselves, if there's no leading `else` comment.
-            if lines_after(full_end, f.context().source()) > 1 {
+            if lines_after_ignoring_trivia(last_preceding.end(), f.context().source()) > 1 {
                 write!(f, [empty_line()])?;
             }
         }

crates/ruff_python_formatter/src/statement/stmt_class_def.rs

@@ -1,6 +1,6 @@
 use ruff_formatter::{write, Buffer};
 use ruff_python_ast::{Ranged, StmtClassDef};
-use ruff_python_trivia::{lines_after, skip_trailing_trivia};
+use ruff_python_trivia::lines_after_ignoring_trivia;
 
 use crate::comments::{leading_comments, trailing_comments};
 use crate::prelude::*;
@@ -41,14 +41,13 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
                 // Write any leading definition comments (between last decorator and the header)
                 // while maintaining the right amount of empty lines between the comment
                 // and the last decorator.
-                let decorator_end =
+                let leading_line =
-                    skip_trailing_trivia(last_decorator.end(), f.context().source());
+                    if lines_after_ignoring_trivia(last_decorator.end(), f.context().source()) <= 1
-
+                    {
-                let leading_line = if lines_after(decorator_end, f.context().source()) <= 1 {
+                        hard_line_break()
-                    hard_line_break()
+                    } else {
-                } else {
+                        empty_line()
-                    empty_line()
+                    };
-                };
 
                 write!(
                     f,

crates/ruff_python_formatter/src/statement/stmt_function_def.rs

@@ -2,7 +2,7 @@ use ruff_python_ast::{Ranged, StmtFunctionDef};
 
 use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule};
 use ruff_python_ast::function::AnyFunctionDefinition;
-use ruff_python_trivia::{lines_after, skip_trailing_trivia};
+use ruff_python_trivia::lines_after_ignoring_trivia;
 
 use crate::comments::{leading_comments, trailing_comments};
 
@@ -54,14 +54,13 @@ impl FormatRule<AnyFunctionDefinition<'_>, PyFormatContext<'_>> for FormatAnyFun
                 // Write any leading definition comments (between last decorator and the header)
                 // while maintaining the right amount of empty lines between the comment
                 // and the last decorator.
-                let decorator_end =
+                let leading_line =
-                    skip_trailing_trivia(last_decorator.end(), f.context().source());
+                    if lines_after_ignoring_trivia(last_decorator.end(), f.context().source()) <= 1
-
+                    {
-                let leading_line = if lines_after(decorator_end, f.context().source()) <= 1 {
+                        hard_line_break()
-                    hard_line_break()
+                    } else {
-                } else {
+                        empty_line()
-                    empty_line()
+                    };
-                };
 
                 write!(
                     f,

crates/ruff_python_formatter/src/statement/suite.rs

@@ -1,7 +1,7 @@
 use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
 use ruff_python_ast::helpers::is_compound_statement;
 use ruff_python_ast::{self as ast, Constant, Expr, Ranged, Stmt, Suite};
-use ruff_python_trivia::{lines_after, lines_before, skip_trailing_trivia};
+use ruff_python_trivia::{lines_after_ignoring_trivia, lines_before};
 
 use crate::context::{NodeLevel, WithNodeLevel};
 use crate::prelude::*;
@@ -180,8 +180,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
                     // it then counts the lines between the statement and the trailing comment, which is
                     // always 0. This is why it skips any trailing trivia (trivia that's on the same line)
                     // and counts the lines after.
-                    let after_trailing_trivia = skip_trailing_trivia(offset, source);
+                    lines_after_ignoring_trivia(offset, source)
-                    lines_after(after_trailing_trivia, source)
                 };
 
                 match node_level {

crates/ruff_python_trivia/src/tokenizer.rs

@@ -70,24 +70,18 @@ pub fn lines_after(offset: TextSize, code: &str) -> u32 {
     newlines
 }
 
-/// Returns the position after skipping any trailing trivia up to, but not including the newline character.
+/// Counts the empty lines after `offset`, ignoring any trailing trivia on the same line as
-pub fn skip_trailing_trivia(offset: TextSize, code: &str) -> TextSize {
+/// `offset`.
-    let tokenizer = SimpleTokenizer::starts_at(offset, code);
+#[allow(clippy::cast_possible_truncation)]
-
+pub fn lines_after_ignoring_trivia(offset: TextSize, code: &str) -> u32 {
-    for token in tokenizer {
+    // SAFETY: We don't support files greater than 4GB, so casting to u32 is safe.
-        match token.kind() {
+    SimpleTokenizer::starts_at(offset, code)
-            SimpleTokenKind::Whitespace
+        .skip_while(|token| token.kind != SimpleTokenKind::Newline && token.kind.is_trivia())
-            | SimpleTokenKind::Comment
+        .take_while(|token| {
-            | SimpleTokenKind::Continuation => {
+            token.kind == SimpleTokenKind::Newline || token.kind == SimpleTokenKind::Whitespace
-                // No op
+        })
-            }
+        .filter(|token| token.kind == SimpleTokenKind::Newline)
-            _ => {
+        .count() as u32
-                return token.start();
-            }
-        }
-    }
-
-    offset
 }
 
 fn is_identifier_start(c: char) -> bool {