retroSoC/scripts/development_environment.py

#!/usr/bin/env python3
"""Create and verify the lock-pinned open-source development environment."""

from __future__ import annotations

import argparse
import hashlib
import json
import shlex
import shutil
import subprocess
import sys
import venv
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from scripts.dependency_lock import DEFAULT_LOCK, lock_digest, load_lock  # noqa: E402
from scripts.install_toolchain import install  # noqa: E402
from scripts.setup_helpers import atomic_write  # noqa: E402


PLATFORM = "ubuntu-22.04"
DEFAULT_TOOLS = (
    "verilator",
    "verible",
    "sv2v",
    "iverilog",
    "yosys",
    "sby",
    "bitwuzla",
    "opensta",
    "openocd",
    "sbt",
    "riscv_gnu",
)
TOOL_EXECUTABLES = {
    "verilator": ("verilator",),
    "verible": ("verible-verilog-format",),
    "sv2v": ("sv2v",),
    "iverilog": ("iverilog", "vvp"),
    "yosys": ("yosys", "yosys-smtbmc"),
    "sby": ("sby",),
    "bitwuzla": ("bitwuzla",),
    "opensta": ("sta",),
    "openocd": ("openocd",),
    "sbt": ("sbt",),
    "riscv_gnu": ("riscv32-unknown-elf-gcc", "riscv32-unknown-elf-gdb"),
}
PYTHON_IMPORTS = ("pytest", "ruff", "yamllint")
PYTHON_EXECUTABLES = ("mbake",)
SYSTEM_EXECUTABLES = (
    "cc",
    "ccache",
    "clang-format-14",
    "git",
    "java",
    "make",
    "mold",
    "numactl",
    "perl",
)
STAMP_NAME = "development-environment.json"
ACTIVATE_NAME = "activate.sh"


def file_digest(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def requirement_paths(root: Path) -> tuple[Path, Path]:
    return root / "requirements/build.txt", root / "requirements/ci.txt"


def parse_tools(values: list[str] | None, available: dict[str, Any]) -> tuple[str, ...]:
    tools = tuple(values or DEFAULT_TOOLS)
    if not tools:
        raise ValueError("at least one tool must be selected")
    if len(tools) != len(set(tools)):
        raise ValueError("selected tools must be unique")
    unknown = sorted(set(tools) - set(available))
    if unknown:
        raise ValueError(f"selected tools are not locked for {PLATFORM}: {', '.join(unknown)}")
    return tools


def tool_paths(cache: Path, tools: tuple[str, ...], available: dict[str, Any]) -> list[Path]:
    return [
        cache / "toolchains" / f"{name}-{available[name]['version']}" / available[name]["path"]
        for name in tools
    ]


def environment_paths(cache: Path, tools: tuple[str, ...], available: dict[str, Any]) -> list[Path]:
    return [cache / "venv" / "bin", *tool_paths(cache, tools, available)]


def stamp_data(
    root: Path, cache: Path, tools: tuple[str, ...], lock: dict[str, Any], lock_path: Path
) -> dict[str, Any]:
    build_requirements, ci_requirements = requirement_paths(root)
    if not build_requirements.is_file() or not ci_requirements.is_file():
        raise ValueError("development Python requirement files are missing")
    available = lock["toolchains"][PLATFORM]
    return {
        "schema_version": 1,
        "cache": str(cache),
        "dependency_lock_sha256": lock_digest(lock_path),
        "python_requirements": {
            str(build_requirements.relative_to(root)): file_digest(build_requirements),
            str(ci_requirements.relative_to(root)): file_digest(ci_requirements),
        },
        "tools": {name: available[name]["version"] for name in tools},
    }


def read_stamp(path: Path) -> dict[str, Any] | None:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    return value if isinstance(value, dict) else None


def has_environment(
    cache: Path, tools: tuple[str, ...], available: dict[str, Any], expected_stamp: dict[str, Any]
) -> bool:
    if read_stamp(cache / STAMP_NAME) != expected_stamp:
        return False
    python = cache / "venv" / "bin" / "python"
    if not python.is_file():
        return False
    result = subprocess.run(
        [str(python), "-c", "import " + ", ".join(PYTHON_IMPORTS)],
        check=False,
        capture_output=True,
        text=True,
    )
    if result.returncode:
        return False
    if any(not (cache / "venv" / "bin" / executable).is_file() for executable in PYTHON_EXECUTABLES):
        return False
    for path in tool_paths(cache, tools, available):
        if not path.is_dir():
            return False
    return True


def render_activation(cache: Path, paths: list[Path]) -> str:
    quoted_path = ":".join(shlex.quote(str(path.resolve())) for path in paths)
    return "\n".join(
        (
            "# Generated by scripts/development_environment.py; do not edit.",
            f"export RETROSOC_DEVELOPMENT_CACHE={shlex.quote(str(cache.resolve()))}",
            f"export VIRTUAL_ENV={shlex.quote(str((cache / 'venv').resolve()))}",
            f"export PATH={quoted_path}:$PATH",
            "",
        )
    )


def write_activation(path: Path, content: str) -> None:
    atomic_write(path, content)
    path.chmod(0o644)


def write_stamp(path: Path, content: dict[str, Any]) -> None:
    atomic_write(path, json.dumps(content, indent=2, sort_keys=True) + "\n")
    path.chmod(0o644)


def run(command: list[str]) -> None:
    print("+ " + " ".join(shlex.quote(argument) for argument in command), flush=True)
    subprocess.run(command, check=True)


def create_virtualenv(root: Path, cache: Path) -> None:
    python = cache / "venv" / "bin" / "python"
    if not python.is_file():
        venv.EnvBuilder(with_pip=True, clear=False).create(cache / "venv")
    build_requirements, ci_requirements = requirement_paths(root)
    run(
        [
            str(python),
            "-m",
            "pip",
            "install",
            "--disable-pip-version-check",
            "--require-hashes",
            "--requirement",
            str(build_requirements),
            "--requirement",
            str(ci_requirements),
        ]
    )


def bootstrap(args: argparse.Namespace) -> int:
    root = args.root.resolve()
    cache = args.cache.resolve()
    lock_path = args.lock.resolve()
    lock = load_lock(lock_path)
    available = lock["toolchains"][PLATFORM]
    tools = parse_tools(args.tool, available)
    expected_stamp = stamp_data(root, cache, tools, lock, lock_path)

    if not has_environment(cache, tools, available, expected_stamp):
        for name in tools:
            install(name, available[name], cache, update=args.update)
        create_virtualenv(root, cache)
        write_stamp(cache / STAMP_NAME, expected_stamp)
    activation = args.output.resolve() if args.output else cache / ACTIVATE_NAME
    write_activation(
        activation,
        render_activation(cache, environment_paths(cache, tools, available)),
    )
    print(f"development environment: {activation}")
    return check(args, quiet=False)


def check(args: argparse.Namespace, *, quiet: bool) -> int:
    cache = args.cache.resolve()
    lock_path = args.lock.resolve()
    lock = load_lock(lock_path)
    available = lock["toolchains"][PLATFORM]
    tools = parse_tools(args.tool, available)
    expected_stamp = stamp_data(args.root.resolve(), cache, tools, lock, lock_path)
    failures: list[str] = []

    if not has_environment(cache, tools, available, expected_stamp):
        failures.append("development environment stamp or installed paths do not match the lock")
    for name in tools:
        for executable in TOOL_EXECUTABLES[name]:
            if not (tool_paths(cache, (name,), available)[0] / executable).is_file():
                failures.append(f"missing executable: {executable}")
    python = cache / "venv" / "bin" / "python"
    if python.is_file():
        result = subprocess.run(
            [str(python), "-c", "import " + ", ".join(PYTHON_IMPORTS)],
            check=False,
            capture_output=True,
            text=True,
        )
        if result.returncode:
            failures.append("locked Python development packages are unavailable")
    else:
        failures.append("missing virtualenv Python executable")
    for executable in SYSTEM_EXECUTABLES:
        if shutil.which(executable) is None:
            failures.append(f"missing system executable: {executable}")

    if not quiet:
        if failures:
            print("development environment check failed:")
            for failure in failures:
                print(f"  {failure}")
        else:
            print(f"development environment check passed: {cache}")
    return int(bool(failures))


def activate(args: argparse.Namespace) -> int:
    activation = args.output.resolve() if args.output else args.cache.resolve() / ACTIVATE_NAME
    if not activation.is_file():
        raise ValueError(f"activation script does not exist: {activation}")
    print(f"source {shlex.quote(str(activation))}")
    return 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=ROOT)
    parser.add_argument("--lock", type=Path, default=DEFAULT_LOCK)
    parser.add_argument("--cache", type=Path, default=ROOT / ".cache/retrosoc/development")
    parser.add_argument("--output", type=Path)
    parser.add_argument("--tool", action="append")
    parser.add_argument("--update", action="store_true")
    subparsers = parser.add_subparsers(dest="command", required=True)
    subparsers.add_parser("bootstrap")
    subparsers.add_parser("check")
    subparsers.add_parser("activate")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        if args.command == "bootstrap":
            return bootstrap(args)
        if args.command == "check":
            return check(args, quiet=False)
        return activate(args)
    except (OSError, RuntimeError, ValueError, subprocess.CalledProcessError) as error:
        print(f"development environment: {error}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())