revng-revng/python/revng/pypeline/pipeline.py

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

from __future__ import annotations

from dataclasses import dataclass, field
from graphlib import TopologicalSorter
from itertools import chain
from typing import Dict, Generator, Iterable, List, Mapping, Optional, Set

import yaml

from revng.pypeline.runner_context import RunnerContext
from revng.pypeline.utils import PypelineException

from .analysis import Analysis, AnalysisList
from .container import Container, ContainerDeclaration
from .graph import Graph
from .model import Model, ModelDiff, ReadOnlyModel
from .object import ObjectID, ObjectSet
from .pipeline_node import PipelineConfiguration, PipelineNode
from .schedule.schedule import Schedule
from .schedule.scheduled_task import PipeScheduledTask, SavepointScheduledTask, ScheduledTask
from .storage.storage_provider import InvalidatedObjects, ObjectsToInvalidate, SavePointsRange
from .storage.storage_provider import StorageProvider
from .task.pipe import Pipe
from .task.requests import Requests
from .task.savepoint import SavePoint
from .task.task import TaskArgumentAccess
from .utils.default_dict_from_key import DefaultDictFromKey
from .utils.logger import pypeline_logger
from .utils.registry import get_singleton


@dataclass(frozen=True, slots=True)
class ArtifactCategory:
    name: str
    show_by_default: bool

    def to_dict(self) -> dict:
        return {"name": self.name, "show_by_default": self.show_by_default}


@dataclass(frozen=True, slots=True)
class Artifact:
    """
    An artifact is a container in a certain point of the pipeline with some extra
    metadata, such as a name.

    It's designed to mark interesting results in the pipeline, so that users an
    easily obtain them.
    """

    name: str
    node: PipelineNode
    container: ContainerDeclaration
    category: ArtifactCategory
    description: Optional[str] = None
    filename: str | None = None
    # TODO: this should be a property of the pipe that inserts the locations,
    # but for now we define it statically as a property of the artifact.
    defined_locations: list[str] = field(default_factory=list, hash=False)
    preferred_artifacts: list[str] = field(default_factory=list, hash=False)

    def pipe_dependencies(self) -> list[str]:
        return sorted({p.name for p in self.node.pipe_dependencies})

    def to_dict(self) -> dict:
        """Convert the artifact to a dictionary representation."""
        if isinstance(self.node.task, SavePoint):
            assert self.node.savepoint_range is not None
            savepoint_id = self.node.savepoint_range.start
        else:
            savepoint_id = -1

        result = {
            "name": self.name,
            "savepoint_id": savepoint_id,
            "container": self.container.name,
            "pipe_dependencies": self.pipe_dependencies(),
            "category": self.category.to_dict(),
            "defined_locations": self.defined_locations,
            "preferred_artifacts": self.preferred_artifacts,
        }
        if self.filename is not None:
            result["filename"] = self.filename
        return result


@dataclass(frozen=True, slots=True)
class AnalysisBinding:
    """Allows to bind an analysis to a pipeline node."""

    analysis: Analysis
    bindings: tuple[ContainerDeclaration, ...]
    node: PipelineNode

    def to_dict(self) -> dict:
        """Convert the data into a dictionary representation."""
        return {
            "name": self.analysis.name,
            "is_available": self.analysis.is_available(),
            "bindings": [
                {
                    "name": binding.name,
                    "container_type": binding.container_type.name,
                }
                for binding in self.bindings
            ],
            "node": self.node.id,
        }


class Pipeline:
    """
    A pipeline is a tree of tasks.

    Given a set of requests, a model and a configuration of the pipes, it can produce a schedule
    that fulfills the requests.
    """

    __slots__ = (
        "declarations",
        "root",
        "artifacts",
        "analyses",
        "analysis_lists",
        "savepoint_id_to_artifact",
        "savepoint_id_to_name",
    )

    def __init__(
        self,
        declarations: set[ContainerDeclaration],
        root: PipelineNode,
        artifacts: Optional[set[Artifact]] = None,
        analyses: Optional[set[AnalysisBinding]] = None,
        analysis_lists: Optional[Iterable[AnalysisList]] = None,
    ):
        self.root = root
        self.declarations = set(declarations)

        self.savepoint_id_to_artifact: dict[int, Artifact] = {}
        self.savepoint_id_to_name: dict[int, str] = {}

        self.artifacts: Mapping[str, Artifact] = {}
        """
        The artifacts, indexed by their name for easy access.
        """
        for artifact in artifacts or set():
            if artifact.name in self.artifacts:
                raise ValueError(
                    f"Artifact {artifact.name} is defined multiple times in the pipeline"
                )
            self.artifacts[artifact.name] = artifact

        self.analyses: Mapping[str, AnalysisBinding] = {}
        """
        The analyses, indexed by their name for easy access.
        """
        for analysis in analyses or set():
            if analysis.analysis.name in self.analyses:
                raise ValueError(
                    f"Analysis {analysis.analysis.name} is defined multiple times in the pipeline"
                )
            self.analyses[analysis.analysis.name] = analysis

        self.analysis_lists: dict[str, AnalysisList] = {}
        """
        Aliases for lists of analyses to execute sequentially.
        """
        for analysis_list in analysis_lists or set():
            # Check that the given analysis list is valid
            if analysis_list.name in self.analyses:
                raise ValueError(
                    f"Analyses list alias {analysis_list.name} conflicts with an existing analysis"
                    " name"
                )
            for analysis_name in analysis_list.analyses:
                if analysis_name not in self.analyses:
                    raise ValueError(
                        f"Analysis {analysis_name} in analyses list {analysis_list.name} is not "
                        "defined in the pipeline"
                    )
            if analysis_list.name in self.analysis_lists:
                raise ValueError(
                    f"Analyses list alias {analysis_list.name} is defined multiple times in the "
                    "pipeline"
                )
            # Add the analysis list
            self.analysis_lists[analysis_list.name] = analysis_list

        # Compute the hash for each `PipelineNode`, this needs to be done
        # before doing any other operation that relies on
        # `walk_pipeline(stable=True)` or `node.sorted_successors`
        self._compute_nodes_hash()
        # Assign the savepoint ranges, this works deterministically by leveraging
        # `node.sorted_successors`
        Pipeline.assign_savepoint_ranges(root)

        id_index = 0
        # Set the dependencies field
        for node in self.walk_pipeline(stable=True):
            node.id = id_index
            id_index += 1

            if node is self.root:
                node.pipe_dependencies = set()
                continue

            node.pipe_dependencies = set(
                chain.from_iterable(n.pipe_dependencies for n in node.predecessors)
            )
            if isinstance(node.task, Pipe):
                node.pipe_dependencies.add(node.task)
            if isinstance(node.task, SavePoint):
                assert node.savepoint_range is not None
                self.savepoint_id_to_name[node.savepoint_range.start] = node.task.name

        for name, artifact in self.artifacts.items():
            if name != artifact.name:
                raise ValueError(
                    f"Artifact name {artifact.name} does not match the key "
                    f"{name} in the artifacts map."
                )
            if isinstance(artifact.node.task, SavePoint):
                assert artifact.node.savepoint_range is not None
                self.savepoint_id_to_artifact[artifact.node.savepoint_range.start] = artifact

    def _compute_nodes_hash(self):
        # Call each node's `_compute_node_hash` in post-order
        sorter: TopologicalSorter[PipelineNode] = TopologicalSorter()
        for node in self.walk_pipeline():
            sorter.add(node, *node.successors)

        for node in sorter.static_order():
            node._compute_node_hash()

    @staticmethod
    def assign_savepoint_ranges(node: PipelineNode, current_id: int = 0) -> int:
        """
        Assigns savepoint ranges to the nodes in the pipeline tree.
        A savepoint range is a pair of integers that represent the start and end of
        the savepoint in a subtree. The idea is to deduplicate things common to a
        savepoint and all its children, so that we can efficiently represent a
        subtree of savepoints as a continuous range of integers.
        """
        # We ID only the savepoints, so we increment the ID only if the current node is a savepoint.
        # This implies that, while the ids start at 0, the first savepoint will have id 1.
        # This is needed in the case the first node is a pipe, which is not a savepoint
        if isinstance(node.task, SavePoint):
            current_id += 1
        # Save the id on preorder visit
        start = current_id
        # Recurse on the children, but do it in a deterministic way
        end = current_id
        for child in node.sorted_successors():
            current_id = Pipeline.assign_savepoint_ranges(child, current_id)
            assert child.savepoint_range is not None, (
                f"Child {child.task.name} does not have a savepoint range assigned:"
                f" {child.savepoint_range}"
            )
            end = max(end, child.savepoint_range.end)
        # The end of the savepoint range is assigned on postorder visit.
        # Therefore the end is inclusive
        assert node.savepoint_range is None, (
            f"SavePoint {node.task.name} already has a savepoint range assigned"
            ": {node.savepoint_range}"
        )
        node.savepoint_range = SavePointsRange(start, end)
        # Return the new id
        return current_id

    def walk_pipeline(
        self, start: Optional[PipelineNode] = None, forward: bool = True, stable: bool = False
    ) -> Generator[PipelineNode, None, None]:
        """BFS walk of pipeline nodes"""
        assert int(not forward) + int(stable) < 2, "forward=False,stable=True is unsupported"

        to_visit: List[PipelineNode] = [start or self.root]
        visited: Set[PipelineNode] = set()

        if forward:

            def successors(node):
                return node.sorted_successors() if stable else node.successors

        else:

            def successors(node):
                return node.predecessors

        while len(to_visit) > 0:
            node = to_visit.pop()
            yield node
            visited.add(node)
            for child_node in successors(node):
                if child_node not in visited:
                    to_visit.append(child_node)

    def graph(self, container_edges=False) -> Graph:
        """A graph for debugging purposes."""

        graph = Graph()

        nodes_map: Dict[PipelineNode, Graph.Node] = {}

        access_to_str = {
            TaskArgumentAccess.READ: "R",
            TaskArgumentAccess.WRITE: "W",
            TaskArgumentAccess.READ_WRITE: "RW",
        }

        for node in self.walk_pipeline():
            if isinstance(node.task, Pipe):
                new_node = Graph.Node(node.task.name)
                for argument in node.arguments:
                    new_node.entries.append(f"{argument.name} [{access_to_str[argument.access]}]")
            else:
                new_node = Graph.Node(node.task.name, color="#2A52BE", bgcolor="#6C83BE")
                for argument in node.arguments:
                    new_node.entries.append(argument.name)

            nodes_map[node] = new_node
            graph.nodes.add(new_node)

            for predecessor in node.predecessors:
                graph.edges.add(
                    Graph.Edge(
                        nodes_map[predecessor],
                        new_node,
                        source_port=-1,
                        destination_port=-1,
                        style="bold",
                    )
                )

            if not container_edges:
                continue

            node_inputs: list[ContainerDeclaration] = list(node.argument_declarations)
            taken_inputs: set[int] = set()

            for parent_node in self.walk_pipeline(node, forward=False):
                if parent_node is node:
                    continue

                for source_index, parent_argument in enumerate(parent_node.arguments):
                    parent_container_decl = parent_argument.declaration()
                    if (
                        parent_argument.access == TaskArgumentAccess.READ
                        or parent_container_decl not in node_inputs
                    ):
                        continue

                    destination_index = node_inputs.index(parent_container_decl)
                    if destination_index in taken_inputs:
                        continue

                    taken_inputs.add(destination_index)
                    if parent_node in node.predecessors:
                        continue

                    graph.edges.add(
                        Graph.Edge(
                            nodes_map[parent_node],
                            new_node,
                            source_port=source_index,
                            destination_port=destination_index,
                            style="dashed",
                            color="#FD6D53",
                        )
                    )

        return graph

    def schedule(
        self,
        model: ReadOnlyModel,
        target_node: PipelineNode,
        requests: Requests,
        configuration: PipelineConfiguration,
        storage_provider: StorageProvider,
    ) -> Schedule:
        def task_generator(pipeline_node: PipelineNode):
            if isinstance(pipeline_node.task, Pipe):
                return PipeScheduledTask(pipeline_node, model, storage_provider, configuration)
            else:
                return SavepointScheduledTask(pipeline_node, model, storage_provider, configuration)

        tasks: DefaultDictFromKey[PipelineNode, ScheduledTask] = DefaultDictFromKey(task_generator)
        # The pipeline is a tree, so we can just unroll the predecessors,
        # When we parallelize, we will make a subclass that overrides this method,
        # and probably it will first call it to produce the initial schedule and
        # then add the parallelization logic

        node: PipelineNode = target_node
        node_outgoing_requests = requests

        while not node_outgoing_requests.empty():
            pypeline_logger.debug_log(f"Scheduling node {node}")
            pypeline_logger.debug_log(f"Outgoing requests: {node_outgoing_requests}")
            orig = repr(node_outgoing_requests)
            # Each node should remove the requests it can handle
            # and add the requests it needs to satisfy the task
            node_ingoing_requests = node.prerequisites_for(
                model=model,
                requests=node_outgoing_requests,
                configuration=configuration,
                storage_provider=storage_provider,
            )
            assert orig == repr(node_outgoing_requests), (
                f"Node {node} modified the outgoing requests, which is not allowed. "
                f"Original: {orig}, modified: {node_outgoing_requests}"
            )
            pypeline_logger.debug_log(f"Computed Ingoing requests: {node_ingoing_requests}")
            # Store the computed requests so that we can use them in the run method
            tasks[node].add_requests(node_ingoing_requests, node_outgoing_requests)
            # If the node has no predecessors, we are done, but
            # we need to check that it has no requests left
            if not node.predecessors:
                assert node_ingoing_requests.empty(), (
                    f"Node {node} has no predecessors, but it still has "
                    f"requests: {node_ingoing_requests.minimize()}"
                )
                break

            assert len(node.predecessors) == 1, (
                f"Node {node} has multiple predecessors, but we assume a tree structure. "
                f"Predecessors: {node.predecessors}"
            )
            predecessor = node.predecessors[0]
            tasks[node].dependencies.append(tasks[predecessor])
            # Recurse on THE predecessor, its outgoing requests will be the
            # ingoing requests of the current node
            node = predecessor
            node_outgoing_requests = node_ingoing_requests

        # Here the task list has been decided, these are now iterated to apply
        # the following optimizations:
        # * Reduce the set of container declarations to those that are actually
        #   necessary to run the schedule
        # * Compact incoming and outgoing of each task to their reduced version
        # * Skip tasks where nothing in outgoing is actually written by the task
        #
        # These could be derived by using a dataflow-based approach to the
        # pipeline, where each pipe binds its inputs to the predecessor's pipe
        # output, but in the general case (with RW pipes) this turns into
        # having a mini-programming language and applying SSA analysis and the
        # like to deduplicate the containers (that, in the general case, need
        # to be duplicated at each node of the dataflow).
        scheduled_task: ScheduledTask | None = tasks[target_node]
        parent_scheduled_task: ScheduledTask | None = None
        used_declatations: set[str] = set()

        while scheduled_task is not None:
            # Assume that the schedule is a straight line, so a task has at
            # most one dependency
            assert len(scheduled_task.dependencies) in (0, 1)

            # Reduce incoming and outgoing
            scheduled_task.incoming = scheduled_task.incoming.minimize()
            scheduled_task.outgoing = scheduled_task.outgoing.minimize()

            # Figure out if a task will actually produce outputs
            written_out = Requests()
            for argument in scheduled_task.node.arguments:
                if argument.access != TaskArgumentAccess.READ:
                    container_decl = argument.declaration()
                    if container_decl in scheduled_task.outgoing:
                        written_out[container_decl] = scheduled_task.outgoing[container_decl]

            if written_out.empty() and parent_scheduled_task is not None:
                # If here the task does not actually need to produce anything
                # and can be skipped, by "glueing" its dependencies onto the parent
                assert [scheduled_task] == parent_scheduled_task.dependencies
                parent_scheduled_task.dependencies = scheduled_task.dependencies
            else:
                used_declatations.update(x.name for x in scheduled_task.node.argument_declarations)
                parent_scheduled_task = scheduled_task

            if len(scheduled_task.dependencies) == 1:
                scheduled_task = scheduled_task.dependencies[0]
            else:
                scheduled_task = None

        return Schedule(
            {v for v in self.declarations if v.name in used_declatations},
            tasks[target_node],
            model,
            storage_provider,
        )

    def get_artifact(
        self,
        model: ReadOnlyModel,
        artifact: Artifact,
        requests: ObjectSet,
        configuration: PipelineConfiguration,
        storage_provider: StorageProvider,
        runner_context: RunnerContext = RunnerContext(),
    ) -> Container:
        schedule = self.schedule(
            model=model,
            target_node=artifact.node,
            requests=Requests({artifact.container: requests}),
            configuration=configuration,
            storage_provider=storage_provider,
        )
        return schedule.run(runner_context)[artifact.container]

    def run_analysis_list(
        self,
        model: ReadOnlyModel,
        analysis_list: AnalysisList,
        configuration: PipelineConfiguration,
        storage_provider: StorageProvider,
        runner_context: RunnerContext = RunnerContext(),
    ) -> tuple[Model, InvalidatedObjects]:
        """
        Run a list of analyses on the pipeline, given a model and a set of requests.
        The analyses will return the new potentially modified model, and set it
        in the storage provider.
        """
        new_model: Model = model.clone()
        total_invalidated: InvalidatedObjects = InvalidatedObjects()

        pypeline_logger.debug_log(f"Running analysis list {analysis_list.name}")

        for analysis_name in analysis_list.analyses:
            if analysis_name not in self.analyses:
                raise PypelineException(f"Analysis {analysis_name} not found in the pipeline")

            if not self.analyses[analysis_name].analysis.is_available():
                raise PypelineException(
                    f"Analysis list {analysis_list.name} cannot be run "
                    f"because analysis {analysis_name} is not available"
                )

        for analysis_name in analysis_list.analyses:
            pypeline_logger.debug_log(f"Running analysis {analysis_name}")
            analysis = self.analyses[analysis_name]
            # Build the requests for the analysis
            requests = Requests()
            for container_decl in analysis.bindings:
                requests[container_decl] = model.all_objects(container_decl.container_type.kind)

            new_model, invalidated = self.run_analysis(
                model=model,
                analysis_name=analysis_name,
                requests=requests,
                configuration=configuration,
                storage_provider=storage_provider,
                runner_context=runner_context,
            )
            for location, objects in invalidated.items():
                if location in total_invalidated:
                    total_invalidated[location] = total_invalidated[location] | objects
                else:
                    total_invalidated[location] = objects
            model = ReadOnlyModel(new_model)

        return new_model, total_invalidated

    def run_analysis(
        self,
        model: ReadOnlyModel,
        analysis_name: str,
        requests: Requests,
        configuration: PipelineConfiguration,
        storage_provider: StorageProvider,
        runner_context: RunnerContext = RunnerContext(),
    ) -> tuple[Model, InvalidatedObjects]:
        """
        Run an analysis on the pipeline, given a model and a set of requests.
        The analysis will return the new potentially modified model, and set it
        in the storage provider.
        """
        if analysis_name not in self.analyses:
            raise PypelineException(f"Analysis {analysis_name} not found in the pipeline")
        analysis_info = self.analyses[analysis_name]

        if len(requests) != len(analysis_info.bindings):
            raise PypelineException(
                f"Expected {len(analysis_info.bindings)} requests for analysis "
                f"{analysis_name}, but got {len(requests)}: {requests}"
            )

        for req in requests:
            if req not in analysis_info.bindings:
                raise PypelineException(
                    f"Request {req} but it's not compatible with in the "
                    f"analysis bindings: {analysis_info.bindings}"
                )

        if not analysis_info.analysis.is_available():
            raise PypelineException(f"Analysis {analysis_name} is not available")

        schedule = self.schedule(
            model=model,
            target_node=analysis_info.node,
            requests=requests,
            configuration=configuration,
            storage_provider=storage_provider,
        )
        all_containers = schedule.run(runner_context)

        new_model = runner_context.run_analysis(
            analysis=analysis_info.analysis,
            model=model,
            containers=[all_containers[decl] for decl in analysis_info.bindings],
            incoming=[requests.get(decl) for decl in analysis_info.bindings],
            configuration=configuration.get(analysis_info.analysis, ""),
        )

        diff = model.diff(ReadOnlyModel(new_model))
        custom_invalidated_objects = self._compute_custom_invalidation(
            configuration, storage_provider, diff
        )
        set_model_result = storage_provider.set_model(
            new_model, diff.paths(), custom_invalidated_objects
        )
        return new_model, set_model_result.invalidated_objects

    def _compute_custom_invalidation(
        self,
        configuration: PipelineConfiguration,
        storage_provider: StorageProvider,
        diff: ModelDiff,
    ) -> list[ObjectsToInvalidate]:
        result: list[ObjectsToInvalidate] = []
        for node in self.walk_pipeline():
            if not isinstance(node.task, Pipe):
                continue

            assert node.savepoint_range is not None

            # Optimization: since few pipes actually implement advanced
            # invalidation and it requires querying storage, check that the
            # method has actually been overridden.
            if not node.task.has_custom_invalidation():
                continue

            # Run the prelimiary check for the pipe
            if not node.task.requires_custom_invalidation(diff):
                continue

            # Fetch the custom invalidation data from storage
            configuration_id = node.configuration_id(configuration)
            invalidation_data = storage_provider.get_custom_invalidation_data(
                node.id, configuration_id
            )
            # If the data returned is empty, it means that:
            # * The pipe did not return any invalidation data
            # * The pipe returned it, but it happened to be empty
            # In both cases it means that the pipe is opting-out of custom invalidation
            if all(len(x) == 0 for x in invalidation_data):
                continue

            # Run the pipe's invalidate
            objects_to_invalidate = node.task.process_custom_invalidation(invalidation_data, diff)
            # Convert the invalidated objects in a pipeline-friendly format
            for index, objects in enumerate(objects_to_invalidate):
                container_decl = node.argument_declarations[index]
                result.append(
                    ObjectsToInvalidate(
                        node.savepoint_range, container_decl.name, configuration_id, objects
                    )
                )

        return result

    def deserialize_schedule(
        self, schedule: str, model: ReadOnlyModel, storage_provider: StorageProvider
    ) -> Schedule:
        schedule_dict = yaml.safe_load(schedule)
        declarations = set()
        for container in schedule_dict["containers"]:
            for declaration in self.declarations:
                if container["name"] == declaration.name:
                    declarations.add(declaration)
                    break
            else:
                raise ValueError()

        pipeline_nodes = list(self.walk_pipeline(stable=True))
        container_map = {x.name: x for x in self.declarations}

        obj_id_type = get_singleton(ObjectID)  # type: ignore[type-abstract]
        configuration: dict[Pipe | Analysis, str] = {}
        scheduled_tasks: list[ScheduledTask] = []
        for task in schedule_dict["tasks"]:
            pipeline_node: PipelineNode = pipeline_nodes[task["node_id"]]
            dependencies = [scheduled_tasks[i] for i in task["dependencies"]]
            outgoing = Requests()
            incoming = Requests()

            if task["type"] == "Pipe":
                assert isinstance(pipeline_node.task, Pipe)
                assert task["name"] == pipeline_node.task.name

                configuration[pipeline_node.task] = task["dynamic_config"]
                for index, arg in enumerate(task["args"]):
                    container_declaration = pipeline_node.bindings[index]
                    assert container_declaration.name == arg["name"]
                    container_kind = container_declaration.container_type.kind

                    incoming[container_declaration] = ObjectSet(
                        container_kind, {obj_id_type.deserialize(x) for x in arg["incoming"]}
                    )
                    outgoing[container_declaration] = ObjectSet(
                        container_kind, {obj_id_type.deserialize(x) for x in arg["outgoing"]}
                    )

                scheduled_tasks.append(
                    PipeScheduledTask(
                        pipeline_node,
                        model,
                        storage_provider,
                        configuration,
                        (incoming, outgoing),
                        dependencies,
                    )
                )
            elif task["type"] == "SavePoint":
                assert isinstance(pipeline_node.task, SavePoint)
                assert task["name"] == pipeline_node.task.name

                for container in task["containers"]:
                    container_declaration = container_map[container["name"]]
                    container_kind = container_declaration.container_type.kind

                    incoming[container_declaration] = ObjectSet(
                        container_kind, {obj_id_type.deserialize(x) for x in container["incoming"]}
                    )
                    outgoing[container_declaration] = ObjectSet(
                        container_kind, {obj_id_type.deserialize(x) for x in container["outgoing"]}
                    )

                scheduled_tasks.append(
                    SavepointScheduledTask(
                        pipeline_node,
                        model,
                        storage_provider,
                        configuration,
                        (incoming, outgoing),
                        dependencies,
                    )
                )
            else:
                raise ValueError(f"Unknown task type: \"{task['type']}\"")

        return Schedule(declarations, scheduled_tasks[-1], model, storage_provider)