Auto merge of #135286 - matthiaskrgr:rollup-sxuq1nh, r=matthiaskrgr

Rollup of 3 pull requests

Successful merges:

 - #134898 (Make it easier to run CI jobs locally)
 - #135195 (Make `lit_to_mir_constant` and `lit_to_const` infallible)
 - #135261 (Account for identity substituted items in symbol mangling)

r? `@ghost`
`@rustbot` modify labels: rollup

1 files changed, 318 insertions, 0 deletions

diff --git a/src/ci/github-actions/ci.py b/src/ci/github-actions/ci.py
new file mode 100755
index 00000000000..97539cf1cdb
--- /dev/null
+++ b/src/ci/github-actions/ci.py
@@ -0,0 +1,318 @@
+#!/usr/bin/env python3
+
+"""
+This script contains CI functionality.
+It can be used to generate a matrix of jobs that should
+be executed on CI, or run a specific CI job locally.
+
+It reads job definitions from `src/ci/github-actions/jobs.yml`.
+"""
+
+import argparse
+import dataclasses
+import json
+import logging
+import os
+import re
+import subprocess
+import typing
+from pathlib import Path
+from typing import List, Dict, Any, Optional
+
+import yaml
+
+CI_DIR = Path(__file__).absolute().parent.parent
+JOBS_YAML_PATH = Path(__file__).absolute().parent / "jobs.yml"
+
+Job = Dict[str, Any]
+
+
+def add_job_properties(jobs: List[Dict], prefix: str) -> List[Job]:
+    """
+    Modify the `name` attribute of each job, based on its base name and the given `prefix`.
+    Add an `image` attribute to each job, based on its image.
+    """
+    modified_jobs = []
+    for job in jobs:
+        # Create a copy of the `job` dictionary to avoid modifying `jobs`
+        new_job = dict(job)
+        new_job["image"] = get_job_image(new_job)
+        new_job["name"] = f"{prefix} - {new_job['name']}"
+        modified_jobs.append(new_job)
+    return modified_jobs
+
+
+def add_base_env(jobs: List[Job], environment: Dict[str, str]) -> List[Job]:
+    """
+    Prepends `environment` to the `env` attribute of each job.
+    The `env` of each job has higher precedence than `environment`.
+    """
+    modified_jobs = []
+    for job in jobs:
+        env = environment.copy()
+        env.update(job.get("env", {}))
+
+        new_job = dict(job)
+        new_job["env"] = env
+        modified_jobs.append(new_job)
+    return modified_jobs
+
+
+@dataclasses.dataclass
+class PRRunType:
+    pass
+
+
+@dataclasses.dataclass
+class TryRunType:
+    custom_jobs: List[str]
+
+
+@dataclasses.dataclass
+class AutoRunType:
+    pass
+
+
+WorkflowRunType = typing.Union[PRRunType, TryRunType, AutoRunType]
+
+
+@dataclasses.dataclass
+class GitHubCtx:
+    event_name: str
+    ref: str
+    repository: str
+    commit_message: Optional[str]
+
+
+def get_custom_jobs(ctx: GitHubCtx) -> List[str]:
+    """
+    Tries to parse names of specific CI jobs that should be executed in the form of
+    try-job: <job-name>
+    from the commit message of the passed GitHub context.
+    """
+    if ctx.commit_message is None:
+        return []
+
+    regex = re.compile(r"^try-job: (.*)", re.MULTILINE)
+    jobs = []
+    for match in regex.finditer(ctx.commit_message):
+        jobs.append(match.group(1))
+    return jobs
+
+
+def find_run_type(ctx: GitHubCtx) -> Optional[WorkflowRunType]:
+    if ctx.event_name == "pull_request":
+        return PRRunType()
+    elif ctx.event_name == "push":
+        try_build = ctx.ref in (
+            "refs/heads/try",
+            "refs/heads/try-perf",
+            "refs/heads/automation/bors/try",
+        )
+
+        # Unrolled branch from a rollup for testing perf
+        # This should **not** allow custom try jobs
+        is_unrolled_perf_build = ctx.ref == "refs/heads/try-perf"
+
+        if try_build:
+            custom_jobs = []
+            if not is_unrolled_perf_build:
+                custom_jobs = get_custom_jobs(ctx)
+            return TryRunType(custom_jobs=custom_jobs)
+
+        if ctx.ref == "refs/heads/auto":
+            return AutoRunType()
+
+    return None
+
+
+def calculate_jobs(run_type: WorkflowRunType, job_data: Dict[str, Any]) -> List[Job]:
+    if isinstance(run_type, PRRunType):
+        return add_base_env(
+            add_job_properties(job_data["pr"], "PR"), job_data["envs"]["pr"]
+        )
+    elif isinstance(run_type, TryRunType):
+        jobs = job_data["try"]
+        custom_jobs = run_type.custom_jobs
+        if custom_jobs:
+            if len(custom_jobs) > 10:
+                raise Exception(
+                    f"It is only possible to schedule up to 10 custom jobs, "
+                    f"received {len(custom_jobs)} jobs"
+                )
+
+            jobs = []
+            unknown_jobs = []
+            for custom_job in custom_jobs:
+                job = [j for j in job_data["auto"] if j["name"] == custom_job]
+                if not job:
+                    unknown_jobs.append(custom_job)
+                    continue
+                jobs.append(job[0])
+            if unknown_jobs:
+                raise Exception(
+                    f"Custom job(s) `{unknown_jobs}` not found in auto jobs"
+                )
+
+        return add_base_env(add_job_properties(jobs, "try"), job_data["envs"]["try"])
+    elif isinstance(run_type, AutoRunType):
+        return add_base_env(
+            add_job_properties(job_data["auto"], "auto"), job_data["envs"]["auto"]
+        )
+
+    return []
+
+
+def skip_jobs(jobs: List[Dict[str, Any]], channel: str) -> List[Job]:
+    """
+    Skip CI jobs that are not supposed to be executed on the given `channel`.
+    """
+    return [j for j in jobs if j.get("only_on_channel", channel) == channel]
+
+
+def get_github_ctx() -> GitHubCtx:
+    event_name = os.environ["GITHUB_EVENT_NAME"]
+
+    commit_message = None
+    if event_name == "push":
+        commit_message = os.environ["COMMIT_MESSAGE"]
+    return GitHubCtx(
+        event_name=event_name,
+        ref=os.environ["GITHUB_REF"],
+        repository=os.environ["GITHUB_REPOSITORY"],
+        commit_message=commit_message,
+    )
+
+
+def format_run_type(run_type: WorkflowRunType) -> str:
+    if isinstance(run_type, PRRunType):
+        return "pr"
+    elif isinstance(run_type, AutoRunType):
+        return "auto"
+    elif isinstance(run_type, TryRunType):
+        return "try"
+    else:
+        raise AssertionError()
+
+
+def get_job_image(job: Job) -> str:
+    """
+    By default, the Docker image of a job is based on its name.
+    However, it can be overridden by its IMAGE environment variable.
+    """
+    env = job.get("env", {})
+    # Return the IMAGE environment variable if it exists, otherwise return the job name
+    return env.get("IMAGE", job["name"])
+
+
+def is_linux_job(job: Job) -> bool:
+    return "ubuntu" in job["os"]
+
+
+def find_linux_job(job_data: Dict[str, Any], job_name: str, pr_jobs: bool) -> Job:
+    candidates = job_data["pr"] if pr_jobs else job_data["auto"]
+    jobs = [job for job in candidates if job.get("name") == job_name]
+    if len(jobs) == 0:
+        available_jobs = "\n".join(
+            sorted(job["name"] for job in candidates if is_linux_job(job))
+        )
+        raise Exception(f"""Job `{job_name}` not found in {'pr' if pr_jobs else 'auto'} jobs.
+The following jobs are available:
+{available_jobs}""")
+    assert len(jobs) == 1
+
+    job = jobs[0]
+    if not is_linux_job(job):
+        raise Exception("Only Linux jobs can be executed locally")
+    return job
+
+
+def run_workflow_locally(job_data: Dict[str, Any], job_name: str, pr_jobs: bool):
+    DOCKER_DIR = Path(__file__).absolute().parent.parent / "docker"
+
+    job = find_linux_job(job_data, job_name=job_name, pr_jobs=pr_jobs)
+
+    custom_env = {}
+    # Replicate src/ci/scripts/setup-environment.sh
+    # Adds custom environment variables to the job
+    if job_name.startswith("dist-"):
+        if job_name.endswith("-alt"):
+            custom_env["DEPLOY_ALT"] = "1"
+        else:
+            custom_env["DEPLOY"] = "1"
+    custom_env.update({k: str(v) for (k, v) in job.get("env", {}).items()})
+
+    args = [str(DOCKER_DIR / "run.sh"), get_job_image(job)]
+    env_formatted = [f"{k}={v}" for (k, v) in sorted(custom_env.items())]
+    print(f"Executing `{' '.join(env_formatted)} {' '.join(args)}`")
+
+    env = os.environ.copy()
+    env.update(custom_env)
+
+    subprocess.run(args, env=env)
+
+
+def calculate_job_matrix(job_data: Dict[str, Any]):
+    github_ctx = get_github_ctx()
+
+    run_type = find_run_type(github_ctx)
+    logging.info(f"Job type: {run_type}")
+
+    with open(CI_DIR / "channel") as f:
+        channel = f.read().strip()
+
+    jobs = []
+    if run_type is not None:
+        jobs = calculate_jobs(run_type, job_data)
+    jobs = skip_jobs(jobs, channel)
+
+    if not jobs:
+        raise Exception("Scheduled job list is empty, this is an error")
+
+    run_type = format_run_type(run_type)
+
+    logging.info(f"Output:\n{yaml.dump(dict(jobs=jobs, run_type=run_type), indent=4)}")
+    print(f"jobs={json.dumps(jobs)}")
+    print(f"run_type={run_type}")
+
+
+def create_cli_parser():
+    parser = argparse.ArgumentParser(
+        prog="ci.py", description="Generate or run CI workflows"
+    )
+    subparsers = parser.add_subparsers(
+        help="Command to execute", dest="command", required=True
+    )
+    subparsers.add_parser(
+        "calculate-job-matrix",
+        help="Generate a matrix of jobs that should be executed in CI",
+    )
+    run_parser = subparsers.add_parser(
+        "run-local", help="Run a CI jobs locally (on Linux)"
+    )
+    run_parser.add_argument(
+        "job_name",
+        help="CI job that should be executed. By default, a merge (auto) "
+        "job with the given name will be executed",
+    )
+    run_parser.add_argument(
+        "--pr", action="store_true", help="Run a PR job instead of an auto job"
+    )
+    return parser
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+
+    with open(JOBS_YAML_PATH) as f:
+        data = yaml.safe_load(f)
+
+    parser = create_cli_parser()
+    args = parser.parse_args()
+
+    if args.command == "calculate-job-matrix":
+        calculate_job_matrix(data)
+    elif args.command == "run-local":
+        run_workflow_locally(data, args.job_name, args.pr)
+    else:
+        raise Exception(f"Unknown command {args.command}")