revng-revng/python/revng/support/__init__.py

#
# This file is distributed under the MIT License. See LICENSE.md for details.
#

import mmap
import os
import re
import sys
import tarfile
from collections.abc import Mapping
from contextlib import contextmanager
from io import BytesIO, IOBase
from pathlib import Path
from shutil import which
from typing import IO, Generator, Generic, Iterable, List, Optional, TypeVar, Union, cast

from llvmcpy import LLVMCPy

T = TypeVar("T")
AnyPath = Union[str, Path]
SingleOrIterable = Union[T, Iterable[T]]
AnyPaths = SingleOrIterable[AnyPath]


def log_error(msg: str):
    sys.stderr.write(msg + "\n")


def read_lines(path: Path) -> List[str]:
    if not path.exists():
        return []
    return path.read_text().strip().split("\n")


# We assume this file is in <root>/lib/python$VERSION/site-packages/revng/
def get_root() -> Path:
    return (Path(__file__) / "../../../../../../").resolve()


additional_bin_paths = read_lines(get_root() / "share/revng/additional-bin-paths")


def collect_files_gen(
    search_prefixes: AnyPaths, path_components: Iterable[AnyPath], pattern: str
) -> Generator[str, None, None]:
    if isinstance(search_prefixes, (str, Path)):
        prefixes: Iterable[AnyPath] = (search_prefixes,)
    else:
        prefixes = search_prefixes

    already_found = set()
    for prefix in prefixes:
        analyses_path = Path(prefix).joinpath(*path_components)
        if not analyses_path.is_dir():
            continue

        # Enumerate all the libraries containing analyses
        for library in analyses_path.glob(pattern):
            if library.is_file() and library.name not in already_found:
                already_found.add(library.name)
                yield str(library.resolve())


def collect_files(
    search_prefixes: AnyPaths, path_components: Iterable[AnyPath], pattern: str
) -> List[str]:
    return list(collect_files_gen(search_prefixes, path_components, pattern))


def collect_one(
    search_prefixes: AnyPaths, path_components: Iterable[AnyPath], pattern: str
) -> Optional[str]:
    return next(collect_files_gen(search_prefixes, path_components, pattern), None)


def _get_command(command: str, search_prefixes: Iterable[str]) -> Optional[str]:
    for additional_bin_path in additional_bin_paths:
        for executable in collect_files(search_prefixes, [additional_bin_path], command):
            return executable

    path = which(command)
    if path is not None:
        return os.path.abspath(path)
    return None


def get_command(command: str, search_prefixes: Iterable[str]) -> str:
    path = _get_command(command, search_prefixes)
    if path is None:
        log_error(f'Couldn\'t find "{command}".')
        raise ValueError()
    return path


def get_llvmcpy():
    return LLVMCPy(_get_command("llvm-config", [get_root()]))


ToBytesInput = Union[str, bytes, IO[str], IO[bytes], mmap.mmap]


@contextmanager
def to_bytes(input_: ToBytesInput) -> Generator[bytes, None, None]:
    """Wrapper that allows to seamlessly treat strings, bytes and files as raw
    bytes. Strings will be encoded in utf-8. Files will be mmap-ed if possible
    otherwise their contents will be read in memory"""
    if isinstance(input_, mmap.mmap):
        yield cast(bytes, input_)
    elif not isinstance(input_, (str, bytes)):
        # Test to exclude things like stdin
        if input_.seekable():
            offset = input_.tell()
            # mmap throws an error if the file is empty, check for the file
            # size by seeking to the end
            size = input_.seek(0, os.SEEK_END)
            input_.seek(offset)
            if size == 0:
                yield b""
            else:
                with mmap.mmap(input_.fileno(), 0, access=mmap.ACCESS_READ, offset=offset) as mm:
                    yield cast(bytes, mm)
        else:
            result = input_.read()
            if isinstance(result, str):
                yield result.encode("utf-8")
            else:
                yield result
    else:
        if isinstance(input_, str):
            yield input_.encode("utf-8")
        else:
            yield input_


def get_example_binary_path() -> str:
    runtime_dir = get_root() / "share/revng/test/tests/runtime"
    for entry in runtime_dir.iterdir():
        if not entry.is_file():
            continue
        if re.match(r"^calc-x86-64-static-revng-qa\.compiled-[0-9a-f]{8}$", entry.name):
            return str(entry.resolve())
    raise ValueError("Could not find calc binary")


class TarDictionary(Mapping):
    """Class that allows wrapping a tar file and treat it as a dictionary, the
    entries are eagerly computed, the content of the files will be lazily
    returned when calling __getitem__."""

    def __init__(self, data: Union[IO[bytes], IOBase, bytes], mode="r"):
        fileobj: IO[bytes] | IOBase
        if isinstance(data, bytes):
            fileobj = BytesIO(data)
        else:
            fileobj = data

        self._tar_file: tarfile.TarFile = tarfile.open(fileobj=fileobj, mode=mode)
        self._keys = {m.name: m for m in self._tar_file.getmembers()}

    def __getitem__(self, key: str) -> bytes:
        if key not in self._keys:
            raise KeyError
        member = self._tar_file.extractfile(self._keys[key])
        assert member is not None
        return member.read()

    def __contains__(self, key) -> bool:
        return key in self._keys

    def __len__(self) -> int:
        return len(self._keys)

    def __iter__(self):
        return iter(self._keys)


class IgnoreDeepCopy(Generic[T]):
    """This class is a wrapper class that allows to have a value that will not
    be deep-copied when using the deepcopy library function.
    """

    def __init__(self, value: T):
        self.value: T = value

    def get(self) -> T:
        return self.value

    def __deepcopy__(self, memo):
        return self