From 51ae47ad56a676646d0a6032b624753bfa5a8b55 Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Mon, 14 Aug 2023 10:25:37 +0200
Subject: [PATCH] Remove lex and parsing from formatter benchmark (#6547)

---
 Cargo.lock                                 |  1 +
 crates/ruff_benchmark/Cargo.toml           |  1 +
 crates/ruff_benchmark/benches/formatter.rs | 38 ++++++++++++++++------
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index ed93d2eb4d..8043bb4875 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2141,6 +2141,7 @@ dependencies = [
  "ruff",
  "ruff_python_ast",
  "ruff_python_formatter",
+ "ruff_python_index",
  "ruff_python_parser",
  "serde",
  "serde_json",
diff --git a/crates/ruff_benchmark/Cargo.toml b/crates/ruff_benchmark/Cargo.toml
index 734471e1d2..f55927fd06 100644
--- a/crates/ruff_benchmark/Cargo.toml
+++ b/crates/ruff_benchmark/Cargo.toml
@@ -37,6 +37,7 @@ ureq = "2.6.2"
 ruff.path = "../ruff"
 ruff_python_ast.path = "../ruff_python_ast"
 ruff_python_formatter = { path = "../ruff_python_formatter" }
+ruff_python_index = { path = "../ruff_python_index" }
 ruff_python_parser = { path = "../ruff_python_parser" }
 criterion = { version = "0.5.1"}
 
diff --git a/crates/ruff_benchmark/benches/formatter.rs b/crates/ruff_benchmark/benches/formatter.rs
index 6895bd2ac1..efcc986e72 100644
--- a/crates/ruff_benchmark/benches/formatter.rs
+++ b/crates/ruff_benchmark/benches/formatter.rs
@@ -1,8 +1,12 @@
-use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
-use ruff_benchmark::{TestCase, TestCaseSpeed, TestFile, TestFileDownloadError};
-use ruff_python_formatter::{format_module, PyFormatOptions};
 use std::path::Path;
-use std::time::Duration;
+
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
+
+use ruff_benchmark::{TestCase, TestFile, TestFileDownloadError};
+use ruff_python_formatter::{format_node, PyFormatOptions};
+use ruff_python_index::CommentRangesBuilder;
+use ruff_python_parser::lexer::lex;
+use ruff_python_parser::{parse_tokens, Mode};
 
 #[cfg(target_os = "windows")]
 #[global_allocator]
@@ -41,19 +45,33 @@ fn benchmark_formatter(criterion: &mut Criterion) {
 
     for case in test_cases {
         group.throughput(Throughput::Bytes(case.code().len() as u64));
-        group.measurement_time(match case.speed() {
-            TestCaseSpeed::Fast => Duration::from_secs(5),
-            TestCaseSpeed::Normal => Duration::from_secs(10),
-            TestCaseSpeed::Slow => Duration::from_secs(20),
-        });
 
         group.bench_with_input(
             BenchmarkId::from_parameter(case.name()),
             &case,
             |b, case| {
+                let mut tokens = Vec::new();
+                let mut comment_ranges = CommentRangesBuilder::default();
+
+                for result in lex(case.code(), Mode::Module) {
+                    let (token, range) = result.expect("Input to be a valid python program.");
+
+                    comment_ranges.visit_token(&token, range);
+                    tokens.push(Ok((token, range)));
+                }
+
+                let comment_ranges = comment_ranges.finish();
+
+                // Parse the AST.
+                let python_ast = parse_tokens(tokens, Mode::Module, "<filename>")
+                    .expect("Input to be a valid python program");
+
                 b.iter(|| {
                     let options = PyFormatOptions::from_extension(Path::new(case.name()));
-                    format_module(case.code(), options).expect("Formatting to succeed")
+                    let formatted = format_node(&python_ast, &comment_ranges, case.code(), options)
+                        .expect("Formatting to succeed");
+
+                    formatted.print().expect("Printing to succeed")
                 });
             },
         );