Keeps all nodes of one scope.
| 24 | |
| 25 | |
| 26 | class Scope: |
| 27 | """Keeps all nodes of one scope.""" |
| 28 | |
| 29 | _graph: ParseGraph |
| 30 | _nodes: StableSet[operator.Operator] |
| 31 | _normal_nodes: StableSet[operator.Operator] |
| 32 | |
| 33 | def __init__(self, graph: ParseGraph) -> None: |
| 34 | self._graph = graph |
| 35 | self._nodes = StableSet() |
| 36 | self._normal_nodes = StableSet() |
| 37 | |
| 38 | def is_empty(self) -> bool: |
| 39 | return not self._nodes |
| 40 | |
| 41 | def add_node(self, node: operator.Operator, *, special: bool = False) -> None: |
| 42 | self._nodes.add(node) |
| 43 | if not special: |
| 44 | self._normal_nodes.add(node) |
| 45 | |
| 46 | @property |
| 47 | def nodes(self) -> Iterator[operator.Operator]: |
| 48 | return iter(self._nodes) |
| 49 | |
| 50 | @property |
| 51 | def normal_nodes(self) -> Iterator[operator.Operator]: |
| 52 | return iter(self._normal_nodes) |
| 53 | |
| 54 | @property |
| 55 | def output_nodes(self) -> Iterator[operator.OutputOperator]: |
| 56 | return ( |
| 57 | node |
| 58 | for node in self.normal_nodes |
| 59 | if isinstance(node, operator.OutputOperator) |
| 60 | ) |
| 61 | |
| 62 | @property |
| 63 | def debug_nodes(self) -> Iterator[operator.DebugOperator]: |
| 64 | return ( |
| 65 | node |
| 66 | for node in self.normal_nodes |
| 67 | if isinstance(node, operator.DebugOperator) |
| 68 | ) |
| 69 | |
| 70 | def relevant_nodes( |
| 71 | self, operators: Iterable[operator.Operator] |
| 72 | ) -> list[operator.Operator]: |
| 73 | stack: list[operator.Operator] = list(operators) |
| 74 | visited: set[operator.Operator] = set(stack) |
| 75 | while stack: |
| 76 | node = stack.pop() |
| 77 | for dependency in node.input_operators(): |
| 78 | if dependency not in visited and dependency in self._nodes: |
| 79 | visited.add(dependency) |
| 80 | stack.append(dependency) |
| 81 | # Iterate over original list of nodes to preserve insertion order |
| 82 | return [node for node in self.nodes if node in visited] |
| 83 |