#![allow(clippy::disallowed_names)]
use ruff_benchmark::criterion;
use ruff_benchmark::real_world_projects::{InstalledProject, RealWorldProject, TY_ECOSYSTEM_PIN};

use std::fmt::Write;
use std::ops::Range;

use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
use rustc_hash::FxHashSet;

use ruff_benchmark::TestFile;
use ruff_db::diagnostic::{Diagnostic, DiagnosticId, Severity};
use ruff_db::files::{File, system_path_to_file};
use ruff_db::source::source_text;
use ruff_db::system::{InMemorySystem, MemoryFileSystem, SystemPath, SystemPathBuf, TestSystem};
use ruff_ranged_value::RangedValue;
use ty_project::metadata::options::{AnalysisOptions, EnvironmentOptions, Options};
use ty_project::metadata::python_version::SupportedPythonVersion;
use ty_project::metadata::value::RelativePathBuf;
use ty_project::watch::{ChangeEvent, ChangedKind};
use ty_project::{CheckMode, Db, ProjectDatabase, ProjectMetadata};

mod ty_shared;

use ty_shared::{Case, setup_micro_case, setup_rayon};

struct FileCase {
    db: ProjectDatabase,
    fs: MemoryFileSystem,
    file: File,
    file_path: SystemPathBuf,
}

// "https://raw.githubusercontent.com/python/cpython/8e8a4baf652f6e1cee7acde9d78c4b6154539748/Lib/tomllib";
static TOMLLIB_FILES: [TestFile; 4] = [
    TestFile::new(
        "tomllib/__init__.py",
        include_str!("../resources/tomllib/__init__.py"),
    ),
    TestFile::new(
        "tomllib/_parser.py",
        include_str!("../resources/tomllib/_parser.py"),
    ),
    TestFile::new(
        "tomllib/_re.py",
        include_str!("../resources/tomllib/_re.py"),
    ),
    TestFile::new(
        "tomllib/_types.py",
        include_str!("../resources/tomllib/_types.py"),
    ),
];

/// A structured set of fields we use to do diagnostic comparisons.
///
/// This helps assert benchmark results. Previously, we would compare
/// the actual diagnostic output, but using `insta` inside benchmarks is
/// problematic, and updating the strings otherwise when diagnostic rendering
/// changes is a PITA.
type KeyDiagnosticFields = (
    DiagnosticId,
    Option<&'static str>,
    Option<Range<usize>>,
    &'static str,
    Severity,
);

static EXPECTED_TOMLLIB_DIAGNOSTICS: &[KeyDiagnosticFields] = &[];

fn tomllib_path(file: &TestFile) -> SystemPathBuf {
    SystemPathBuf::from("src").join(file.name())
}

fn setup_tomllib_case() -> FileCase {
    let system = TestSystem::default();
    let fs = system.memory_file_system().clone();

    fs.write_files_all(
        TOMLLIB_FILES
            .iter()
            .map(|file| (tomllib_path(file), file.code().to_string())),
    )
    .unwrap();

    let src_root = SystemPath::new("/src");
    let mut metadata = ProjectMetadata::discover(src_root, &system).unwrap();
    metadata.apply_override_options(Options {
        environment: Some(EnvironmentOptions {
            python_version: Some(RangedValue::cli(SupportedPythonVersion::Py312)),
            ..EnvironmentOptions::default()
        }),
        analysis: Some(AnalysisOptions {
            respect_type_ignore_comments: Some(false),
            ..AnalysisOptions::default()
        }),
        ..Options::default()
    });

    let mut db = ProjectDatabase::fallible(metadata, system).unwrap();
    let mut tomllib_files = FxHashSet::default();
    let mut re: Option<File> = None;

    for test_file in &TOMLLIB_FILES {
        let file = system_path_to_file(&db, tomllib_path(test_file)).unwrap();
        if test_file.name().ends_with("_re.py") {
            re = Some(file);
        }
        tomllib_files.insert(file);
    }

    let re = re.unwrap();

    db.set_check_mode(CheckMode::OpenFiles);
    db.project().set_open_files(&mut db, tomllib_files);

    let re_path = re.path(&db).as_system_path().unwrap().to_owned();
    FileCase {
        db,
        fs,
        file: re,
        file_path: re_path,
    }
}

fn benchmark_incremental(criterion: &mut Criterion) {
    fn setup() -> FileCase {
        let case = setup_tomllib_case();

        let result: Vec<_> = case.db.check();

        assert_diagnostics(&case.db, &result, EXPECTED_TOMLLIB_DIAGNOSTICS);

        case.fs
            .write_file_all(
                &case.file_path,
                format!(
                    "{}\n# A comment\n",
                    source_text(&case.db, case.file).as_str()
                ),
            )
            .unwrap();

        case
    }

    fn incremental(case: &mut FileCase) {
        let FileCase { db, .. } = case;

        db.apply_changes(&[ChangeEvent::Changed {
            path: case.file_path.clone(),
            kind: ChangedKind::FileContent,
        }]);

        let result = db.check();

        assert_eq!(result.len(), EXPECTED_TOMLLIB_DIAGNOSTICS.len());
    }

    setup_rayon();

    criterion.bench_function("ty_check_file[incremental]", |b| {
        b.iter_batched_ref(setup, incremental, BatchSize::SmallInput);
    });
}

fn benchmark_cold(criterion: &mut Criterion) {
    setup_rayon();

    criterion.bench_function("ty_check_file[cold]", |b| {
        b.iter_batched_ref(
            setup_tomllib_case,
            |case| {
                let FileCase { db, .. } = case;
                let result: Vec<_> = db.check();

                assert_diagnostics(db, &result, EXPECTED_TOMLLIB_DIAGNOSTICS);
            },
            BatchSize::SmallInput,
        );
    });
}

#[track_caller]
fn assert_diagnostics(db: &dyn Db, diagnostics: &[Diagnostic], expected: &[KeyDiagnosticFields]) {
    let normalized: Vec<_> = diagnostics
        .iter()
        .map(|diagnostic| {
            (
                diagnostic.id(),
                diagnostic
                    .primary_span()
                    .map(|span| span.expect_ty_file())
                    .map(|file| file.path(db).as_str()),
                diagnostic
                    .primary_span()
                    .and_then(|span| span.range())
                    .map(Range::<usize>::from),
                diagnostic.headline_message(),
                diagnostic.severity(),
            )
        })
        .collect();
    assert_eq!(&normalized, expected);
}

fn benchmark_many_string_assignments(criterion: &mut Criterion) {
    setup_rayon();

    criterion.bench_function("ty_micro[many_string_assignments]", |b| {
        b.iter_batched_ref(
            || {
                // This is a micro benchmark, but it is effectively identical to a code sample
                // observed "in the wild":
                setup_micro_case(
                    r#"
                    def f(x) -> str:
                        s = ""
                        # Each conditional doubles the size of the union of string literal types,
                        # so if we go up to attr10, we have 2**10 = 1024 string literal types
                        if x.attr1:
                            s += "attr1"
                        if x.attr2:
                            s += "attr2"
                        if x.attr3:
                            s += "attr3"
                        if x.attr4:
                            s += "attr4"
                        if x.attr5:
                            s += "attr5"
                        if x.attr6:
                            s += "attr6"
                        if x.attr7:
                            s += "attr7"
                        if x.attr8:
                            s += "attr8"
                        if x.attr9:
                            s += "attr9"
                        if x.attr10:
                            s += "attr10"
                        # The above checked how fast we are in building the union; this checks how
                        # we manage it once it is built. If implemented naively, this has to check
                        # each member of the union for compatibility with the Sized protocol.
                        if len(s) > 0:
                            s = s[:-3]
                        return s
                    "#,
                )
            },
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

fn benchmark_many_tuple_assignments(criterion: &mut Criterion) {
    setup_rayon();

    criterion.bench_function("ty_micro[many_tuple_assignments]", |b| {
        b.iter_batched_ref(
            || {
                // This is a micro benchmark, but it is effectively identical to a code sample
                // observed in https://github.com/astral-sh/ty/issues/362
                setup_micro_case(
                    r#"
                    def flag() -> bool:
                        return True

                    t = ()
                    if flag():
                        t += (1,)
                    if flag():
                        t += (2,)
                    if flag():
                        t += (3,)
                    if flag():
                        t += (4,)
                    if flag():
                        t += (5,)
                    if flag():
                        t += (6,)
                    if flag():
                        t += (7,)
                    if flag():
                        t += (8,)

                    # Perform some kind of operation on the union type
                    print(1 in t)
                    "#,
                )
            },
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

fn benchmark_tuple_implicit_instance_attributes(criterion: &mut Criterion) {
    setup_rayon();

    criterion.bench_function("ty_micro[many_tuple_assignments]", |b| {
        b.iter_batched_ref(
            || {
                // This is a regression benchmark for a case that used to hang:
                // https://github.com/astral-sh/ty/issues/765
                setup_micro_case(
                    r#"
                    from typing import Any

                    class A:
                        foo: tuple[Any, ...]

                    class B(A):
                        def __init__(self, parent: "C", x: tuple[Any]):
                            self.foo = parent.foo + x

                    class C(A):
                        def __init__(self, parent: B, x: tuple[Any]):
                            self.foo = parent.foo + x
                    "#,
                )
            },
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

fn benchmark_complex_constrained_attributes_1(criterion: &mut Criterion) {
    setup_rayon();

    criterion.bench_function("ty_micro[complex_constrained_attributes_1]", |b| {
        b.iter_batched_ref(
            || {
                // This is a regression benchmark for https://github.com/astral-sh/ty/issues/627.
                // Before this was fixed, the following sample would take >1s to type check.
                setup_micro_case(
                    r#"
                    class C:
                        def f(self: "C"):
                            if isinstance(self.a, str):
                                return

                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                    "#,
                )
            },
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert!(!result.is_empty());
            },
            BatchSize::SmallInput,
        );
    });
}

fn benchmark_complex_constrained_attributes_2(criterion: &mut Criterion) {
    setup_rayon();

    criterion.bench_function("ty_micro[complex_constrained_attributes_2]", |b| {
        b.iter_batched_ref(
            || {
                // This is similar to the case above, but now the attributes are actually defined.
                // https://github.com/astral-sh/ty/issues/711
                setup_micro_case(
                    r#"
                    class C:
                        def f(self: "C"):
                            if isinstance(self.a, str):
                                return

                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return
                            if isinstance(self.b, str):
                                return

                            self.a = ""
                            self.b = ""
                    "#,
                )
            },
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

fn benchmark_complex_constrained_attributes_3(criterion: &mut Criterion) {
    setup_rayon();

    criterion.bench_function("ty_micro[complex_constrained_attributes_3]", |b| {
        b.iter_batched_ref(
            || {
                // This is a regression test for https://github.com/astral-sh/ty/issues/758
                setup_micro_case(
                    r#"
                    class GridOut:
                        def __init__(self: "GridOut") -> None:
                            self._buffer = b""
                            self._position = 0

                        def _read_size_or_line(self: "GridOut", size: int = -1):
                            if size > self._position:
                                size = self._position
                                pass
                            if size == 0:
                                return bytes()

                            while size > 0:
                                if self._buffer:
                                    buf = self._buffer
                                    self._buffer = b""
                                else:
                                    buf = b""

                                if len(buf) > size:
                                    self._buffer = buf
                                    self._position -= len(self._buffer)
                    "#,
                )
            },
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

fn benchmark_many_enum_members(criterion: &mut Criterion) {
    const NUM_ENUM_MEMBERS: usize = 512;

    setup_rayon();

    let mut code = "from enum import Enum\n".to_string();

    code.push_str("class E(Enum):\n");
    for i in 0..NUM_ENUM_MEMBERS {
        writeln!(&mut code, "    m{i} = {i}").ok();
    }
    code.push('\n');

    code.push_str("print((");
    for i in 0..NUM_ENUM_MEMBERS {
        write!(&mut code, "E.m{i}, ").ok();
    }
    code.push_str("))");

    criterion.bench_function("ty_micro[many_enum_members]", |b| {
        b.iter_batched_ref(
            || setup_micro_case(&code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

/// Regression benchmark for membership narrowing over a large enum.
///
/// The right-hand-side union should be evaluated one listed member at a time. Evaluating it as a
/// single union expands the complete enum once per function instead.
fn benchmark_large_enum_membership(criterion: &mut Criterion) {
    const NUM_ENUM_MEMBERS: usize = 512;
    const NUM_FUNCTIONS: usize = 128;

    setup_rayon();

    let mut code = "from enum import Enum\n\nclass E(Enum):\n".to_string();
    for i in 0..NUM_ENUM_MEMBERS {
        writeln!(&mut code, "    m{i} = {i}").ok();
    }
    code.push('\n');

    for i in 0..NUM_FUNCTIONS {
        writeln!(
            &mut code,
            "def check_{i}(value: E) -> E:\n    if value in (E.m0, E.m1):\n        return value\n    return value\n"
        )
        .ok();
    }

    criterion.bench_function("ty_micro[large_enum_membership]", |b| {
        b.iter_batched_ref(
            || setup_micro_case(&code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

fn benchmark_enum_comparison(criterion: &mut Criterion, name: &str, code: &str) {
    setup_rayon();

    criterion.bench_function(name, |b| {
        b.iter_batched_ref(
            || setup_micro_case(code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

/// Regression benchmark for <https://github.com/astral-sh/ty/issues/3830>.
fn benchmark_narrowed_str_enum_comparison(criterion: &mut Criterion) {
    const NUM_ENUM_MEMBERS: usize = 256;

    let mut code = "from enum import StrEnum\n\nclass LargeEnum(StrEnum):\n".to_string();
    for index in 0..NUM_ENUM_MEMBERS {
        writeln!(&mut code, "    VALUE_{index} = \"value_{index}\"").ok();
    }
    code.push_str(
        "\n\ndef compare(left: LargeEnum, right: LargeEnum):\n    if right and left != right:\n        return\n    return left == right\n",
    );

    benchmark_enum_comparison(criterion, "ty_micro[narrowed_str_enum_comparison]", &code);
}

/// Regression benchmark for <https://github.com/astral-sh/ty/issues/4069>.
///
/// Compare a large enum with optional enum fields in a chain of conditions.
fn benchmark_optional_str_enum_comparison(criterion: &mut Criterion) {
    const NUM_ENUM_MEMBERS: usize = 256;

    let mut code =
        "from dataclasses import dataclass\nfrom enum import StrEnum\n\nclass ModelSlug(StrEnum):\n"
            .to_string();
    for index in 0..NUM_ENUM_MEMBERS {
        writeln!(&mut code, "    M{index} = \"m{index}\"").ok();
    }
    code.push_str(
        r#"

@dataclass
class Category:
    default_model: ModelSlug | None = None
    browsing_model: ModelSlug | None = None
    code_interpreter_model: ModelSlug | None = None
    plugins_model: ModelSlug | None = None
    dalle_model: ModelSlug | None = None


def belongs(slug: ModelSlug, category: Category) -> bool:
    return (
        category.default_model == slug
        or category.browsing_model == slug
        or category.code_interpreter_model == slug
        or category.plugins_model == slug
        or category.dalle_model == slug
    )
"#,
    );

    benchmark_enum_comparison(criterion, "ty_micro[optional_str_enum_comparison]", &code);
}

/// Ensure explicit enum-literal unions are compared as value sets, not member pairs.
fn benchmark_enum_literal_union_comparison(criterion: &mut Criterion) {
    const NUM_ENUM_MEMBERS: usize = 256;

    let mut code =
        "from enum import StrEnum\nfrom typing import Literal\n\nclass LargeEnum(StrEnum):\n"
            .to_string();
    for index in 0..NUM_ENUM_MEMBERS {
        writeln!(&mut code, "    VALUE_{index} = \"value_{index}\"").ok();
    }
    code.push_str("\nLeft = Literal[\n");
    for index in 0..NUM_ENUM_MEMBERS / 2 {
        writeln!(&mut code, "    LargeEnum.VALUE_{index},").ok();
    }
    code.push_str("]\nRight = Literal[\n");
    for index in NUM_ENUM_MEMBERS / 2..NUM_ENUM_MEMBERS {
        writeln!(&mut code, "    LargeEnum.VALUE_{index},").ok();
    }
    code.push_str("]\n\n\ndef compare(left: Left, right: Right):\n    return left == right\n");

    benchmark_enum_comparison(criterion, "ty_micro[enum_literal_union_comparison]", &code);
}

/// Ensure the comparison profile is reused instead of scanning every member per expression.
fn benchmark_repeated_str_enum_comparisons(criterion: &mut Criterion) {
    const NUM_ENUM_MEMBERS: usize = 1_024;
    const NUM_COMPARISONS: usize = 1_000;

    let mut code = "from enum import StrEnum\n\nclass LargeEnum(StrEnum):\n".to_string();
    for index in 0..NUM_ENUM_MEMBERS {
        writeln!(&mut code, "    VALUE_{index} = \"value_{index}\"").ok();
    }
    code.push_str("\n\ndef compare(left: LargeEnum, right: LargeEnum):\n");
    for _ in 0..NUM_COMPARISONS {
        code.push_str("    left == right\n");
    }

    benchmark_enum_comparison(criterion, "ty_micro[repeated_str_enum_comparisons]", &code);
}

/// Ensure comparison constraints between two large scalar enums do not compare every member pair.
fn benchmark_cross_str_enum_comparison(criterion: &mut Criterion) {
    const NUM_ENUM_MEMBERS: usize = 256;

    let mut code = "from enum import StrEnum\n".to_string();
    for side in ["Left", "Right"] {
        writeln!(&mut code, "\nclass {side}(StrEnum):").ok();
        for index in 0..NUM_ENUM_MEMBERS {
            writeln!(&mut code, "    VALUE_{index} = \"value_{index}\"").ok();
        }
    }
    code.push_str(
        "\n\ndef compare(left: Left, right: Right):\n    if left != right:\n        return\n    return left == right\n",
    );

    benchmark_enum_comparison(criterion, "ty_micro[cross_str_enum_comparison]", &code);
}

/// Ensure comparisons of unions spanning several scalar enum classes avoid member-pair expansion.
fn benchmark_mixed_str_enum_comparison(criterion: &mut Criterion) {
    const NUM_ENUM_CLASSES: usize = 8;
    const NUM_ENUM_MEMBERS: usize = 32;

    let mut code = "from enum import StrEnum\n".to_string();
    for side in ["Left", "Right"] {
        for class_index in 0..NUM_ENUM_CLASSES {
            writeln!(&mut code, "\nclass {side}{class_index}(StrEnum):").ok();
            for member_index in 0..NUM_ENUM_MEMBERS {
                writeln!(
                    &mut code,
                    "    VALUE_{member_index} = \"class_{class_index}_value_{member_index}\""
                )
                .ok();
            }
        }
    }
    let class_union = |side| {
        (0..NUM_ENUM_CLASSES)
            .map(|index| format!("{side}{index}"))
            .collect::<Vec<_>>()
            .join(" | ")
    };
    writeln!(
        &mut code,
        "\ndef compare(left: {}, right: {}):\n    if left != right:\n        return\n    return left == right",
        class_union("Left"),
        class_union("Right"),
    )
    .ok();

    benchmark_enum_comparison(criterion, "ty_micro[mixed_str_enum_comparison]", &code);
}

/// Micro-benchmark that tests our performance when slicing and unpacking
/// a very large tuple that has many varied literal strings inside it.
///
/// Adapted from <https://github.com/MMD-Blender/blender_mmd_tools/blob/6ae13d6039763b6813832622c13da464983e93b3/mmd_tools/m17n.py>
fn benchmark_very_large_tuple(criterion: &mut Criterion) {
    /// Number of entries in the tuple -- must be >64 to trigger the literal-promotion optimization
    const NUM_ENTRIES: usize = 65;

    setup_rayon();

    let mut code = "translations_tuple = (\n".to_string();
    for i in 0..NUM_ENTRIES {
        writeln!(
            &mut code,
            r#"    (("*", "Description {i}"), (("ref.path.{i}",), ()), ("ja_JP", "翻訳{i}", (False, ())), ("zh_HANS", "翻译{i}", (False, ()))),"#
        )
        .ok();
    }
    code.push_str(")\n\n");

    // This slice operation (msg[2:]) is what triggers the expensive type inference:
    // ty must compute the union of all possible slice results.
    code.push_str(
        "\
for msg in translations_tuple:
    for lang, trans, (is_fuzzy, comments) in msg[2:]:
        pass
",
    );

    criterion.bench_function("ty_micro[very_large_tuple]", |b| {
        b.iter_batched_ref(
            || setup_micro_case(&code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

fn benchmark_many_enum_members_2(criterion: &mut Criterion) {
    const NUM_ENUM_MEMBERS: usize = 48;

    setup_rayon();

    let mut code = "\
from enum import Enum
from typing_extensions import assert_never

class E(Enum):
"
    .to_string();

    for i in 0..NUM_ENUM_MEMBERS {
        writeln!(&mut code, "    m{i} = {i}").ok();
    }

    code.push_str(
        "
    def method(self):
        match self:",
    );

    for i in 0..NUM_ENUM_MEMBERS {
        write!(
            &mut code,
            "
            case E.m{i}:
                pass"
        )
        .ok();
    }

    code.push_str(
        "
            case _:
                assert_never(self)",
    );

    criterion.bench_function("ty_micro[many_enum_members_2]", |b| {
        b.iter_batched_ref(
            || setup_micro_case(&code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

/// Regression benchmark for protocol-to-protocol non-assignability when the target protocol
/// has one more member than the source protocol.
///
/// Without the member-count gate, repeated return-type checks spend time comparing every shared
/// protocol member before eventually failing on the missing one.
fn benchmark_many_protocol_members_mismatch(criterion: &mut Criterion) {
    const NUM_MEMBERS: usize = 800;
    const NUM_FUNCTIONS: usize = 400;

    setup_rayon();

    let mut code = "from typing import Protocol\n\nclass Small(Protocol):\n".to_string();
    for i in 0..NUM_MEMBERS {
        writeln!(&mut code, "    member_{i}: int").ok();
    }

    code.push_str("\nclass Big(Protocol):\n");
    for i in 0..=NUM_MEMBERS {
        writeln!(&mut code, "    member_{i}: int").ok();
    }

    code.push('\n');
    for i in 0..NUM_FUNCTIONS {
        writeln!(
            &mut code,
            "def check_{i}(value: Small) -> Big:\n    return value\n"
        )
        .ok();
    }

    criterion.bench_function("ty_micro[many_protocol_members_mismatch]", |b| {
        b.iter_batched_ref(
            || setup_micro_case(&code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), NUM_FUNCTIONS);
            },
            BatchSize::SmallInput,
        );
    });
}

/// Regression benchmark for large calls to a gradual variadic tail.
///
/// Without the gradual-call shortcut, every positional argument type is folded into the same
/// `*args` parameter type, making the call checker repeatedly grow a large union.
fn benchmark_gradual_vararg_call(criterion: &mut Criterion) {
    const NUM_ARGUMENTS: usize = 256;

    setup_rayon();

    let mut code = "\
from typing import Any

def accepts_anything(first: int, *args: Any, **kwargs: Any) -> None: ...

accepts_anything(
    0,
"
    .to_string();

    for i in 0..NUM_ARGUMENTS {
        writeln!(&mut code, r#"    ("field_{i}", {i}),"#).ok();
    }

    code.push_str(")\n");

    criterion.bench_function("ty_micro[gradual_vararg_call]", |b| {
        b.iter_batched_ref(
            || setup_micro_case(&code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

/// Regression benchmark for many precise arguments flowing into one variadic parameter.
///
/// The declared type is intentionally non-gradual, so argument checks still run. The important
/// part is that the call checker should not repeatedly rebuild a growing union for `*args`.
fn benchmark_vararg_parameter_type_accumulation(criterion: &mut Criterion) {
    const NUM_ARGUMENTS: usize = 256;

    setup_rayon();

    let mut code = "\
def accepts_objects(first: int, *args: object) -> None: ...

accepts_objects(
    0,
"
    .to_string();

    for i in 0..NUM_ARGUMENTS {
        writeln!(&mut code, r#"    ("field_{i}", {i}),"#).ok();
    }

    code.push_str(")\n");

    criterion.bench_function("ty_micro[vararg_parameter_type_accumulation]", |b| {
        b.iter_batched_ref(
            || setup_micro_case(&code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

/// Benchmark for narrowing a large union type through multiple match statements.
///
/// This is extracted from egglog-python's `pretty.py`, where a ~30-class union type
/// (`AllDecls`) is narrowed by exhaustive match statements.
///
/// Sample code structure:
/// ```python
/// from __future__ import annotations
/// from dataclasses import dataclass
///
/// @dataclass
/// class C0:
///     value: int
/// ...
///
/// AllDecls = C0 | C1 | ...
///
/// def process(decl: AllDecls) -> None:
///     match decl:
///         case C0(): pass
///         ...
///         case _: pass
/// ```
fn benchmark_large_union_narrowing(criterion: &mut Criterion) {
    const NUM_CLASSES: usize = 30;
    const NUM_MATCH_BRANCHES: usize = 29;

    setup_rayon();

    let mut code =
        "from __future__ import annotations\nfrom dataclasses import dataclass\n\n".to_string();

    for i in 0..NUM_CLASSES {
        writeln!(&mut code, "@dataclass\nclass C{i}:\n    value: int\n").ok();
    }

    code.push_str("AllDecls = ");
    for i in 0..NUM_CLASSES {
        if i > 0 {
            code.push_str(" | ");
        }
        write!(&mut code, "C{i}").ok();
    }
    code.push_str("\n\n");

    code.push_str("def process(decl: AllDecls) -> None:\n    match decl:\n");
    for i in 0..NUM_MATCH_BRANCHES {
        writeln!(&mut code, "        case C{i}():\n            pass").ok();
    }
    code.push_str("        case _:\n            pass\n\n");

    criterion.bench_function("ty_micro[large_union_narrowing]", |b| {
        b.iter_batched_ref(
            || setup_micro_case(&code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

/// Benchmark for narrowing through a long `isinstance` elif chain.
///
/// This pattern is common in visitor-style dispatch code (e.g. koda-validate's
/// `generate_schema_validator`) where a base class parameter is narrowed through
/// many sequential `isinstance` checks.
///
/// Sample code structure:
/// ```python
/// class Base: ...
/// class C0(Base): ...
/// class C1(Base): ...
/// ...
///
/// def f(obj: Base) -> None:
///     if isinstance(obj, C0):
///         pass
///     elif isinstance(obj, C1):
///         pass
///     ...
/// ```
fn benchmark_large_isinstance_narrowing(criterion: &mut Criterion) {
    const NUM_CLASSES: usize = 50;

    setup_rayon();

    let mut code = "class Base: ...\n".to_string();
    for i in 0..NUM_CLASSES {
        writeln!(&mut code, "class C{i}(Base): ...").ok();
    }
    code.push('\n');

    code.push_str("def f(obj: Base) -> None:\n");
    for i in 0..NUM_CLASSES {
        if i == 0 {
            writeln!(&mut code, "    if isinstance(obj, C{i}):").ok();
        } else {
            writeln!(&mut code, "    elif isinstance(obj, C{i}):").ok();
        }
        code.push_str("        pass\n");
    }

    criterion.bench_function("ty_micro[large_isinstance_narrowing]", |b| {
        b.iter_batched_ref(
            || setup_micro_case(&code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

const NUM_LITERAL_FALLTHROUGH_ARMS: usize = 40;
const NUM_LITERAL_OR_PATTERN_ALTERNATIVES: usize = 23;

fn is_literal_fallthrough_arm(index: usize) -> bool {
    index.is_multiple_of(8) || index == NUM_LITERAL_FALLTHROUGH_ARMS - 4
}

fn literal_fallthrough_value(index: usize) -> String {
    let prefix = if is_literal_fallthrough_arm(index) {
        'f'
    } else {
        'v'
    };
    format!("{prefix}{index:02}")
}

fn literal_fallthrough_prelude(guarded_any: bool) -> String {
    let imports = if guarded_any {
        "from typing import Any, Literal, TypeGuard\n\n"
    } else {
        "from typing import Literal\n\n"
    };
    let mut code = imports.to_string();

    code.push_str("Value = Literal[\n");
    for index in 0..NUM_LITERAL_FALLTHROUGH_ARMS {
        writeln!(&mut code, r#"    "{}","#, literal_fallthrough_value(index)).ok();
    }
    code.push_str("]\n\n");

    if guarded_any {
        code.push_str("def keep_value(value: object) -> TypeGuard[Value]:\n    return True\n\n");
    }

    code
}

fn literal_match_fallthrough_code(guarded_any: bool) -> String {
    let parameter_ty = if guarded_any { "Value | Any" } else { "Value" };
    let mut code = literal_fallthrough_prelude(guarded_any);
    writeln!(
        &mut code,
        "def check_literal_match(value: {parameter_ty}) -> None:\n    match value:"
    )
    .ok();

    for index in 0..NUM_LITERAL_FALLTHROUGH_ARMS {
        let guard = if guarded_any && index == 0 {
            " if keep_value(value)"
        } else {
            ""
        };
        writeln!(
            &mut code,
            r#"        case "{}"{guard}:"#,
            literal_fallthrough_value(index),
        )
        .ok();
        code.push_str(if is_literal_fallthrough_arm(index) {
            "            pass\n"
        } else {
            "            return\n"
        });
    }

    code.push_str("\n    repr(value)\n");
    code
}

fn literal_equality_fallthrough_code() -> String {
    let mut code = literal_fallthrough_prelude(true);
    code.push_str("def check_literal_equality(value: Value | Any) -> None:\n");

    for index in 0..NUM_LITERAL_FALLTHROUGH_ARMS {
        let keyword = if index == 0 { "if" } else { "elif" };
        write!(
            &mut code,
            r#"    {keyword} value == "{}""#,
            literal_fallthrough_value(index)
        )
        .ok();
        if index == 0 {
            code.push_str(" and keep_value(value)");
        }
        code.push_str(":\n");
        code.push_str(if is_literal_fallthrough_arm(index) {
            "        pass\n"
        } else {
            "        return\n"
        });
    }

    code.push_str("\n    repr(value)\n");
    code
}

fn literal_or_pattern_reachability_code() -> String {
    let mut code = "from typing import Any\n\ndef check(item: Any) -> None:\n    x: int\n    match item:\n        case ".to_string();

    for index in 0..NUM_LITERAL_OR_PATTERN_ALTERNATIVES {
        if index > 0 {
            code.push_str(" | ");
        }
        write!(&mut code, "{index}").ok();
    }

    code.push_str(":\n            x = 1\n");
    code
}

/// Benchmarks `match` fallthrough over a finite `Literal` subject.
///
/// This generates a large version of this Python shape:
/// ```python
/// def check(value: Value) -> None:
///     match value:
///         case "f00":
///             pass
///         case "v01":
///             return
///         case "f08":
///             pass
///
///     repr(value)
/// ```
///
/// Some cases return and some continue after the `match`. This measures how long it takes to infer
/// the type of `value` after many literal match cases when the subject has a finite literal type.
fn benchmark_literal_match_fallthrough(criterion: &mut Criterion) {
    benchmark_literal_fallthrough(
        criterion,
        "ty_micro[literal_match_fallthrough]",
        &literal_match_fallthrough_code(false),
    );
}

/// Benchmarks guarded `match` fallthrough over a gradual literal subject.
///
/// This generates a large version of this Python shape:
/// ```python
/// def keep_value(value: object) -> TypeGuard[Value]:
///     return True
///
/// def check(value: Value | Any) -> None:
///     match value:
///         case "f00" if keep_value(value):
///             pass
///         case "v01":
///             return
///         case "f08":
///             pass
///
///     repr(value)
/// ```
///
/// This is the same match shape, but the subject can also be `Any` and the first continuing case
/// uses a `TypeGuard`. It measures how that guard affects the type of `value` after many match
/// cases.
fn benchmark_literal_match_fallthrough_guarded_any(criterion: &mut Criterion) {
    benchmark_literal_fallthrough(
        criterion,
        "ty_micro[literal_match_fallthrough_guarded_any]",
        &literal_match_fallthrough_code(true),
    );
}

/// Benchmarks guarded fallthrough through many equality checks.
///
/// This generates a large version of this Python shape:
/// ```python
/// def check(value: Value | Any) -> None:
///     if value == "f00" and keep_value(value):
///         pass
///     elif value == "v01":
///         return
///     elif value == "f08":
///         pass
///
///     repr(value)
/// ```
///
/// This is the guarded gradual case written with `if`/`elif` equality checks instead of `match`.
/// Some branches return and some continue after the conditional, so the benchmark measures the
/// type of `value` after many comparisons.
fn benchmark_literal_equality_fallthrough_guarded_any(criterion: &mut Criterion) {
    benchmark_literal_fallthrough(
        criterion,
        "ty_micro[literal_equality_fallthrough_guarded_any]",
        &literal_equality_fallthrough_code(),
    );
}

/// Regression benchmark for <https://github.com/astral-sh/ty/issues/3880>.
///
/// Reachability analysis for a large literal OR pattern on `Any` used to rebuild the remaining
/// subject from all preceding alternatives. Assigning to an annotation-only local forces the case
/// predicate to be evaluated while resolving the live declaration, exposing the exponential
/// growth.
fn benchmark_literal_or_pattern_reachability(criterion: &mut Criterion) {
    setup_rayon();

    let code = literal_or_pattern_reachability_code();
    criterion.bench_function("ty_micro[literal_or_pattern_reachability]", |b| {
        b.iter_batched_ref(
            || setup_micro_case(&code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

fn benchmark_literal_fallthrough(criterion: &mut Criterion, name: &str, code: &str) {
    setup_rayon();

    criterion.bench_function(name, |b| {
        b.iter_batched_ref(
            || setup_micro_case(code),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

/// Regression benchmark for <https://github.com/astral-sh/ty/issues/3120>.
///
/// Sequential (`TypeIs`) narrowing on a large `Literal` union, combined with
/// `match`/`assert_never` on another `Literal` union, caused a combinatorial
/// explosion when the `PredicateNode::IsNonTerminalCall` optimization was
/// removed.
fn benchmark_typeis_narrowing(criterion: &mut Criterion) {
    setup_rayon();

    criterion.bench_function("ty_micro[typeis_narrowing]", |b| {
        b.iter_batched_ref(
            || setup_micro_case(include_str!("../resources/typeis_narrowing.py")),
            |case| {
                let Case { db } = case;
                let result = db.check();
                assert_eq!(result.len(), 0);
            },
            BatchSize::SmallInput,
        );
    });
}

/// Regression benchmark for <https://github.com/astral-sh/ty/issues/3986>.
///
/// Non-terminal-call predicates must gate later narrowing. Keeping these scope-wide constraints in
/// an append-only tree avoids eagerly rewriting every live place state after each call. Exercise
/// unnarrowed, already-narrowed, and fixed-reachability bindings because each takes a different
/// path through reachability and narrowing evaluation.
fn benchmark_repeated_statement_calls(criterion: &mut Criterion) {
    setup_rayon();

    let cases = [
        (
            "ty_micro[repeated_statement_calls]",
            String::from("def f() -> None:\n    value = 'abc'\n"),
            "    value.upper()\n",
        ),
        (
            "ty_micro[repeated_statement_calls_pre_narrowed]",
            String::from(
                "def f(value: str | None) -> None:\n    if value is None:\n        return\n",
            ),
            "    value.upper()\n",
        ),
        (
            "ty_micro[repeated_statement_calls_fixed_reachability]",
            String::from("def f(value: str, flag: bool) -> None:\n    if flag:\n"),
            "        value.upper()\n",
        ),
    ];

    for (name, mut code, statement) in cases {
        code.push_str(&statement.repeat(1_500));
        criterion.bench_function(name, |b| {
            b.iter_batched_ref(
                || setup_micro_case(&code),
                |case| {
                    let Case { db } = case;
                    let result = db.check();
                    assert_eq!(result.len(), 0);
                },
                BatchSize::SmallInput,
            );
        });
    }
}

struct ProjectBenchmark<'a> {
    project: InstalledProject<'a>,
    fs: MemoryFileSystem,
    max_diagnostics: usize,
    freeze_inputs: bool,
}

impl<'a> ProjectBenchmark<'a> {
    fn new(project: RealWorldProject<'a>, max_diagnostics: usize) -> Self {
        let setup_project = project.setup().expect("Failed to setup project");
        let fs = setup_project
            .copy_to_memory_fs()
            .expect("Failed to copy project to memory fs");

        Self {
            project: setup_project,
            fs,
            max_diagnostics,
            freeze_inputs: false,
        }
    }

    fn freeze_inputs(mut self) -> Self {
        self.freeze_inputs = true;
        self
    }

    fn setup_iteration(&self) -> ProjectDatabase {
        let system = TestSystem::new(InMemorySystem::from_memory_fs(self.fs.clone()));

        let src_root = SystemPath::new("/");
        let mut metadata = ProjectMetadata::discover(src_root, &system).unwrap();

        metadata.apply_override_options(Options {
            environment: Some(EnvironmentOptions {
                python_version: Some(RangedValue::cli(self.project.config.python_version)),
                python: Some(RelativePathBuf::cli(SystemPath::new(".venv"))),
                ..EnvironmentOptions::default()
            }),
            ..Options::default()
        });

        let mut db = ProjectDatabase::fallible(metadata, system).unwrap();

        db.project().set_included_paths(
            &mut db,
            self.project
                .check_paths()
                .iter()
                .map(|path| SystemPathBuf::from(*path))
                .collect(),
        );

        if self.freeze_inputs {
            db.freeze();
        }

        db
    }
}

#[track_caller]
fn bench_project(benchmark: &ProjectBenchmark, criterion: &mut Criterion) {
    bench_project_named(benchmark, criterion, benchmark.project.config.name);
}

#[track_caller]
fn bench_project_named(
    benchmark: &ProjectBenchmark,
    criterion: &mut Criterion,
    benchmark_name: &str,
) {
    fn check_project(db: &mut ProjectDatabase, project_name: &str, max_diagnostics: usize) {
        let result = db.check();
        let diagnostics = result.len();

        if diagnostics > max_diagnostics {
            let details = result
                .into_iter()
                .map(|diagnostic| diagnostic.concise_message().to_string())
                .collect::<Vec<_>>()
                .join("\n  ");
            assert!(
                diagnostics <= max_diagnostics,
                "{project_name}: Expected <={max_diagnostics} diagnostics but got {diagnostics}:\n  {details}",
            );
        }
    }

    setup_rayon();

    let mut group = criterion.benchmark_group("project");
    group.sampling_mode(criterion::SamplingMode::Flat);
    group.bench_function(benchmark_name, |b| {
        b.iter_batched_ref(
            || benchmark.setup_iteration(),
            |db| check_project(db, benchmark_name, benchmark.max_diagnostics),
            BatchSize::SmallInput,
        );
    });
}

fn hydra(criterion: &mut Criterion) {
    let benchmark = ProjectBenchmark::new(
        RealWorldProject {
            name: "hydra-zen",
            repository: "https://github.com/mit-ll-responsible-ai/hydra-zen",
            commit: "03a01096ea6a7c574fdf0b9990056506e566df2d",
            paths: &["src", "tests/annotations"],
            dependencies: &["pydantic", "beartype", "hydra-core"],
            max_dep_date: TY_ECOSYSTEM_PIN,
            python_version: SupportedPythonVersion::Py311,
        },
        520,
    );

    bench_project(&benchmark, criterion);
}

fn attrs(criterion: &mut Criterion) {
    let benchmark = ProjectBenchmark::new(
        RealWorldProject {
            name: "attrs",
            repository: "https://github.com/python-attrs/attrs",
            commit: "89fae8300f484544c1b7678cea5efe58c551fbb9",
            paths: &["src/attrs", "src/attr", "typing-examples"],
            dependencies: &[],
            max_dep_date: TY_ECOSYSTEM_PIN,
            python_version: SupportedPythonVersion::Py311,
        },
        104,
    );

    bench_project(&benchmark, criterion);

    // Keep one real-world benchmark frozen to catch regressions from newly added inputs.
    let frozen_benchmark = benchmark.freeze_inputs();
    bench_project_named(&frozen_benchmark, criterion, "attrs (frozen inputs)");
}

fn anyio(criterion: &mut Criterion) {
    let benchmark = ProjectBenchmark::new(
        RealWorldProject {
            name: "anyio",
            repository: "https://github.com/agronholm/anyio",
            commit: "ffe91331adb912c5d150f5d373f7cd28a0e96a62",
            paths: &["src"],
            dependencies: &["exceptiongroup", "idna", "pytest"],
            max_dep_date: TY_ECOSYSTEM_PIN,
            python_version: SupportedPythonVersion::Py311,
        },
        110,
    );

    bench_project(&benchmark, criterion);
}

fn datetype(criterion: &mut Criterion) {
    let benchmark = ProjectBenchmark::new(
        RealWorldProject {
            name: "DateType",
            repository: "https://github.com/glyph/DateType",
            commit: "a6ebb954cd18302a031a29b2f65e077b8e7776d4",
            paths: &["src"],
            dependencies: &[],
            max_dep_date: TY_ECOSYSTEM_PIN,
            python_version: SupportedPythonVersion::Py311,
        },
        17,
    );

    bench_project(&benchmark, criterion);
}

criterion_group!(check_file, benchmark_cold, benchmark_incremental);
criterion_group!(
    micro,
    benchmark_many_string_assignments,
    benchmark_many_tuple_assignments,
    benchmark_tuple_implicit_instance_attributes,
    benchmark_complex_constrained_attributes_1,
    benchmark_complex_constrained_attributes_2,
    benchmark_complex_constrained_attributes_3,
    benchmark_gradual_vararg_call,
    benchmark_large_enum_membership,
    benchmark_many_enum_members,
    benchmark_narrowed_str_enum_comparison,
    benchmark_optional_str_enum_comparison,
    benchmark_enum_literal_union_comparison,
    benchmark_repeated_str_enum_comparisons,
    benchmark_cross_str_enum_comparison,
    benchmark_mixed_str_enum_comparison,
    benchmark_many_enum_members_2,
    benchmark_many_protocol_members_mismatch,
    benchmark_vararg_parameter_type_accumulation,
    benchmark_very_large_tuple,
    benchmark_large_union_narrowing,
    benchmark_large_isinstance_narrowing,
    benchmark_literal_match_fallthrough,
    benchmark_literal_match_fallthrough_guarded_any,
    benchmark_literal_equality_fallthrough_guarded_any,
    benchmark_literal_or_pattern_reachability,
    benchmark_typeis_narrowing,
    benchmark_repeated_statement_calls,
);
criterion_group!(project, anyio, attrs, hydra, datetype);
criterion_main!(check_file, micro, project);
