revng-revng/python/revng/api/manager.py

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

from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Dict, Generator, Iterable, List, Optional, Union

from revng.support import AnyPath, AnyPaths, to_iterable
from revng.support.collect import collect_files

from ._capi import ROOT, _api, ffi
from .analysis import Analysis
from .container import Container, ContainerIdentifier
from .error_list import Expected, run_with_el
from .exceptions import RevngException
from .invalidations import Invalidations, ResultWithInvalidations
from .kind import Kind
from .step import Step
from .string_map import StringMap
from .target import Target, TargetsList
from .utils import make_c_string, make_generator, make_python_string

INVALID_INDEX = 0xFFFFFFFFFFFFFFFF


class Manager:
    def __init__(
        self,
        workdir: Optional[AnyPath] = None,
        pipeline_files: Optional[AnyPaths] = None,
        flags: Iterable[str] = (),
    ):
        if pipeline_files is None:
            pipeline_files = collect_files(ROOT, ["share", "revng", "pipelines"], "*.yml")
        else:
            pipeline_files = to_iterable(pipeline_files)

        pipelines = [Path(f) for f in pipeline_files]
        assert len(pipelines) > 0, "Pipelines must have len > 0"
        assert all(x.is_file() for x in pipelines), "Pipeline files must exist"

        if workdir is None:
            self.temporary_workdir = TemporaryDirectory(prefix="revng-manager-workdir-")
            self.workdir = Path(self.temporary_workdir.name)
        else:
            self.workdir = Path(workdir)

        self.workdir.mkdir(parents=True, exist_ok=True)

        _flags = [make_c_string(s) for s in flags]
        _workdir = make_c_string(str(self.workdir))
        _pipelines_paths = [make_c_string(str(s.resolve())) for s in pipelines]

        # Ensures that the _manager property is always defined even if the API call fails
        self._manager = None

        self._manager = _api.rp_manager_create(
            len(_pipelines_paths),
            _pipelines_paths,
            len(_flags),
            _flags,
            _workdir,
        )

        assert self._manager, "Failed to instantiate manager"

    @property
    def uid(self) -> int:
        return int(ffi.cast("uintptr_t", self._manager))

    def save(self, destination_directory: Optional[Union[Path, str]] = None):
        if destination_directory is None:
            _dir_path = ffi.NULL
        else:
            dir_path = Path(destination_directory)
            if not dir_path.is_dir():
                dir_path.mkdir()
            _dir_path = make_c_string(str(dir_path.resolve()))
        return _api.rp_manager_save(self._manager, _dir_path)

    def save_context(self, destination_directory: Union[Path, str]):
        dest_dir = Path(destination_directory).resolve()
        _dest_dir = make_c_string(str(dest_dir))
        return _api.rp_manager_save_context(self._manager, _dest_dir)

    # Kind-related Functions

    @property
    def kinds_count(self) -> int:
        return _api.rp_manager_kinds_count(self._manager)

    def _get_kind_from_index(self, idx: int) -> Optional[Kind]:
        _kind = _api.rp_manager_get_kind(self._manager, idx)
        return Kind(_kind) if _kind != ffi.NULL else None

    def kinds(self) -> Generator[Kind, None, None]:
        return make_generator(self.kinds_count, self._get_kind_from_index)

    def kind_from_name(self, kind_name: str) -> Optional[Kind]:
        _kind_name = make_c_string(kind_name)
        kind = _api.rp_manager_get_kind_from_name(self._manager, _kind_name)
        return Kind(kind) if kind != ffi.NULL else None

    # Target-related Functions

    def deserialize_target(self, serialized_target: str) -> Optional[Target]:
        step_name, container_name, target = serialized_target.split("/", 2)
        target_path, kind_name = target.split(":", 1)
        targets = self.get_targets(step_name, container_name)
        return next((t for t in targets if t.serialize() == target), None)

    def _produce_target(
        self,
        step: Step,
        target: Union[Target, List[Target]],
        container: Container,
    ) -> Dict[str, str]:
        if isinstance(target, Target):
            targets = [
                target,
            ]
        else:
            targets = target
        _step = step._step
        _container = container._container
        product = _api.rp_manager_produce_targets(
            self._manager, len(targets), [t._target for t in targets], _step, _container
        )
        if not product:
            # TODO: we really should be able to provide a detailed error here
            raise RevngException("Failed to produce targets")

        return {t.serialize(): t.extract() for t in targets}

    def produce_target(
        self,
        step_name: str,
        target: Union[None, str, List[str]],
        container_name: Optional[str] = None,
        only_if_ready=False,
    ) -> Dict[str, str]:
        step = self.get_step(step_name)
        if step is None:
            raise RevngException(f"Invalid step {step_name}")

        if container_name is not None:
            container_identifier = self.get_container_with_name(container_name)
            if container_identifier is None:
                raise RevngException(f"Invalid container {container_name}")

            container = step.get_container(container_identifier)
            if container is None:
                raise RevngException(f"Step {step_name} does not use container {container_name}")
        else:
            container = step.get_artifacts_container()
            if container is None:
                raise RevngException(f"Step {step_name} does not have an artifacts container")

        if target is None:
            _targets = [
                "",
            ]
        elif isinstance(target, str):
            _targets = [
                target,
            ]
        else:
            _targets = target

        targets: List[Target] = []
        for _target_elem in _targets:
            if container_name is not None:
                targets.append(self.create_target(_target_elem, container))
            else:
                targets.append(self.create_target(_target_elem, container, step))

        if only_if_ready and any(not t.is_ready for t in targets):
            raise RevngException("Requested production of unready targets")

        product = self._produce_target(step, targets, container)
        if not product:
            # TODO: we really should be able to provide a detailed error here
            raise RevngException("Failed to produce target")
        return product

    def create_target(
        self,
        target_path: str,
        container: Container,
        step: Optional[Step] = None,
    ) -> Target:
        if step is not None:
            path = target_path
            kind = step.get_artifacts_kind()
            if kind is None:
                raise RevngException("Step does not have an artifacts kind")
        else:
            components = target_path.split(":")
            path, kind_name = ":".join(components[:-1]), components[-1]

            kind = self.kind_from_name(kind_name)
            if kind is None:
                raise RevngException("Invalid kind")

        path_components = path.split("/") if path != "" else []

        if kind.rank is not None and len(path_components) != kind.rank.depth:
            raise RevngException("Path components need to equal kind rank")

        target = Target.create(kind, container, path_components)
        if target is None:
            raise RevngException("Invalid target")
        return target

    def get_targets_list(self, container: Container) -> Optional[TargetsList]:
        targets_list = _api.rp_manager_get_container_targets_list(
            self._manager, container._container
        )

        return TargetsList(targets_list, container) if targets_list != ffi.NULL else None

    # Container-related functions

    def container_path(self, step_name: str, container_name: str) -> Optional[str]:
        _step_name = make_c_string(step_name)
        _container_name = make_c_string(container_name)
        _path = _api.rp_manager_create_container_path(self._manager, _step_name, _container_name)
        if not _path:
            return None
        return make_python_string(_path)

    @property
    def containers_count(self) -> int:
        return _api.rp_manager_containers_count(self._manager)

    def containers(self) -> Generator[ContainerIdentifier, None, None]:
        return make_generator(self.containers_count, self._get_container_identifier)

    def get_container_with_name(self, name) -> Optional[ContainerIdentifier]:
        for container in self.containers():
            if container.name == name:
                return container
        return None

    def deserialize_container(self, step: Step, container_name: str, content: Union[str, bytes]):
        _content = make_c_string(content)
        return _api.rp_manager_container_deserialize(
            self._manager, step._step, make_c_string(container_name), _content, len(_content)
        )

    def _get_container_identifier(self, idx: int) -> Optional[ContainerIdentifier]:
        _container_identifier = _api.rp_manager_get_container_identifier(self._manager, idx)
        if _container_identifier != ffi.NULL:
            return ContainerIdentifier(_container_identifier)
        return None

    # Step-related functions

    @property
    def steps_count(self) -> int:
        return _api.rp_manager_steps_count(self._manager)

    def get_step(self, step_name: str) -> Optional[Step]:
        step_index = self._step_name_to_index(step_name)
        if step_index is None:
            return None
        return self._get_step_from_index(step_index)

    def steps(self) -> Generator[Step, None, None]:
        return make_generator(self.steps_count, self._get_step_from_index)

    def _step_name_to_index(self, name: str) -> Optional[int]:
        _name = make_c_string(name)
        index = _api.rp_manager_step_name_to_index(self._manager, _name)
        return index if index != INVALID_INDEX else None

    def _get_step_from_index(self, idx: int) -> Optional[Step]:
        step = _api.rp_manager_get_step(self._manager, idx)
        return Step(step) if step != ffi.NULL else None

    # Utility target functions

    def get_targets(self, step_name: str, container_name: str) -> List[Target]:
        step = self.get_step(step_name)
        if step is None:
            raise RevngException("Invalid step name")

        container_identifier = self.get_container_with_name(container_name)
        if container_identifier is None:
            raise RevngException("Invalid container name")

        container = step.get_container(container_identifier)
        if container is None:
            raise RevngException(f"Step {step_name} does not use container {container_name}")

        targets_list = self.get_targets_list(container)
        return list(targets_list.targets()) if targets_list is not None else []

    def get_targets_from_step(self, step_name: str) -> Dict[str, List[Target]]:
        step = self.get_step(step_name)
        if step is None:
            raise RevngException("Invalid step name")

        containers = []
        for container_id in self.containers():
            containers.append(step.get_container(container_id))

        ret = {}
        for container in containers:
            if container is not None:
                targets = self.get_targets(step_name, container.name)
                ret[container.name] = targets
        return ret

    def get_all_targets(self) -> Dict[str, Dict[str, List[Target]]]:
        targets: Dict[str, Dict[str, List[Target]]] = {}
        container_ids = list(self.containers())
        for step in self.steps():
            targets[step.name] = {}
            containers = [step.get_container(cid) for cid in container_ids]
            for container in [c for c in containers if c is not None]:
                target_list = self.get_targets_list(container)
                if target_list is not None:
                    target_dicts = list(target_list.targets())
                    targets[step.name][container.name] = target_dicts
        return targets

    # Analysis handling

    def run_analysis(
        self,
        step_name: str,
        analysis_name: str,
        target_mapping: Dict[str, List[str]],
        options: Dict[str, str] = {},
    ) -> ResultWithInvalidations[Dict[str, str]]:
        step = self.get_step(step_name)
        if step is None:
            raise RevngException(f"Invalid step {step_name}")

        analysis = next((a for a in step.analyses() if a.name == analysis_name), None)
        if analysis is None:
            raise RevngException(f"Invalid Analysis {analysis_name}")

        concrete_target_mapping = {}
        for container_name, target_list in target_mapping.items():
            container_identifier = self.get_container_with_name(container_name)
            if container_identifier is None:
                raise RevngException(f"Invalid container {container_name}")

            container = step.get_container(container_identifier)
            if container is None:
                raise RevngException(f"Step {step_name} does not use container {container_name}")

            concrete_targets = [
                self.create_target(target, container, None) for target in target_list
            ]

            analysis_argument = next(a for a in analysis.arguments() if a.name == container_name)
            for target in concrete_targets:
                if target.kind is None or target.kind.name not in [
                    k.name for k in analysis_argument.acceptable_kinds()
                ]:
                    raise RevngException(
                        f"Wrong kind for analysis: found '{target.kind}', "
                        + f"expected: {[a.name for a in analysis_argument.acceptable_kinds()]}"
                    )
            concrete_target_mapping[container] = concrete_targets

        options_map = StringMap(options)
        analysis_result = self._run_analysis(step, analysis, concrete_target_mapping, options_map)
        if analysis_result.result is None:
            raise RevngException("Failed to run analysis")
        return analysis_result  # type: ignore

    def _run_analysis(
        self,
        step: Step,
        analysis: Analysis,
        target_mapping: Dict[Container, List[Target]],
        options: StringMap,
    ) -> ResultWithInvalidations[Optional[Dict[str, str]]]:
        first_key = list(target_mapping.keys())[0]
        targets = target_mapping[first_key]
        invalidations = Invalidations()
        result = _api.rp_manager_run_analysis(
            self._manager,
            len(targets),
            [t._target for t in targets],
            make_c_string(step.name),
            make_c_string(analysis.name),
            first_key._container,
            invalidations._invalidations,
            options._string_map,
        )

        if result != ffi.NULL and not _api.rp_diff_map_is_empty(result):
            self.save()
        return ResultWithInvalidations(
            self.parse_diff_map(result) if result != ffi.NULL else None, invalidations
        )

    def run_all_analyses(
        self, options: Dict[str, str] = {}
    ) -> ResultWithInvalidations[Optional[Dict[str, str]]]:
        options_map = StringMap(options)
        invalidations = Invalidations()
        result = _api.rp_manager_run_all_analyses(
            self._manager, invalidations._invalidations, options_map._string_map
        )
        if result != ffi.NULL and not _api.rp_diff_map_is_empty(result):
            self.save()
        return ResultWithInvalidations(
            self.parse_diff_map(result) if result != ffi.NULL else None, invalidations
        )

    def parse_diff_map(self, diff_map) -> Dict[str, str]:
        result = {}
        for global_name in self.globals_list():
            _diff_value = _api.rp_diff_map_get_diff(diff_map, make_c_string(global_name))
            if _diff_value != ffi.NULL:
                result[global_name] = make_python_string(_diff_value)
        return result

    # Global Handling & misc.

    def get_global(self, name) -> str:
        _name = make_c_string(name)
        _out = _api.rp_manager_create_global_copy(self._manager, _name)
        return make_python_string(_out)

    def set_global(self, name, content) -> ResultWithInvalidations[Expected[bool]]:
        _name = make_c_string(name)
        _content = make_c_string(content)
        invalidations = Invalidations()
        res = run_with_el(
            _api.rp_manager_set_global, self._manager, _content, _name, invalidations._invalidations
        )
        return ResultWithInvalidations(res, invalidations)

    def verify_global(self, name, content) -> Expected[bool]:
        _name = make_c_string(name)
        _content = make_c_string(content)
        return run_with_el(_api.rp_manager_verify_global, self._manager, _content, _name)

    def apply_diff(self, name, diff) -> ResultWithInvalidations[Expected[bool]]:
        _name = make_c_string(name)
        _diff = make_c_string(diff)
        invalidations = Invalidations()
        res = run_with_el(
            _api.rp_manager_apply_diff, self._manager, _diff, _name, invalidations._invalidations
        )
        return ResultWithInvalidations(res, invalidations)

    def verify_diff(self, name, diff) -> Expected[bool]:
        _name = make_c_string(name)
        _diff = make_c_string(diff)
        return run_with_el(_api.rp_manager_verify_diff, self._manager, _diff, _name)

    def get_model(self) -> str:
        return self.get_global("model.yml")

    def globals_count(self) -> int:
        return _api.rp_manager_get_globals_count(self._manager)

    def _get_global_from_index(self, idx: int) -> str:
        _name = _api.rp_manager_get_global_name(self._manager, idx)
        return make_python_string(_name)

    def globals_list(self) -> Generator[str, None, None]:
        return make_generator(self.globals_count(), self._get_global_from_index)

    def set_input(self, container_name: str, content: bytes, _key=None):
        step = self.get_step("begin")
        if step is None:
            raise RevngException('Step "begin" not found')

        success = self.deserialize_container(step, container_name, content)

        if not success:
            raise RevngException(
                f"Failed loading user provided input for container {container_name}"
            )