revng-revng/python/revng/pypeline/pipeline_parser.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 pathlib import Path
from typing import Any, Optional

import jsonschema
import yaml

from .analysis import Analysis, AnalysisList
from .container import Container, ContainerDeclaration
from .pipeline import AnalysisBinding, Artifact, ArtifactCategory, Pipeline
from .pipeline_node import DummyPipelineNode, PipelineNode
from .task.pipe import Pipe
from .task.savepoint import SavePoint
from .utils.registry import get_registry


@dataclass(slots=True)
class Node:
    """
    We use this to represent the graph of branches in the pipeline,
    so we can compute the loading order of the branches.
    """

    content: Any
    is_root: bool = True
    successors: set[str] = field(default_factory=set)


def schema() -> dict[str, Any]:
    """
    Return the jsonschema for the pipeline.
    """
    root = Path(__file__).resolve().parent
    with open(root / "schema.yml", "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


class PipelineParser:
    def __init__(self, values: Any):
        # This setup is needed so we can have multiple schemas in the same file
        # so we can reuse the definitions
        full_schema = schema()
        resolver = jsonschema.RefResolver.from_schema(full_schema)
        pipeline_schema = full_schema["$defs"]["pipeline"]
        validator = jsonschema.Draft7Validator(pipeline_schema, resolver=resolver)
        validator.validate(values)

        self.values = values
        self.container_decls: dict[str, ContainerDeclaration] = {}
        self.artifacts: set[Artifact] = set()
        self.analyses: set[AnalysisBinding] = set()
        self.artifact_categories: dict[str, ArtifactCategory] = {}

    def parse(self) -> Pipeline:
        # Parse all the container declarations
        self.container_decls = self._parse_container_decls(self.values["containers"])

        # Parse all the artifact categories
        self.artifact_categories = self._parse_artifact_categories(
            self.values["artifact-categories"]
        )
        # Check that there is at least one artifact category
        assert len(self.artifact_categories) > 0

        # Then parse analyses lists, we will just do structural parsing, the actual validation
        # will be done in the Pipeline's __init__
        analysis_lists: list[AnalysisList] = []
        for analysis_list in self.values.get("analysis-lists", []):
            analysis_lists.append(
                AnalysisList(
                    name=analysis_list["name"],
                    analyses=analysis_list["analyses"],
                    description=analysis_list.get("description"),
                )
            )

        roots: list[PipelineNode] = self._parse_branches(branches=self.values["branches"])

        root_analyses_raw = self.values.get("analyses", [])
        if len(roots) == 1 and len(root_analyses_raw) == 0:
            root = roots[0]
        else:
            root = DummyPipelineNode()
            for node in roots:
                root.add_successor(node)
            root_analyses_dict = [{"analysis": x, "containers": []} for x in root_analyses_raw]
            self._parse_analyses(root_analyses_dict, root)

        # Check artifact consistency
        artifact_names = {a.name for a in self.artifacts}
        for artifact in self.artifacts:
            # TODO: also check for locations once those are statically defined
            for preferred_artifact in artifact.preferred_artifacts:
                if preferred_artifact not in artifact_names:
                    raise ValueError(
                        f"Artifact {artifact.name} defines {preferred_artifact} "
                        "as a preferred artifact but it does not exist"
                    )

        return Pipeline(
            declarations=set(self.container_decls.values()),
            root=root,
            artifacts=self.artifacts,
            analyses=self.analyses,
            analysis_lists=analysis_lists,
        )

    def _parse_pipe(self, task: dict[str, Any]):
        """Parse a pipe task from the JSON value."""
        pipes = get_registry(Pipe)  # type: ignore[type-abstract]
        pipe_name = task["pipe"]
        if pipe_name not in pipes:
            raise ValueError(
                f"Pipe {pipe_name} is not registered, available pipes: " f"{sorted(pipes.keys())}"
            )
        pipe_type = pipes[pipe_name]

        pipe_args = task.get("arguments", [])
        bindings = []
        for arg in pipe_args:
            if arg not in self.container_decls:
                raise ValueError(
                    f'While parsing {pipe_name}\'s arguments found container "{arg}" '
                    "that is not declared in the pipeline"
                )
            bindings.append(self.container_decls[arg])

        configuration = task.get("configuration")
        if configuration is None:
            configuration_string = ""
        elif isinstance(configuration, str):
            configuration_string = configuration
        else:
            configuration_string = yaml.safe_dump(configuration)

        return PipelineNode(
            task=pipe_type(configuration_string),
            bindings=bindings,
        )

    def _parse_savepoint(self, task: dict[str, Any]):
        """Parse a savepoint task from the JSON value."""
        name = task["savepoint"]
        containers = task["containers"]
        args = []
        for container_name in containers:
            if container_name not in self.container_decls:
                raise ValueError(f"Container {container_name} is not declared in the pipeline")
            args.append(self.container_decls[container_name])
        return PipelineNode(SavePoint(name=name, to_save=args))

    def _parse_artifacts(self, node_artifacts: list[Any], target_node: PipelineNode):
        """
        Parse artifacts from the node artifacts list and populate the artifacts dictionary.
        """
        for artifact in node_artifacts:
            name = artifact["name"]
            container = artifact["container"]
            if container not in self.container_decls:
                raise ValueError(
                    f"Artifact {name} references container {container} that is not "
                    "declared in the pipeline"
                )

            # Artifacts need to have a category associated to them
            if artifact["category"] not in self.artifact_categories:
                raise ValueError(f"Artifact category {artifact["category"]} not found")

            self.artifacts.add(
                Artifact(
                    name=name,
                    node=target_node,
                    container=self.container_decls[container],
                    description=artifact.get("description"),
                    category=self.artifact_categories[artifact["category"]],
                    filename=artifact.get("filename"),
                    defined_locations=artifact.get("defined_locations", []),
                    preferred_artifacts=artifact.get("preferred_artifacts", []),
                )
            )

    def _parse_analyses(self, node_analyses: list[Any], target_node: PipelineNode):
        """
        Parse analyses from the node analyses list and populate the analyses dictionary.
        """
        for analysis in node_analyses:
            analysis_name = analysis["analysis"]
            if analysis_name in self.analyses:
                raise ValueError(
                    f"Analysis {analysis_name} is defined multiple times in the pipeline"
                )
            containers = analysis["containers"]
            bindings = []
            for container in containers:
                if container not in self.container_decls:
                    raise ValueError(
                        f"Analysis {analysis_name} references container {container} that is "
                        "not declared in the pipeline"
                    )
                bindings.append(self.container_decls[container])

            analyses_registry = get_registry(Analysis)  # type: ignore[type-abstract]
            if analysis_name not in analyses_registry:
                raise ValueError(
                    f"Analysis {analysis_name} is not registered, available analyses: "
                    f"{sorted(analyses_registry.keys())}"
                )
            analysis_type: type[Analysis] = analyses_registry[analysis_name]
            self.analyses.add(
                AnalysisBinding(
                    analysis=analysis_type(),
                    bindings=tuple(bindings),
                    node=target_node,
                )
            )

    def _parse_task(self, task: Any, parent: Optional[PipelineNode] = None) -> PipelineNode:
        """
        Parse a single task from the JSON value.
        The JSON value should contain a dictionary with the branch content.
        """

        # Create the PipelineNode
        res: PipelineNode
        if "savepoint" in task:
            res = self._parse_savepoint(task)
        else:
            res = self._parse_pipe(task)

        # Parse artifacts
        node_artifacts = task.get("artifacts", [])
        self._parse_artifacts(node_artifacts=node_artifacts, target_node=res)

        # Parse analyses
        node_analyses = task.get("analyses", [])
        self._parse_analyses(node_analyses=node_analyses, target_node=res)

        # Connect
        if parent is not None:
            parent.add_successor(res)
        return res

    def _parse_branch(
        self, node: Node, graph: dict[str, PipelineNode]
    ) -> tuple[PipelineNode, PipelineNode]:
        """
        Parse the branch and return the last node in the branch.
        """
        parent: Optional[PipelineNode] = None
        if "from" in node.content:
            parent_name = node.content["from"]
            if parent_name not in graph:
                raise ValueError(f"Branch {parent_name} is not defined in the pipeline")
            parent = graph[parent_name]

        # Parse the nodes
        tasks = node.content.get("tasks", [])
        head = None
        for task in tasks:
            if "pipe" not in task and "savepoint" not in task:
                raise ValueError("Task must have either a pipe or a savepoint")
            if "pipe" in task and "savepoint" in task:
                raise ValueError("Task cannot have both a pipe and a savepoint")

            # Parse the node
            res = self._parse_task(task=task, parent=parent)
            if head is None:
                head = res
            parent = res

        assert parent is not None, "A branch cannot be empty and not have a parent"
        assert head is not None, "A branch needs to have at least one task"
        return (head, parent)

    def _parse_branches(self, branches: dict[str, Any]) -> list[PipelineNode]:
        """
        Parse the branches from the JSON value.
        The JSON value should contain a dictionary of branches with their names and tasks.
        """

        # Find the root branch
        graph: dict[str, Node] = {}
        roots: set[str] = set()
        for name, branch in branches.items():
            node = Node(content=branch)
            graph[name] = node
            if "from" in branch:
                node.is_root = False
                from_node = graph[branch["from"]]
                from_node.successors.add(name)
            else:
                roots.add(name)

        sorter: TopologicalSorter[str] = TopologicalSorter()
        for name, node in graph.items():
            sorter.add(name)
            for successor in node.successors:
                sorter.add(successor, name)

        # Parse the branches in the correct order (DFS)
        node_tips: dict[str, PipelineNode] = {}
        result: list[PipelineNode] = []
        for name in sorter.static_order():
            first_node, last_node = self._parse_branch(node=graph[name], graph=node_tips)
            node_tips[name] = last_node
            if name in roots:
                result.append(first_node)

        return result

    def _parse_container_decls(self, containers: list[Any]) -> dict[str, ContainerDeclaration]:
        """
        Parse the container declarations from the JSON value.
        The JSON value should contain a list of containers with their names and types.
        """
        containers_registry = get_registry(Container)  # type: ignore[type-abstract]
        container_decls: dict[str, ContainerDeclaration] = {}
        for container in containers:
            name = container["name"]
            ty = container["type"]
            if ty not in containers_registry:
                raise ValueError(
                    f"Container type {ty} is not registered, the available types "
                    f"are: {sorted(containers_registry.keys())}"
                )

            container_decls[name] = ContainerDeclaration(
                name=name,
                container_type=containers_registry[ty],
            )
        return container_decls

    def _parse_artifact_categories(
        self, artifact_categories: list[Any]
    ) -> dict[str, ArtifactCategory]:
        result: dict[str, ArtifactCategory] = {}
        for element in artifact_categories:
            if element["name"] in result:
                raise ValueError(f"Artifact category {element["name"]} is duplicated")
            result[element["name"]] = ArtifactCategory(element["name"], element["show-by-default"])
        return result


def load_pipeline_yaml(yaml_data: str) -> Pipeline:
    """
    Load a pipeline from a YAML string.
    """
    values = yaml.safe_load(yaml_data)
    return PipelineParser(values).parse()


def load_pipeline_yaml_file(file: str) -> Pipeline:
    """
    Load a pipeline from a YAML file."""
    with open(file, "r", encoding="utf-8") as f:
        values = yaml.safe_load(f)
    return PipelineParser(values).parse()