llvm_tools/patch_utils.py

# Copyright 2022 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Provides patch utilities for PATCHES.json file handling."""

import collections
import contextlib
import dataclasses
import json
from pathlib import Path
import re
import subprocess
import sys
from typing import Any, Dict, IO, List, Optional, Tuple, Union


CHECKED_FILE_RE = re.compile(r"^checking file\s+(.*)$")
HUNK_FAILED_RE = re.compile(r"^Hunk #(\d+) FAILED at.*")
HUNK_HEADER_RE = re.compile(r"^@@\s+-(\d+),(\d+)\s+\+(\d+),(\d+)\s+@@")
HUNK_END_RE = re.compile(r"^--\s*$")
PATCH_SUBFILE_HEADER_RE = re.compile(r"^\+\+\+ [ab]/(.*)$")


@contextlib.contextmanager
def atomic_write(fp: Union[Path, str], mode="w", *args, **kwargs):
    """Write to a filepath atomically.

    This works by a temp file swap, created with a .tmp suffix in
    the same directory briefly until being renamed to the desired
    filepath.

    Args:
      fp: Filepath to open.
      mode: File mode; can be 'w', 'wb'. Default 'w'.
      *args: Passed to Path.open as nargs.
      **kwargs: Passed to Path.open as kwargs.

    Raises:
      ValueError when the mode is invalid.
    """
    if isinstance(fp, str):
        fp = Path(fp)
    if mode not in ("w", "wb"):
        raise ValueError(f"mode {mode} not accepted")
    temp_fp = fp.with_suffix(fp.suffix + ".tmp")
    try:
        with temp_fp.open(mode, *args, **kwargs) as f:
            yield f
    except:
        if temp_fp.is_file():
            temp_fp.unlink()
        raise
    temp_fp.rename(fp)


@dataclasses.dataclass
class Hunk:
    """Represents a patch Hunk."""

    hunk_id: int
    """Hunk ID for the current file."""
    orig_start: int
    orig_hunk_len: int
    patch_start: int
    patch_hunk_len: int
    patch_hunk_lineno_begin: int
    patch_hunk_lineno_end: Optional[int]


def parse_patch_stream(patch_stream: IO[str]) -> Dict[str, List[Hunk]]:
    """Parse a patch file-like into Hunks.

    Args:
      patch_stream: A IO stream formatted like a git patch file.

    Returns:
      A dictionary mapping filenames to lists of Hunks present
      in the patch stream.
    """

    current_filepath = None
    current_hunk_id = 0
    current_hunk = None
    out = collections.defaultdict(list)
    for lineno, line in enumerate(patch_stream.readlines()):
        subfile_header = PATCH_SUBFILE_HEADER_RE.match(line)
        if subfile_header:
            current_filepath = subfile_header.group(1)
            if not current_filepath:
                raise RuntimeError("Could not get file header in patch stream")
            # Need to reset the hunk id, as it's per-file.
            current_hunk_id = 0
            continue
        hunk_header = HUNK_HEADER_RE.match(line)
        if hunk_header:
            if not current_filepath:
                raise RuntimeError(
                    "Parsed hunk before file header in patch stream"
                )
            if current_hunk:
                # Already parsing a hunk
                current_hunk.patch_hunk_lineno_end = lineno
            current_hunk_id += 1
            current_hunk = Hunk(
                hunk_id=current_hunk_id,
                orig_start=int(hunk_header.group(1)),
                orig_hunk_len=int(hunk_header.group(2)),
                patch_start=int(hunk_header.group(3)),
                patch_hunk_len=int(hunk_header.group(4)),
                patch_hunk_lineno_begin=lineno + 1,
                patch_hunk_lineno_end=None,
            )
            out[current_filepath].append(current_hunk)
            continue
        if current_hunk and HUNK_END_RE.match(line):
            current_hunk.patch_hunk_lineno_end = lineno
    return out


def parse_failed_patch_output(text: str) -> Dict[str, List[int]]:
    current_file = None
    failed_hunks = collections.defaultdict(list)
    for eline in text.split("\n"):
        checked_file_match = CHECKED_FILE_RE.match(eline)
        if checked_file_match:
            current_file = checked_file_match.group(1)
            continue
        failed_match = HUNK_FAILED_RE.match(eline)
        if failed_match:
            if not current_file:
                raise ValueError("Input stream was not parsable")
            hunk_id = int(failed_match.group(1))
            failed_hunks[current_file].append(hunk_id)
    return failed_hunks


@dataclasses.dataclass(frozen=True)
class PatchResult:
    """Result of a patch application."""

    succeeded: bool
    failed_hunks: Dict[str, List[Hunk]] = dataclasses.field(
        default_factory=dict
    )

    def __bool__(self):
        return self.succeeded

    def failure_info(self) -> str:
        if self.succeeded:
            return ""
        s = ""
        for file, hunks in self.failed_hunks.items():
            s += f"{file}:\n"
            for h in hunks:
                s += f"Lines {h.orig_start} to {h.orig_start + h.orig_hunk_len}\n"
            s += "--------------------\n"
        return s


@dataclasses.dataclass
class PatchEntry:
    """Object mapping of an entry of PATCHES.json."""

    workdir: Path
    """Storage location for the patches."""
    metadata: Optional[Dict[str, Any]]
    platforms: Optional[List[str]]
    rel_patch_path: str
    version_range: Optional[Dict[str, Optional[int]]]
    _parsed_hunks = None

    def __post_init__(self):
        if not self.workdir.is_dir():
            raise ValueError(f"workdir {self.workdir} is not a directory")

    @classmethod
    def from_dict(cls, workdir: Path, data: Dict[str, Any]):
        """Instatiate from a dictionary.

        Dictionary must have at least the following key:

          {
            'rel_patch_path': '<relative patch path to workdir>',
          }

        Returns:
          A new PatchEntry.
        """
        return cls(
            workdir,
            data.get("metadata"),
            data.get("platforms"),
            data["rel_patch_path"],
            data.get("version_range"),
        )

    def to_dict(self) -> Dict[str, Any]:
        out = {
            "metadata": self.metadata,
            "rel_patch_path": self.rel_patch_path,
            "version_range": self.version_range,
        }
        if self.platforms:
            # To match patch_sync, only serialized when
            # non-empty and non-null.
            out["platforms"] = sorted(self.platforms)
        return out

    def parsed_hunks(self) -> Dict[str, List[Hunk]]:
        # Minor caching here because IO is slow.
        if not self._parsed_hunks:
            with self.patch_path().open(encoding="utf-8") as f:
                self._parsed_hunks = parse_patch_stream(f)
        return self._parsed_hunks

    def patch_path(self) -> Path:
        return self.workdir / self.rel_patch_path

    def can_patch_version(self, svn_version: int) -> bool:
        """Is this patch meant to apply to `svn_version`?"""
        # Sometimes the key is there, but it's set to None.
        if not self.version_range:
            return True
        from_v = self.version_range.get("from") or 0
        until_v = self.version_range.get("until")
        if until_v is None:
            until_v = sys.maxsize
        return from_v <= svn_version < until_v

    def is_old(self, svn_version: int) -> bool:
        """Is this patch old compared to `svn_version`?"""
        if not self.version_range:
            return False
        until_v = self.version_range.get("until")
        # Sometimes the key is there, but it's set to None.
        if until_v is None:
            until_v = sys.maxsize
        return svn_version >= until_v

    def apply(
        self, root_dir: Path, extra_args: Optional[List[str]] = None
    ) -> PatchResult:
        """Apply a patch to a given directory."""
        if not extra_args:
            extra_args = []
        # Cmd to apply a patch in the src unpack path.
        abs_patch_path = self.patch_path().absolute()
        if not abs_patch_path.is_file():
            raise RuntimeError(
                f"Cannot apply: patch {abs_patch_path} is not a file"
            )
        cmd = [
            "patch",
            "-d",
            root_dir.absolute(),
            "-f",
            "-p1",
            "--no-backup-if-mismatch",
            "-i",
            abs_patch_path,
        ] + extra_args
        try:
            subprocess.run(
                cmd, encoding="utf-8", check=True, stdout=subprocess.PIPE
            )
        except subprocess.CalledProcessError as e:
            parsed_hunks = self.parsed_hunks()
            failed_hunks_id_dict = parse_failed_patch_output(e.stdout)
            failed_hunks = {}
            for path, failed_hunk_ids in failed_hunks_id_dict.items():
                hunks_for_file = parsed_hunks[path]
                failed_hunks[path] = [
                    hunk
                    for hunk in hunks_for_file
                    if hunk.hunk_id in failed_hunk_ids
                ]
            return PatchResult(succeeded=False, failed_hunks=failed_hunks)
        return PatchResult(succeeded=True)

    def test_apply(self, root_dir: Path) -> PatchResult:
        """Dry run applying a patch to a given directory."""
        return self.apply(root_dir, ["--dry-run"])

    def title(self) -> str:
        if not self.metadata:
            return ""
        return self.metadata.get("title", "")


@dataclasses.dataclass(frozen=True)
class PatchInfo:
    """Holds info for a round of patch applications."""

    # str types are legacy. Patch lists should
    # probably be PatchEntries,
    applied_patches: List[PatchEntry]
    failed_patches: List[PatchEntry]
    # Can be deleted once legacy code is removed.
    non_applicable_patches: List[str]
    # Can be deleted once legacy code is removed.
    disabled_patches: List[str]
    # Can be deleted once legacy code is removed.
    removed_patches: List[str]
    # Can be deleted once legacy code is removed.
    modified_metadata: Optional[str]

    def _asdict(self):
        return dataclasses.asdict(self)


def json_to_patch_entries(workdir: Path, json_fd: IO[str]) -> List[PatchEntry]:
    """Convert a json IO object to List[PatchEntry].

    Examples:
      >>> f = open('PATCHES.json')
      >>> patch_entries = json_to_patch_entries(Path(), f)
    """
    return [PatchEntry.from_dict(workdir, d) for d in json.load(json_fd)]


def _print_failed_patch(pe: PatchEntry, failed_hunks: Dict[str, List[Hunk]]):
    """Print information about a single failing PatchEntry.

    Args:
      pe: A PatchEntry that failed.
      failed_hunks: Hunks for pe which failed as dict:
        filepath: [Hunk...]
    """
    print(f"Could not apply {pe.rel_patch_path}: {pe.title()}", file=sys.stderr)
    for fp, hunks in failed_hunks.items():
        print(f"{fp}:", file=sys.stderr)
        for h in hunks:
            print(
                f"- {pe.rel_patch_path} "
                f"l:{h.patch_hunk_lineno_begin}...{h.patch_hunk_lineno_end}",
                file=sys.stderr,
            )


def apply_all_from_json(
    svn_version: int,
    llvm_src_dir: Path,
    patches_json_fp: Path,
    continue_on_failure: bool = False,
) -> PatchInfo:
    """Attempt to apply some patches to a given LLVM source tree.

    This relies on a PATCHES.json file to be the primary way
    the patches are applied.

    Args:
      svn_version: LLVM Subversion revision to patch.
      llvm_src_dir: llvm-project root-level source directory to patch.
      patches_json_fp: Filepath to the PATCHES.json file.
      continue_on_failure: Skip any patches which failed to apply,
        rather than throw an Exception.
    """
    with patches_json_fp.open(encoding="utf-8") as f:
        patches = json_to_patch_entries(patches_json_fp.parent, f)
    skipped_patches = []
    failed_patches = []
    applied_patches = []
    for pe in patches:
        applied, failed_hunks = apply_single_patch_entry(
            svn_version, llvm_src_dir, pe
        )
        if applied:
            applied_patches.append(pe)
            continue
        if failed_hunks is not None:
            if continue_on_failure:
                failed_patches.append(pe)
                continue
            else:
                _print_failed_patch(pe, failed_hunks)
                raise RuntimeError(
                    "failed to apply patch " f"{pe.patch_path()}: {pe.title()}"
                )
        # Didn't apply, didn't fail, it was skipped.
        skipped_patches.append(pe)
    return PatchInfo(
        non_applicable_patches=skipped_patches,
        applied_patches=applied_patches,
        failed_patches=failed_patches,
        disabled_patches=[],
        removed_patches=[],
        modified_metadata=None,
    )


def apply_single_patch_entry(
    svn_version: int,
    llvm_src_dir: Path,
    pe: PatchEntry,
    ignore_version_range: bool = False,
) -> Tuple[bool, Optional[Dict[str, List[Hunk]]]]:
    """Try to apply a single PatchEntry object.

    Returns:
      Tuple where the first element indicates whether the patch applied,
      and the second element is a faild hunk mapping from file name to lists of
      hunks (if the patch didn't apply).
    """
    # Don't apply patches outside of the version range.
    if not ignore_version_range and not pe.can_patch_version(svn_version):
        return False, None
    # Test first to avoid making changes.
    test_application = pe.test_apply(llvm_src_dir)
    if not test_application:
        return False, test_application.failed_hunks
    # Now actually make changes.
    application_result = pe.apply(llvm_src_dir)
    if not application_result:
        # This should be very rare/impossible.
        return False, application_result.failed_hunks
    return True, None


def is_git_dirty(git_root_dir: Path) -> bool:
    """Return whether the given git directory has uncommitted changes."""
    if not git_root_dir.is_dir():
        raise ValueError(f"git_root_dir {git_root_dir} is not a directory")
    cmd = ["git", "ls-files", "-m", "--other", "--exclude-standard"]
    return (
        subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            check=True,
            cwd=git_root_dir,
            encoding="utf-8",
        ).stdout
        != ""
    )


def clean_src_tree(src_path):
    """Cleans the source tree of the changes made in 'src_path'."""

    reset_src_tree_cmd = ["git", "-C", src_path, "reset", "HEAD", "--hard"]

    subprocess.run(reset_src_tree_cmd, check=True)

    clean_src_tree_cmd = ["git", "-C", src_path, "clean", "-fd"]

    subprocess.run(clean_src_tree_cmd, check=True)


@contextlib.contextmanager
def git_clean_context(git_root_dir: Path):
    """Cleans up a git directory when the context exits."""
    if is_git_dirty(git_root_dir):
        raise RuntimeError("Cannot setup clean context; git_root_dir is dirty")
    try:
        yield
    finally:
        clean_src_tree(git_root_dir)