source_manager.py

#
# Copyright (C) 2021 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Package to manage Rust source files when building a toolchain distributable.
"""

from pathlib import Path
import shutil
import subprocess
import sys
from typing import Any

from paths import OUT_PATH_PATCHS_LOG, PATCH_PATH, OUT_PATH_PATCHED_FILES
from utils import prepare_command


def ensure_unique_patch_numbers(patch_list: list[Path]) -> None:
    patches_seen = set()
    for filepath in patch_list:
        patch_num = filepath.name.split("-", 3)[1]
        if patch_num in patches_seen:
            raise RuntimeError(f"Duplicate patch number {patch_num} for {filepath.name}")
        patches_seen.add(patch_num)


def apply_patches(code_dir: Path, patch_list: list[Path], patch_abort: bool = False) -> None:
    count_padding = len(str(len(patch_list)))

    # We will overwrite the log file if it already existed.
    with OUT_PATH_PATCHS_LOG.open("w") as f:
        for idx, filepath in enumerate(patch_list):
            print(f"\33[2K\rApplying patch ({(idx + 1):>{count_padding}}/{len(patch_list)}): {filepath.name}", end="")

            command_list: list[str] = prepare_command([PATCH_PATH, "-p1", "-l", "-i", filepath])
            # We collect the stdout and stderr output and then print it to the
            # log so that when an error is encountered we can display the
            # relevent message and don't have to refer the user to the log.
            # The log is intended as a resource for investigating patch
            # drift.
            result = subprocess.run(command_list, cwd=code_dir, capture_output=True)
            if result.stdout or result.stderr:
                f.write(f"Patch {filepath.name}:\n")
                if result.stdout:
                    f.write(result.stdout.decode("UTF-8") + "\n")
                if result.stderr:
                    f.write(result.stderr.decode("UTF-8") + "\n")

            if result.returncode != 0 and patch_abort:
                print(f"\nBuild failed when applying patch {filepath}")
                print("If developing locally, try the --no-patch-abort flag")
                if result.stdout:
                    print("\nOutput (stdout):")
                    print(result.stdout.decode("UTF-8"))
                if result.stderr:
                    print("\nOutput (stderr):")
                    print(result.stderr.decode("UTF-8"))

                print("Failed")

                f.truncate()
                sys.exit(result.returncode)

        # Remove any possible leftovers from the previous log
        f.truncate()

    # If all patches applied cleanly we need to advance to the next line in the
    # terminal
    print()


def setup_files(input_dir: Path, output_dir: Path, patches_dir: Path, patch_abort: bool = False, repatch: bool = False) -> None:
    """Copy source and apply patches in a performant and fault-tolerant manner.

    This function creates a copy of the source directory and
    applies the patches contained in the patch directory.
    """

    patch_list = sorted(patches_dir.glob("**/rustc-*"), key=lambda p: p.name)
    ensure_unique_patch_numbers(patch_list)

    if not output_dir.parent.exists():
        output_dir.parent.mkdir(parents=True)

    if repatch and OUT_PATH_PATCHED_FILES.exists():
        with OUT_PATH_PATCHED_FILES.open("r") as previously_patched_files:
            for line in previously_patched_files.readlines():
                file = line.strip()
                print(f"\33[2K\rRestoring {file}", end="")
                if (input_dir / file).exists():
                    shutil.copy2(input_dir / file, output_dir / file)
                else:
                    (output_dir / file).unlink(missing_ok=True)
        print()
    else:
        print("Creating copy of Rust source")
        if output_dir.exists():
            shutil.rmtree(output_dir)
        shutil.copytree(input_dir, output_dir, ignore=shutil.ignore_patterns(".git"))
        (output_dir / "Android.mk").unlink()

    with OUT_PATH_PATCHED_FILES.open("w") as patched_files:
        for filepath in patch_list:
            with filepath.open("r") as patchfile:
                for line in patchfile:
                    if line.startswith('+++ b/'):
                        file = line.removeprefix('+++ b/').split()[0].strip()
                        patched_files.write(f"{file}\n")

    apply_patches(output_dir, patch_list, patch_abort=patch_abort)