[ty] Implement Duboc's TDD optimization for unions of constraint sets

crates/ruff_benchmark/benches/ty.rs

@@ -1,6 +1,9 @@
 #![allow(clippy::disallowed_names)]
 use ruff_benchmark::criterion;
-use ruff_benchmark::real_world_projects::{InstalledProject, RealWorldProject};
+use ruff_benchmark::real_world_projects::{
+    InstalledProject, RealWorldProject, copy_directory_recursive, get_project_cache_dir,
+    install_dependencies_to_cache,
+};
 
 use std::fmt::Write;
 use std::ops::Range;
@@ -219,9 +222,40 @@ fn assert_diagnostics(db: &dyn Db, diagnostics: &[Diagnostic], expected: &[KeyDi
 }
 
 fn setup_micro_case(code: &str) -> Case {
+    setup_micro_case_inner(code, None)
+}
+
+fn setup_micro_case_with_dependencies(name: &str, dependencies: &[&str], code: &str) -> Case {
+    setup_micro_case_inner(code, Some((name, dependencies)))
+}
+
+fn setup_micro_case_inner(code: &str, dependencies: Option<(&str, &[&str])>) -> Case {
     let system = TestSystem::default();
     let fs = system.memory_file_system().clone();
 
+    let python = dependencies.map(|(name, dependencies)| {
+        let cache_dir = get_project_cache_dir(name).expect("Failed to get cache directory");
+        std::fs::create_dir_all(&cache_dir).expect("Failed to create cache directory");
+
+        let venv_path = cache_dir.join(".venv");
+        install_dependencies_to_cache(
+            name,
+            dependencies,
+            &venv_path,
+            PythonVersion::PY312,
+            "2025-06-17",
+        )
+        .expect("Failed to install dependencies");
+
+        // Copy the on-disk venv into the in-memory filesystem.
+        // ProjectMetadata::discover walks up from /src and uses / as the project root,
+        // so the venv must be at /.venv for the `python = ".venv"` option to resolve correctly.
+        copy_directory_recursive(&fs, &venv_path, SystemPath::new("/.venv"))
+            .expect("Failed to copy venv to memory filesystem");
+
+        RelativePathBuf::cli(SystemPath::new(".venv"))
+    });
+
     let file_path = "src/test.py";
     fs.write_file_all(
         SystemPathBuf::from(file_path),
@@ -234,6 +268,7 @@ fn setup_micro_case(code: &str) -> Case {
     metadata.apply_options(Options {
         environment: Some(EnvironmentOptions {
             python_version: Some(RangedValue::cli(PythonVersion::PY312)),
+            python,
             ..EnvironmentOptions::default()
         }),
         ..Options::default()
@@ -777,6 +812,39 @@ fn benchmark_large_isinstance_narrowing(criterion: &mut Criterion) {
     });
 }
 
+fn benchmark_pandas_tdd(criterion: &mut Criterion) {
+    setup_rayon();
+
+    // This example was reported in https://github.com/astral-sh/ty/issues/3039.
+    criterion.bench_function("ty_micro[pandas_tdd]", |b| {
+        b.iter_batched_ref(
+            || {
+                setup_micro_case_with_dependencies(
+                    "pandas_tdd",
+                    &["pandas-stubs"],
+                    r#"
+                    import pandas as pd
+
+                    df = pd.DataFrame({
+                        "a": [1, 2, 3],
+                        "b": [4, 5, 6],
+                        "c": [7, 8, 9],
+                    })
+                    df["d"] = df["a"] + df["b"] + df["c"] + 1 + (
+                        df["a"] ** 2 + df["b"] ** 2 + df["c"] ** 2)
+                    "#,
+                )
+            },
+            |case| {
+                let Case { db, .. } = case;
+                let result = db.check();
+                assert_eq!(result.len(), 0);
+            },
+            BatchSize::SmallInput,
+        );
+    });
+}
+
 struct ProjectBenchmark<'a> {
     project: InstalledProject<'a>,
     fs: MemoryFileSystem,
@@ -941,6 +1009,7 @@ criterion_group!(
     benchmark_very_large_tuple,
     benchmark_large_union_narrowing,
     benchmark_large_isinstance_narrowing,
+    benchmark_pandas_tdd,
 );
 criterion_group!(project, anyio, attrs, hydra, datetype);
 criterion_main!(check_file, micro, project);

crates/ruff_benchmark/src/real_world_projects.rs

@@ -146,7 +146,7 @@ impl<'a> InstalledProject<'a> {
 }
 
 /// Get the cache directory for a project in the cargo target directory
-fn get_project_cache_dir(project_name: &str) -> Result<std::path::PathBuf> {
+pub fn get_project_cache_dir(project_name: &str) -> Result<std::path::PathBuf> {
     let target_dir = cargo_target_directory()
         .cloned()
         .unwrap_or_else(|| PathBuf::from("target"));
@@ -250,8 +250,18 @@ fn clone_repository(repo_url: &str, target_dir: &Path, commit: &str) -> Result<(
     Ok(())
 }
 
-/// Install dependencies using uv with date constraints
-fn install_dependencies(checkout: &Checkout) -> Result<()> {
+/// Install dependencies into a cached virtual environment.
+///
+/// Creates a venv under `target/benchmark_cache/<name>/.venv` and installs the given
+/// dependencies using `uv pip install` with `--exclude-newer` for reproducibility. The venv is
+/// reused across benchmark runs. Returns the path to the venv directory.
+pub fn install_dependencies_to_cache(
+    name: &str,
+    dependencies: &[&str],
+    venv_path: &PathBuf,
+    python_version: PythonVersion,
+    max_dep_date: &str,
+) -> Result<()> {
     // Check if uv is available
     let uv_check = Command::new("uv")
         .arg("--version")
@@ -264,12 +274,11 @@ fn install_dependencies(checkout: &Checkout) -> Result<()> {
         );
     }
 
-    let venv_path = checkout.venv_path();
-    let python_version_str = checkout.project().python_version.to_string();
+    let python_version_str = python_version.to_string();
 
     let output = Command::new("uv")
         .args(["venv", "--python", &python_version_str, "--allow-existing"])
-        .arg(&venv_path)
+        .arg(venv_path)
         .output()
         .context("Failed to execute uv venv command")?;
 
@@ -279,11 +288,8 @@ fn install_dependencies(checkout: &Checkout) -> Result<()> {
         String::from_utf8_lossy(&output.stderr)
     );
 
-    if checkout.project().dependencies.is_empty() {
-        tracing::debug!(
-            "No dependencies to install for project '{}'",
-            checkout.project().name
-        );
+    if dependencies.is_empty() {
+        tracing::debug!("No dependencies to install for project '{}'", name);
         return Ok(());
     }
 
@@ -295,9 +301,9 @@ fn install_dependencies(checkout: &Checkout) -> Result<()> {
         "--python",
         venv_path.to_str().unwrap(),
         "--exclude-newer",
-        checkout.project().max_dep_date,
+        max_dep_date,
     ])
-    .args(checkout.project().dependencies);
+    .args(dependencies);
 
     let output = cmd
         .output()
@@ -312,8 +318,21 @@ fn install_dependencies(checkout: &Checkout) -> Result<()> {
     Ok(())
 }
 
+/// Install dependencies using uv with date constraints
+fn install_dependencies(checkout: &Checkout) -> Result<()> {
+    let venv_path = checkout.venv_path();
+    let project = checkout.project();
+    install_dependencies_to_cache(
+        project.name,
+        project.dependencies,
+        &venv_path,
+        project.python_version,
+        project.max_dep_date,
+    )
+}
+
 /// Recursively load a directory into the memory filesystem
-fn copy_directory_recursive(
+pub fn copy_directory_recursive(
     fs: &MemoryFileSystem,
     source_path: &Path,
     dest_path: &SystemPath,