from ....metamodel.rewriter import Rewriter
from ....metamodel.metamodel import Task, Update
from .....config import Config
from ..annotations import require_seminaive
from .lqp_pass import Pass

class RecursionConfig(Pass):
    """
    A rewrite pass that adds the `@recursion_config(:require_seminaive)` annotation to all
    `Update` nodes.
    """
    def rewrite_task(self, node: Task, config:Config) -> Task:
        if not config.compiler.safe_recursion:
            return node
        rewriter = RecursionConfigRewriter()
        return rewriter.add_recursion_config(node)

class RecursionConfigRewriter(Rewriter):
    """
    Rewriter that adds the `@recursion_config(:require_seminaive)` annotation to all
    `Update` nodes.
    """
    def add_recursion_config(self, node: Task):
        return self(node)

    def update(self, node: Update):
        if require_seminaive in node.annotations:
            return node

        new_annotations = list(node.annotations) + [require_seminaive]
        return node.mut(annotations=tuple(new_annotations))
