Simple Red-Knot benchmark

Cargo.toml

@@ -35,6 +35,7 @@ ruff_source_file = { path = "crates/ruff_source_file" }
 ruff_text_size = { path = "crates/ruff_text_size" }
 ruff_workspace = { path = "crates/ruff_workspace" }
 
+red_knot = { path = "crates/red_knot" }
 red_knot_module_resolver = { path = "crates/red_knot_module_resolver" }
 red_knot_python_semantic = { path = "crates/red_knot_python_semantic" }
 

crates/red_knot/src/program/check.rs

@@ -11,15 +11,19 @@ impl Program {
         self.with_db(|db| {
             let mut result = Vec::new();
             for open_file in db.workspace.open_files() {
-                result.extend_from_slice(&db.check_file(open_file));
+                result.extend_from_slice(&db.check_file_impl(open_file));
             }
 
             result
         })
     }
 
     #[tracing::instrument(level = "debug", skip(self))]
-    fn check_file(&self, file: VfsFile) -> Diagnostics {
+    pub fn check_file(&self, file: VfsFile) -> Result<Diagnostics, Cancelled> {
+        self.with_db(|db| db.check_file_impl(file))
+    }
+
+    fn check_file_impl(&self, file: VfsFile) -> Diagnostics {
         let mut diagnostics = Vec::new();
         diagnostics.extend_from_slice(lint_syntax(self, file));
         diagnostics.extend_from_slice(lint_semantic(self, file));

crates/ruff_benchmark/Cargo.toml

@@ -31,6 +31,10 @@ harness = false
 name = "formatter"
 harness = false
 
+[[bench]]
+name = "red_knot"
+harness = false
+
 [dependencies]
 once_cell = { workspace = true }
 serde = { workspace = true }
@@ -41,11 +45,14 @@ criterion = { workspace = true, default-features = false }
 codspeed-criterion-compat = { workspace = true, default-features = false, optional = true }
 
 [dev-dependencies]
+ruff_db = { workspace = true }
 ruff_linter = { workspace = true }
 ruff_python_ast = { workspace = true }
 ruff_python_formatter = { workspace = true }
 ruff_python_parser = { workspace = true }
 ruff_python_trivia = { workspace = true }
+red_knot = { workspace = true }
+red_knot_module_resolver = { workspace = true }
 
 [lints]
 workspace = true

crates/ruff_benchmark/benches/red_knot.rs

@@ -0,0 +1,143 @@
+#![allow(clippy::disallowed_names)]
+
+use criterion::BenchmarkId;
+use red_knot::program::Program;
+use red_knot::Workspace;
+use red_knot_module_resolver::{set_module_resolution_settings, ModuleResolutionSettings};
+use ruff_benchmark::criterion::{
+    criterion_group, criterion_main, BatchSize, Criterion, Throughput,
+};
+use ruff_db::file_system::{FileSystemPath, MemoryFileSystem};
+use ruff_db::parsed::parsed_module;
+use ruff_db::vfs::{system_path_to_file, VfsFile};
+use ruff_db::Upcast;
+
+static FOO_CODE: &str = r#"
+import typing
+
+from bar import Bar
+
+class Foo(Bar):
+    def foo() -> str:
+        return "foo"
+
+    @typing.override
+    def bar() -> str:
+        return "foo_bar"
+"#;
+
+static BAR_CODE: &str = r#"
+class Bar:
+    def bar() -> str:
+        return "bar"
+
+    def random(arg: int) -> int:
+        if arg == 1:
+            return 48472783
+        if arg < 10:
+            return 20
+        return 36673
+"#;
+
+static TYPING_CODE: &str = r#"
+def override(): ...
+"#;
+
+struct Case {
+    program: Program,
+    foo: VfsFile,
+    bar: VfsFile,
+    typing: VfsFile,
+}
+
+fn setup_case() -> Case {
+    let fs = MemoryFileSystem::new();
+    let foo_path = FileSystemPath::new("/src/foo.py");
+    let bar_path = FileSystemPath::new("/src/bar.py");
+    let typing_path = FileSystemPath::new("/src/typing.pyi");
+    fs.write_files([
+        (foo_path, FOO_CODE),
+        (bar_path, BAR_CODE),
+        (typing_path, TYPING_CODE),
+    ])
+    .unwrap();
+
+    let workspace_root = FileSystemPath::new("/src");
+    let workspace = Workspace::new(workspace_root.to_path_buf());
+
+    let mut program = Program::new(workspace, fs);
+    let foo = system_path_to_file(&program, foo_path).unwrap();
+
+    set_module_resolution_settings(
+        &mut program,
+        ModuleResolutionSettings {
+            extra_paths: vec![],
+            workspace_root: workspace_root.to_path_buf(),
+            site_packages: None,
+            custom_typeshed: None,
+        },
+    );
+
+    program.workspace_mut().open_file(foo);
+
+    let bar = system_path_to_file(&program, bar_path).unwrap();
+    let typing = system_path_to_file(&program, typing_path).unwrap();
+
+    Case {
+        program,
+        foo,
+        bar,
+        typing,
+    }
+}
+
+fn benchmark_warm(criterion: &mut Criterion) {
+    let mut group = criterion.benchmark_group("red_knot/check_file");
+    group.throughput(Throughput::Bytes(FOO_CODE.len() as u64));
+
+    group.bench_function(BenchmarkId::from_parameter("warm"), |b| {
+        b.iter_batched(
+            || {
+                let case = setup_case();
+                // Pre-parse the module to only measure the semantic time.
+                parsed_module(case.program.upcast(), case.foo);
+                parsed_module(case.program.upcast(), case.bar);
+                parsed_module(case.program.upcast(), case.typing);
+                case
+            },
+            |case| {
+                let Case { program, foo, .. } = case;
+                let result = program.check_file(foo).unwrap();
+
+                assert_eq!(result.as_slice(), [] as [String; 0]);
+            },
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.finish();
+}
+
+fn benchmark_cold(criterion: &mut Criterion) {
+    let mut group = criterion.benchmark_group("red_knot/check_file");
+    group.throughput(Throughput::Bytes(FOO_CODE.len() as u64));
+
+    group.bench_function(BenchmarkId::from_parameter("cold"), |b| {
+        b.iter_batched(
+            setup_case,
+            |case| {
+                let Case { program, foo, .. } = case;
+                let result = program.check_file(foo).unwrap();
+
+                assert_eq!(result.as_slice(), [] as [String; 0]);
+            },
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.finish();
+}
+
+criterion_group!(cold, benchmark_warm);
+criterion_group!(warm, benchmark_cold);
+criterion_main!(warm, cold);