"""
    Support for traversing the IR, often to search for information.
"""
from dataclasses import dataclass, field
from typing import Callable, Optional, TypeVar, cast, Tuple, Union as PyUnion, Generic
from abc import abstractmethod
from .util import OrderedSet, ordered_set, FrozenOrderedSet, rewrite_list, rewrite_set
from . import ir

#--------------------------------------------------
# Visitor Abstraction
#--------------------------------------------------

Result = TypeVar('Result')
class GenericVisitor(Generic[Result]):
    """
    Abstract visitor with handlers for each IR node type.
    Each handler should return a value of type `Result`.

    Actual behavior (e.g., AST traversal) should be implemented in
    subclasses of `GenericVisitor`.
    """

    ##################################################
    # Model

    @abstractmethod
    def visit_model(self, node: ir.Model, parent: Optional[ir.Node]) -> Result:
        pass

    ##################################################
    # Engine

    @abstractmethod
    def visit_capability(self, node: ir.Capability, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_engine(self, node: ir.Engine, parent: Optional[ir.Node]) -> Result:
        pass

    ##################################################
    # Types

    @abstractmethod
    def visit_scalartype(self, node: ir.ScalarType, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_listtype(self, node: ir.ListType, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_settype(self, node: ir.SetType, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_uniontype(self, node: ir.UnionType, parent: Optional[ir.Node]) -> Result:
        pass

    ##################################################
    # Relations

    @abstractmethod
    def visit_field(self, node: ir.Field, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_relation(self, node: ir.Relation, parent: Optional[ir.Node]) -> Result:
        pass

    ##################################################
    # Variables and values

    @abstractmethod
    def visit_var(self, node: ir.Var, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_default(self, node: ir.Default, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_literal(self, node: ir.Literal, parent: Optional[ir.Node]) -> Result:
        pass

    ##################################################
    # Composite tasks

    @abstractmethod
    def visit_logical(self, node: ir.Logical, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_union(self, node: ir.Union, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_sequence(self, node: ir.Sequence, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_match(self, node: ir.Match, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_until(self, node: ir.Until, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_wait(self, node: ir.Wait, parent: Optional[ir.Node]) -> Result:
        pass

    ##################################################
    # Logical tasks

    @abstractmethod
    def visit_not(self, node: ir.Not, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_exists(self, node: ir.Exists, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_forall(self, node: ir.ForAll, parent: Optional[ir.Node]) -> Result:
        pass

    ##################################################
    # Iteration tasks

    @abstractmethod
    def visit_loop(self, node: ir.Loop, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_break(self, node: ir.Break, parent: Optional[ir.Node]) -> Result:
        pass

    ##################################################
    # Leaf tasks

    @abstractmethod
    def visit_update(self, node: ir.Update, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_annotation(self, node: ir.Annotation, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_lookup(self, node: ir.Lookup, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_output(self, node: ir.Output, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_construct(self, node: ir.Construct, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_aggregate(self, node: ir.Aggregate, parent: Optional[ir.Node]) -> Result:
        pass

    @abstractmethod
    def visit_task(self, node: ir.Task, parent: Optional[ir.Node]) -> Result:
        pass

@dataclass
class Visitor(GenericVisitor[None]):
    """
    A visitor that just walks the AST, performing no changes.
    """

    def enter(self, node: ir.Node, parent: Optional[ir.Node]=None) -> "Visitor":
        """ Visit the node, possibly returning a new visitor to be used to visit the node's children. """
        return self

    def prune(self):
        """ Return a new Visitor that just returns each node unchanged without traversing its children. """
        return PruneVisitor()

    def leave(self, node: ir.Node, parent: Optional[ir.Node]=None) -> ir.Node:
        """ Visit the node after visiting it and its children. """
        return node

    def _walk_node(self, node: ir.Node, parent: Optional[ir.Node]=None):
        # Enter the node, returning a new context.
        v = self.enter(node, parent)

        # Dispatch to the handler for the node.
        # By default the handler traverses the node's children and reconstructs the node.
        node.accept(v, parent)

        self.leave(node, parent)

    # Avoid dispatch overhead for some types.
    # These should never be called directly.
    def _walk_var(self, node: ir.Var, parent: Optional[ir.Node]):
        v = self.enter(node, parent)
        v.visit_var(node, parent)
        self.leave(node, parent)

    def _walk_var_or_default(self, node: PyUnion[ir.Var, ir.Default], parent: Optional[ir.Node]):
        v = self.enter(node, parent)
        if isinstance(node, ir.Var):
            v.visit_var(node, parent)
        else:
            v.visit_default(node, parent)
        self.leave(node, parent)

    def _walk_value(self, node: ir.Value, parent: Optional[ir.Node]):
        if node is None:
            return None
        elif isinstance(node, ir.Var):
            return self._walk_var(node, parent)
        elif isinstance(node, ir.Relation):
            return self._walk_relation(node, parent)
        elif isinstance(node, ir.Node):
            return self._walk_node(node, parent)
        elif isinstance(node, Tuple):
            for item in node:
                self._walk_value(item, parent)
        elif isinstance(node, FrozenOrderedSet):
            for item in node:
                self._walk_value(item, parent)

    def _walk_engine(self, node: ir.Engine, parent: Optional[ir.Node]):
        v = self.enter(node, parent)
        v.visit_engine(node, parent)
        self.leave(node, parent)

    def _walk_relation(self, node: ir.Relation, parent: Optional[ir.Node]):
        v = self.enter(node, parent)
        v.visit_relation(node, parent)
        self.leave(node, parent)

    def _walk_type(self, node: ir.Type, parent: Optional[ir.Node]):
        # The if-isinstance chain saves about 10% on a traversal vs. node.accept(v, parent)
        v = self.enter(node, parent)
        if isinstance(node, ir.ScalarType):
            v.visit_scalartype(node, parent)
        elif isinstance(node, ir.ListType):
            v.visit_listtype(node, parent)
        elif isinstance(node, ir.SetType):
            v.visit_settype(node, parent)
        elif isinstance(node, ir.UnionType):
            v.visit_uniontype(node, parent)
        else:
            raise NotImplementedError(f"visit_type not implemented for {type(node)}")
        self.leave(node, parent)

    ##################################################
    # Model

    def visit_model(self, node: ir.Model, parent: Optional[ir.Node]):
        for c in node.engines:
            self._walk_engine(c, node)
        for c in node.relations:
            self._walk_relation(c, node)
        for c in node.types:
            self._walk_type(c, node)
        self._walk_node(node.root, node)

    ##################################################
    # Engine

    def visit_capability(self, node: ir.Capability, parent: Optional[ir.Node]):
        pass

    def visit_engine(self, node: ir.Engine, parent: Optional[ir.Node]):
        for c in node.capabilities:
            self._walk_node(c, node)
        for c in node.relations:
            self._walk_relation(c, node)

    ##################################################
    # Types

    def visit_scalartype(self, node: ir.ScalarType, parent: Optional[ir.Node]):
        for t in node.super_types:
            self._walk_type(t, node)

    def visit_listtype(self, node: ir.ListType, parent: Optional[ir.Node]):
        self._walk_type(node.element_type, node)

    def visit_settype(self, node: ir.SetType, parent: Optional[ir.Node]):
        self._walk_type(node.element_type, node)

    def visit_uniontype(self, node: ir.UnionType, parent: Optional[ir.Node]):
        for t in node.types:
            self._walk_type(t, node)

    ##################################################
    # Relations

    def visit_field(self, node: ir.Field, parent: Optional[ir.Node]):
        self._walk_type(node.type, node)

    def visit_relation(self, node: ir.Relation, parent: Optional[ir.Node]):
        for f in node.fields:
            self._walk_node(f, node)
        for r in node.requires:
            self._walk_node(r, node)

    ##################################################
    # Variables and values

    def visit_var(self, node: ir.Var, parent: Optional[ir.Node]):
        self._walk_type(node.type, node)

    def visit_default(self, node: ir.Default, parent: Optional[ir.Node]):
        self._walk_var(node.var, node)
        self._walk_value(node.value, node)

    def visit_literal(self, node: ir.Literal, parent: Optional[ir.Node]):
        self._walk_type(node.type, node)

    ##################################################
    # Composite tasks

    def visit_logical(self, node: ir.Logical, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        for h in node.hoisted:
            self._walk_var_or_default(h, node)
        for b in node.body:
            self._walk_node(b, node)

    def visit_union(self, node: ir.Union, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        for h in node.hoisted:
            self._walk_var_or_default(h, node)
        for t in node.tasks:
            self._walk_node(t, node)

    def visit_sequence(self, node: ir.Sequence, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        for h in node.hoisted:
            self._walk_var_or_default(h, node)
        for t in node.tasks:
            self._walk_node(t, node)

    def visit_match(self, node: ir.Match, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        for h in node.hoisted:
            self._walk_var_or_default(h, node)
        for t in node.tasks:
            self._walk_node(t, node)

    def visit_until(self, node: ir.Until, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        for h in node.hoisted:
            self._walk_var_or_default(h, node)
        self._walk_node(node.check, node)
        self._walk_node(node.body, node)

    def visit_wait(self, node: ir.Wait, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        for h in node.hoisted:
            self._walk_var_or_default(h, node)
        self._walk_node(node.check, node)

    ##################################################
    # Logical tasks

    def visit_not(self, node: ir.Not, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        self._walk_node(node.task, node)

    def visit_exists(self, node: ir.Exists, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        for v in node.vars:
            self._walk_var(v, node)
        self._walk_node(node.task, node)

    def visit_forall(self, node: ir.ForAll, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        for v in node.vars:
            self._walk_var(v, node)
        self._walk_node(node.task, node)

    ##################################################
    # Iteration tasks

    def visit_loop(self, node: ir.Loop, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        for h in node.hoisted:
            self._walk_var_or_default(h, node)
        self._walk_var(node.iter, node)
        self._walk_node(node.body, node)

    def visit_break(self, node: ir.Break, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        self._walk_node(node.check, node)

    ##################################################
    # Leaf tasks

    def visit_update(self, node: ir.Update, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        self._walk_relation(node.relation, node)
        for a in node.args:
            self._walk_value(a, node)
        for a in node.annotations:
            self._walk_node(a, node)

    def visit_annotation(self, node: ir.Annotation, parent: Optional[ir.Node]):
        self._walk_relation(node.relation, node)
        for a in node.args:
            self._walk_value(a, node)

    def visit_lookup(self, node: ir.Lookup, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        self._walk_relation(node.relation, node)
        for a in node.args:
            self._walk_value(a, node)

    def visit_output(self, node: ir.Output, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        for _, v in node.aliases:
            self._walk_var(v, node)

    def visit_construct(self, node: ir.Construct, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        for v in node.values:
            self._walk_value(v, node)
        self._walk_var(node.id_var, node)

    def visit_aggregate(self, node: ir.Aggregate, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)
        self._walk_relation(node.aggregation, node)
        for p in node.projection:
            self._walk_var(p, node)
        for g in node.group:
            self._walk_var(g, node)
        for a in node.args:
            self._walk_value(a, node)

    def visit_task(self, node: ir.Task, parent: Optional[ir.Node]):
        if node.engine is not None:
            self._walk_engine(node.engine, node)

#--------------------------------------------------
# Some generally useful visitors
#--------------------------------------------------

@dataclass
class PruneVisitor(Visitor):
    """ A Visitor that does not walk children. """

    def _walk_node(self, node: ir.Node, parent: Optional[ir.Node]=None):
        pass
    def _walk_var(self, node: ir.Var, parent: Optional[ir.Node]):
        pass
    def _walk_var_or_default(self, node: ir.VarOrDefault, parent: Optional[ir.Node]):
        pass
    def _walk_engine(self, node: ir.Engine, parent: Optional[ir.Node]):
        pass
    def _walk_relation(self, node: ir.Relation, parent: Optional[ir.Node]):
        pass
    def _walk_type(self, node: ir.Type, parent: Optional[ir.Node]):
        pass

@dataclass
class DAGVisitor(Visitor):
    """ A visitor that just walks the AST, performing no changes. Nodes are visited at most once. """
    seen: set[ir.Node] = field(default_factory=set)

    def _walk_node(self, node: ir.Node, parent: Optional[ir.Node]=None):
        if node in self.seen:
            return
        self.seen.add(node)
        return super()._walk_node(node, parent)
    def _walk_var(self, node: ir.Var, parent: Optional[ir.Node]):
        if node in self.seen:
            return
        self.seen.add(node)
        return super()._walk_var(node, parent)
    def _walk_var_or_default(self, node: ir.VarOrDefault, parent: Optional[ir.Node]):
        if node in self.seen:
            return
        self.seen.add(node)
        return super()._walk_var_or_default(node, parent)
    def _walk_engine(self, node: ir.Engine, parent: Optional[ir.Node]):
        if node in self.seen:
            return
        self.seen.add(node)
        return super()._walk_engine(node, parent)
    def _walk_relation(self, node: ir.Relation, parent: Optional[ir.Node]):
        if node in self.seen:
            return
        self.seen.add(node)
        return super()._walk_relation(node, parent)
    def _walk_type(self, node: ir.Type, parent: Optional[ir.Node]):
        if node in self.seen:
            return
        self.seen.add(node)
        return super()._walk_type(node, parent)

class Collector(Visitor):
    """ A visitor that collects instances that match a predicate. """
    def __init__(self, predicate: Callable[[ir.Node, Optional[ir.Node]], bool]):
        super().__init__()
        self.elements: OrderedSet = ordered_set()
        self.predicate = predicate

    def enter(self, node: ir.Node, parent: Optional[ir.Node]=None):
        if self.predicate(node, parent):
            self.elements.add(node)
        return self

def collect(predicate: Callable[[ir.Node, Optional[ir.Node]], bool], *nodes: ir.Node) -> OrderedSet[ir.Node]:
    """ Collect children of node that match the predicate. """
    c = Collector(predicate)
    for n in nodes:
        n.accept(c)
    return c.elements

T = TypeVar('T')
def collect_by_type(t: PyUnion[type[T], tuple[type[T], ...]], *nodes: ir.Node) -> OrderedSet[T]:
    """ Collect instances of the type t by traversing this node and its children. """
    return cast(OrderedSet[T],
        collect(lambda n, parent: isinstance(n, t), *nodes)
    )

def collect_vars(*nodes: ir.Node) -> OrderedSet[ir.Var]:
    """ Collect all Vars starting at this node. """
    return cast(OrderedSet[ir.Var],
        collect_by_type(ir.Var, *nodes)
    )

def collect_quantified_vars(*nodes: ir.Node) -> OrderedSet[ir.Var]:
    """ Collect all Vars that are children of Exists and ForAll. """
    return cast(OrderedSet[ir.Var],
        collect(
            lambda n, parent: isinstance(n, ir.Var) and isinstance(parent, (ir.Exists, ir.ForAll)),
            *nodes)
    )

def collect_aggregate_vars(*nodes: ir.Node) -> OrderedSet[ir.Var]:
    """ Collect vars that are declared by aggregates in Rel (projection + over). """
    return cast(OrderedSet[ir.Var],
        # TODO - when dealing with multiple aggregations we will need to consider groupbys
        collect(
            lambda n, parent:
                # parent of var is an aggregate and either
                isinstance(parent, ir.Aggregate) and isinstance(n, ir.Var) and (
                # var is in the projection
                n in parent.projection or
                # var is an input to the aggregation's relation
                is_aggregate_input(n, parent)
                ),
            *nodes)
    )

def is_aggregate_input(var: ir.Var, agg: ir.Aggregate):
    # field at the var index is input (the +2 is for projection + group)
    return (var in agg.args and agg.aggregation.fields[agg.args.index(var) + 2].input)

def collect_implicit_vars(*nodes: ir.Node) -> OrderedSet[ir.Var]:
    """ Collect vars except the quantified and aggregate vars. """
    if not nodes:
        return ordered_set()
    return collect_vars(*nodes) - collect_quantified_vars(*nodes) - collect_aggregate_vars(*nodes)



@dataclass
class ReadWriteVisitor(Visitor):
    """
    Compute the set of reads and writes for Logical nodes.

    Note that reads are Lookups and writes are Updates. We don't consider Output a write
    because it is not targeting a relation.
    """
    # TODO: we currently only compute for Logical nodes, but it may be useful for other nodes
    _reads: dict[ir.Logical, OrderedSet[ir.Relation]] = field(default_factory=dict)
    _writes: dict[ir.Logical, OrderedSet[ir.Relation]] = field(default_factory=dict)

    def reads(self, key: ir.Logical):
        # TODO - use a singleton empty set
        return self._reads[key] if key in self._reads else ordered_set()

    def writes(self, key: ir.Logical):
        return self._writes[key] if key in self._writes else ordered_set()

    _stack: list[ir.Logical] = field(default_factory=list)

    def visit_logical(self, node: ir.Logical, parent: Optional[ir.Node]):
        self._stack.append(node)
        super().visit_logical(node, parent)
        self._stack.pop()

    def visit_lookup(self, node: ir.Lookup, parent: Optional[ir.Node]):
        for lu in self._stack:
            if lu not in self._reads:
                self._reads[lu] = ordered_set()
            self._reads[lu].add(node.relation)
        return super().visit_lookup(node, parent)

    def visit_aggregate(self, node: ir.Aggregate, parent: Optional[ir.Node]):
        for lu in self._stack:
            if lu not in self._reads:
                self._reads[lu] = ordered_set()
            self._reads[lu].add(node.aggregation)
        return super().visit_aggregate(node, parent)

    def visit_update(self, node: ir.Update, parent: Optional[ir.Node]):
        for lu in self._stack:
            if lu not in self._writes:
                self._writes[lu] = ordered_set()
            self._writes[lu].add(node.relation)
        return super().visit_update(node, parent)

@dataclass
class DeprecatedPass():
    """
    This is the old Pass class. To be replaced by a Rewriter.

    """

    rewritten: dict[ir.Node, ir.Node] = field(default_factory=dict, init=False, compare=False, hash=False)
    handler_cache: dict[str, Callable[[ir.Node, Optional[ir.Node]], ir.Node]] = field(default_factory=dict, init=False, compare=False, hash=False)

    T = TypeVar('T', bound=PyUnion[ir.Node, ir.Value])
    def walk(self, node: T, parent=None) -> T:
        # if node is a value, just return it
        if not isinstance(node, ir.Node):
            return node

        result = self.rewritten.get(node, None)
        if result:
            return cast(T, result) # type: ignore[reportReturnType]

        # node is actually some Node type, handle with the appropriate handler
        handler = self.handler_cache.get(node.kind, None)
        if not handler:
            handler = getattr(self, f"handle_{node.kind}", None)
            if handler:
                self.handler_cache[node.kind] = handler
        if handler:
            result = cast(ir.Node, handler(node, parent))
            self.rewritten[node] = result
            return cast(T, result) # type: ignore[reportReturnType]
        else:
            raise NotImplementedError(f"walk: {node.kind}")

    def walk_set(self, items: FrozenOrderedSet[T], parent=None) -> FrozenOrderedSet[T]:
        return ordered_set(*[self.walk(n, parent) for n in items]).frozen()

    def walk_list(self, items: Tuple[T, ...], parent=None) -> Tuple[T, ...]:
        return tuple([self.walk(n, parent) for n in items])

    U = TypeVar('U', bound=PyUnion[ir.Node, ir.Value, FrozenOrderedSet[ir.Relation]])

    def walk_dict(self, items: dict[str, U], parent=None) -> dict[str, U]:
        d = dict()
        for k, v in items.items():
            if isinstance(v, FrozenOrderedSet):
                d[k] = self.walk_set(v, parent) # type: ignore
            else:
                d[k] = self.walk(v, parent)
        return d



    #-------------------------------------------------
    # Public Types - Model
    #-------------------------------------------------

    def handle_model(self, model: ir.Model, parent: None):
        engines = rewrite_set(ir.Engine, lambda n: self.walk(n, model), model.engines)
        relations = rewrite_set(ir.Relation, lambda n: self.walk(n, model), model.relations)
        types = rewrite_set(ir.Type, lambda n: self.walk(n, model), model.types)
        root = self.walk(model.root, model)

        return model.reconstruct(engines, relations, types, root)

    #-------------------------------------------------
    # Public Types - Engine
    #-------------------------------------------------

    def handle_capability(self, node: ir.Capability, parent: ir.Node):
        return node

    def handle_engine(self, node: ir.Engine, parent: ir.Node):
        return node

    #-------------------------------------------------
    # Public Types - Data Model
    #-------------------------------------------------

    def handle_scalartype(self, node: ir.ScalarType, parent: ir.Node):
        return node

    def handle_listtype(self, node: ir.ListType, parent: ir.Node):
        return node

    def handle_settype(self, node: ir.SetType, parent: ir.Node):
        return node

    def handle_uniontype(self, node: ir.UnionType, parent: ir.Node):
        # TODO - we could traverse the children
        return node

    def handle_field(self, node: ir.Field, parent: ir.Node):
        return node

    def handle_relation(self, node: ir.Relation, parent: ir.Node):
        return node

    #-------------------------------------------------
    # Public Types - Tasks
    #-------------------------------------------------

    def handle_task(self, node: ir.Task, parent: ir.Node):
        return node

    #
    # Task composition
    #

    def handle_logical(self, node: ir.Logical, parent: ir.Node):
        hoisted = rewrite_list(ir.VarOrDefault, lambda n: self.walk(n, node), node.hoisted)
        body = rewrite_list(ir.Task, lambda n: self.walk(n, node), node.body)

        return node.reconstruct(node.engine, hoisted, body)

    def handle_union(self, node: ir.Union, parent: ir.Node):
        hoisted = rewrite_list(ir.VarOrDefault, lambda n: self.walk(n, node), node.hoisted)
        tasks = rewrite_list(ir.Task, lambda n: self.walk(n, node), node.tasks)

        return node.reconstruct(node.engine, hoisted, tasks)

    def handle_sequence(self, node: ir.Sequence, parent: ir.Node):
        hoisted = rewrite_list(ir.VarOrDefault, lambda n: self.walk(n, node), node.hoisted)
        tasks = rewrite_list(ir.Task, lambda n: self.walk(n, node), node.tasks)

        return node.reconstruct(node.engine, hoisted, tasks)

    def handle_match(self, node: ir.Match, parent: ir.Node):
        hoisted = rewrite_list(ir.VarOrDefault, lambda n: self.walk(n, node), node.hoisted)
        tasks = rewrite_list(ir.Task, lambda n: self.walk(n, node), node.tasks)

        return node.reconstruct(node.engine, hoisted, tasks)

    def handle_until(self, node: ir.Until, parent: ir.Node):
        hoisted = rewrite_list(ir.VarOrDefault, lambda n: self.walk(n, node), node.hoisted)
        check = self.walk(node.check, node)
        body = self.walk(node.body, node)

        return node.reconstruct(node.engine, hoisted, check, body)

    def handle_wait(self, node: ir.Wait, parent: ir.Node):
        hoisted = rewrite_list(ir.VarOrDefault, lambda n: self.walk(n, node), node.hoisted)
        check = self.walk(node.check, node)

        return node.reconstruct(node.engine, hoisted, check)

    #
    # Relational Operations
    #

    def handle_var(self, node: ir.Var, parent: ir.Node):
        type_val = self.walk(node.type, node)
        name_val = node.name or f"v{node.id}"
        return node.reconstruct(type_val, name_val)

    def handle_default(self, node: ir.Default, parent: ir.Node):
        var = self.walk(node.var, node)
        return node.reconstruct(var, node.value)

    def handle_literal(self, node: ir.Literal, parent: ir.Node):
        type_val = self.walk(node.type, node)
        return node.reconstruct(type_val, node.value)

    def handle_annotation(self, node: ir.Annotation, parent: ir.Node):
        return node

    def handle_update(self, node: ir.Update, parent: ir.Node):
        relation = self.walk(node.relation, node)
        args = rewrite_list(ir.Value, lambda n: self.walk(n, node), node.args)
        return node.reconstruct(node.engine, relation, args, node.effect, node.annotations)

    def handle_lookup(self, node: ir.Lookup, parent: ir.Node):
        relation = self.walk(node.relation, node)
        args = rewrite_list(ir.Value, lambda n: self.walk(n, node), node.args)
        return node.reconstruct(node.engine, relation, args)

    def handle_output(self, node: ir.Output, parent: ir.Node):
        def rewrite_alias(pair: Tuple[str, ir.Var]):
            name, x = pair
            new_x = self.walk(x, node)
            if new_x is not x:
                return (name, new_x)
            else:
                return pair
        # rewrite_set can't be passed a generic type, so just use Tuple and cast the result.
        aliases = cast(
            FrozenOrderedSet[Tuple[str, ir.Var]],
            rewrite_set(
                Tuple,
                rewrite_alias,
                node.aliases))
        return node.reconstruct(node.engine, aliases)

    def handle_construct(self, node: ir.Construct, parent: ir.Node):
        values = rewrite_list(ir.Value, lambda n: self.walk(n, node), node.values)
        id_var = self.walk(node.id_var, node)
        return node.reconstruct(node.engine, values, id_var)

    def handle_aggregate(self, node: ir.Aggregate, parent: ir.Node):
        aggregation = self.walk(node.aggregation, node)
        projection = rewrite_list(ir.Var, lambda n: self.walk(n, node), node.projection)
        group = rewrite_list(ir.Var, lambda n: self.walk(n, node), node.group)
        args = rewrite_list(ir.Value, lambda n: self.walk(n, node), node.args)
        return node.reconstruct(node.engine, aggregation, projection, group, args)

    #
    # Logical Quantifiers
    #

    def handle_not(self, node: ir.Not, parent: ir.Node):
        task = self.walk(node.task, node)
        return node.reconstruct(node.engine, task)

    def handle_exists(self, node: ir.Exists, parent: ir.Node):
        vars = rewrite_list(ir.Var, lambda n: self.walk(n, node), node.vars)
        task = self.walk(node.task, node)
        return node.reconstruct(node.engine, vars, task)

    def handle_forall(self, node: ir.ForAll, parent: ir.Node):
        vars = rewrite_list(ir.Var, lambda n: self.walk(n, node), node.vars)
        task = self.walk(node.task, node)
        return node.reconstruct(node.engine, vars, task)

    #
    # Iteration (Loops)
    #
    def handle_loop(self, node: ir.Loop, parent: ir.Node):
        hoisted = rewrite_list(ir.VarOrDefault, lambda n: self.walk(n, node), node.hoisted)
        iter_val = self.walk(node.iter, node)
        body = self.walk(node.body, node)
        return node.reconstruct(node.engine, hoisted, iter_val, body)

    def handle_break(self, node: ir.Break, parent: ir.Node):
        check = self.walk(node.check, node)
        return node.reconstruct(node.engine, check)
