test_harness.py

from __future__ import annotations

import hashlib
import os
import re
import shutil
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

from stdlib_overlay import (
    STDLIB_OVERLAY_ROOT,
    STDLIB_PATCH_ROOT,
    prepare_stdlib_source,
    rustc_info,
)

ROOT = Path(__file__).resolve().parent
TARGET_SPEC = ROOT / "jvm-unknown-jvm.json"
TEST_TARGET_DIR = ROOT / "target" / "test-suite"
TEST_CONFIG = ROOT / "config.toml"
CORE_BUILD_MANIFEST = ROOT / "tests" / "support" / "core_build" / "Cargo.toml"
TEST_TYPES = ("binary", "multicrate", "integration", "kotlin", "cargo_jvm")
CACHE_TAG = (
    "Signature: 8a477f597d28d172789f06886806bc55\n"
    "# This file is a cache directory tag created by rustc_codegen_jvm.\n"
)
_STDLIB_SOURCE: Path | None = None
_STDLIB_FINGERPRINT: str | None = None


@dataclass(frozen=True)
class TestCase:
    directory: Path
    kind: str
    package_name: str

    @property
    def name(self) -> str:
        return self.directory.name

    @property
    def artifact_name(self) -> str:
        return self.package_name.replace("-", "_")


def cpu_count() -> int:
    return os.cpu_count() or 1


def resolve_workers(requested: int | None) -> int:
    if requested is not None:
        if requested < 1:
            raise ValueError("--jobs must be at least 1")
        return requested
    return min(cpu_count(), 4)


def cargo_jobs(workers: int) -> int:
    return max(1, cpu_count() // workers)


def package_name(manifest: Path) -> str:
    in_package = False
    for line in manifest.read_text(encoding="utf-8").splitlines():
        stripped = line.strip()
        if stripped.startswith("["):
            in_package = stripped == "[package]"
            continue
        if in_package:
            match = re.fullmatch(r'name\s*=\s*"([^"]+)"', stripped)
            if match:
                return match.group(1)
    raise ValueError(f"missing package name in {manifest}")


def discover_tests(
    only_run: set[str] | None = None,
    dont_run: set[str] | None = None,
) -> list[TestCase]:
    excluded = dont_run or set()
    tests: list[TestCase] = []
    for kind in TEST_TYPES:
        parent = ROOT / "tests" / kind
        if not parent.is_dir():
            continue
        for directory in sorted(path for path in parent.iterdir() if path.is_dir()):
            manifest = directory / "Cargo.toml"
            if not manifest.exists():
                continue
            if only_run is not None and directory.name not in only_run:
                continue
            if directory.name in excluded:
                continue
            tests.append(TestCase(directory, kind, package_name(manifest)))
    return tests


def validate_configuration() -> None:
    missing = [path for path in (TARGET_SPEC, TEST_CONFIG) if not path.exists()]
    if missing:
        paths = ", ".join(str(path) for path in missing)
        raise RuntimeError(
            f"missing generated test configuration: {paths}; run `python3 build.py all`"
        )


def prepare_stdlib() -> tuple[Path, str]:
    global _STDLIB_SOURCE, _STDLIB_FINGERPRINT
    if _STDLIB_SOURCE is None or _STDLIB_FINGERPRINT is None:
        sysroot, _, commit = rustc_info()
        _STDLIB_SOURCE, overlay_hash = prepare_stdlib_source(sysroot, commit)
        _STDLIB_FINGERPRINT = f"{commit}:{overlay_hash}"
    return _STDLIB_SOURCE, _STDLIB_FINGERPRINT


def stdlib_build_environment(base: dict[str, str] | None = None) -> dict[str, str]:
    source, _ = prepare_stdlib()
    environment = (base if base is not None else os.environ).copy()
    environment["__CARGO_TESTS_ONLY_SRC_ROOT"] = str(source)
    return environment


def prepare_shared_cache() -> bool:
    """Invalidate all test artifacts when the compiler toolchain inputs change."""
    validate_configuration()
    _, stdlib_fingerprint = prepare_stdlib()
    inputs = [
        Path(__file__),
        TARGET_SPEC,
        TEST_CONFIG,
        ROOT / "runtime" / "build" / "libs" / "runtime-0.1.0.jar",
        ROOT / "java-linker" / "target" / "release" / (
            "java-linker.exe" if os.name == "nt" else "java-linker"
        ),
        ROOT / "cargo-jvm" / "target" / "release" / (
            "cargo-jvm.exe" if os.name == "nt" else "cargo-jvm"
        ),
    ]
    inputs.extend(path for path in STDLIB_OVERLAY_ROOT.rglob("*") if path.is_file())
    inputs.extend(path for path in STDLIB_PATCH_ROOT.rglob("*.patch") if path.is_file())
    backend_candidates = list((ROOT / "target" / "release").glob("*rustc_codegen_jvm.*"))
    inputs.extend(path for path in backend_candidates if path.suffix in {".dll", ".dylib", ".so"})

    digest = hashlib.sha256()
    digest.update(stdlib_fingerprint.encode("utf-8"))
    for path in sorted(inputs):
        if not path.exists():
            raise RuntimeError(f"missing test toolchain input: {path}; run `python3 build.py all`")
        digest.update(str(path).encode())
        digest.update(path.read_bytes())
    fingerprint = digest.hexdigest()
    marker = TEST_TARGET_DIR / ".harness-fingerprint"
    if marker.exists() and marker.read_text(encoding="utf-8") == fingerprint:
        (TEST_TARGET_DIR / "CACHEDIR.TAG").write_text(CACHE_TAG, encoding="utf-8")
        return False

    if TEST_TARGET_DIR.exists():
        shutil.rmtree(TEST_TARGET_DIR)
    TEST_TARGET_DIR.mkdir(parents=True)
    (TEST_TARGET_DIR / "CACHEDIR.TAG").write_text(CACHE_TAG, encoding="utf-8")
    marker.write_text(fingerprint, encoding="utf-8")
    return True


def cargo_build_command(manifest: Path, release: bool, jobs: int) -> list[str]:
    command = [
        "cargo",
        "build",
        "--manifest-path",
        str(manifest),
        "--target",
        str(TARGET_SPEC),
        "-Zjson-target-spec",
        "-Zbuild-std=std,panic_unwind",
        "-Zbuild-std-features=panic-unwind",
        "--target-dir",
        str(TEST_TARGET_DIR),
        "--config",
        str(TEST_CONFIG),
        "--jobs",
        str(jobs),
    ]
    if release:
        command.extend(
            ["--release", "--config", 'profile.release.debug="line-tables-only"']
        )
    return command


def run_command(
    command: list[str],
    *,
    cwd: Path | None = None,
    env: dict[str, str] | None = None,
    input_text: str | None = None,
) -> subprocess.CompletedProcess[str]:
    if input_text is None:
        return subprocess.run(
            command,
            cwd=cwd or ROOT,
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            encoding="utf-8",
            errors="replace",
        )

    # Text-mode pipes translate LF to CRLF on Windows. Feed UTF-8 bytes so a
    # test fixture reaches Rust's stdin unchanged on every host.
    completed = subprocess.run(
        command,
        cwd=cwd or ROOT,
        env=env,
        input=input_text.encode("utf-8"),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    return subprocess.CompletedProcess(
        completed.args,
        completed.returncode,
        completed.stdout.decode("utf-8", errors="replace"),
        completed.stderr.decode("utf-8", errors="replace"),
    )


def prime_core(release: bool) -> subprocess.CompletedProcess[str]:
    validate_configuration()
    return run_command(
        cargo_build_command(CORE_BUILD_MANIFEST, release, cpu_count()),
        env=stdlib_build_environment(),
    )


def build_test(
    test: TestCase,
    release: bool,
    jobs: int,
    *,
    env: dict[str, str] | None = None,
) -> subprocess.CompletedProcess[str]:
    if test.kind == "cargo_jvm":
        cargo_jvm = ROOT / "cargo-jvm" / "target" / "release" / (
            "cargo-jvm.exe" if os.name == "nt" else "cargo-jvm"
        )
        command = [
            sys.executable,
            str(test.directory / "run.py"),
            "--cargo-jvm",
            str(cargo_jvm),
            "--backend",
            str(ROOT),
            "--target-dir",
            str(TEST_TARGET_DIR),
        ]
        if release:
            command.append("--release")
        return run_command(command, env=stdlib_build_environment(env))

    if test.kind in ("integration", "kotlin"):
        cargo_jvm = ROOT / "cargo-jvm" / "target" / "release" / (
            "cargo-jvm.exe" if os.name == "nt" else "cargo-jvm"
        )
        command = [
            str(cargo_jvm),
            "package",
            "--manifest-path",
            str(test.directory / "Cargo.toml"),
            "--target-dir",
            str(TEST_TARGET_DIR),
            "--jobs",
            str(jobs),
            "--output",
            str(jar_path(test, release)),
        ]
        if release:
            command.extend(
                ["--release", "--config", 'profile.release.debug="line-tables-only"']
            )
        environment = stdlib_build_environment(env)
        environment["CARGO_JVM_BACKEND_PATH"] = str(ROOT)
        return run_command(command, env=environment)

    return run_command(
        cargo_build_command(test.directory / "Cargo.toml", release, jobs),
        env=stdlib_build_environment(env),
    )


def jar_path(test: TestCase, release: bool) -> Path:
    profile = "release" if release else "debug"
    profile_dir = TEST_TARGET_DIR / "jvm-unknown-jvm" / profile
    direct = profile_dir / f"{test.artifact_name}.jar"
    if direct.exists():
        return direct

    candidates = list((profile_dir / "deps").glob(f"{test.artifact_name}-*.jar"))
    if candidates:
        return max(candidates, key=lambda path: path.stat().st_mtime_ns)
    return direct