diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index 6a414bd67f8..2be4ec6ae64 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -665,12 +665,15 @@ materialization. iterator ancestors because the source path crosses a collector, prepared parent iteration paths are still used to materialize one downstream exec node for each preserved collector path. -- For an IfInvocation with a wired condition, prepare condition inputs first. Readiness is constrained to condition - source nodes and their ancestors; iterator ancestry on a branch cannot hold the gate open. After the condition resolves, - activation dependencies admit selected branch contexts and attach only the selected branch input. Literal conditions - use their known value and selected-branch ancestry. Unsupported fresh shapes and legacy snapshots use the compatibility - controller. Rejected branch sources are retired empty without invocation. Pending If state is reconstructed from durable - execution edges, prepared mappings, and the activation token after rehydration. +- For an IfInvocation with a wired condition, prepare condition inputs first. Readiness waits for condition source nodes + and their ancestors, and for the iterators in the `If`'s own iterator ancestry, which give it its frames and are never + gated by it; other iterator ancestry on a branch cannot hold the gate open. After the condition resolves, activation + dependencies admit selected branch contexts and attach only the selected branch input. A literal condition is gated by + the ancestry of its known selected branch; an iterator that only feeds the unselected branch, such as one closed by a + branch `Collect`, cannot hold it back. Its frames still wait for the iterators open at it. Unsupported fresh shapes + and legacy snapshots use the compatibility controller. Rejected branch sources are retired empty without invocation. + Pending If state is reconstructed from durable execution edges, prepared mappings, and the activation token after + rehydration. - A branch `Collect` whose `item` source combines exactly two sibling `Iterate` axes can close the axis beside the `If` condition axis. The compatibility materializer expands that sibling `Iterate` beneath each resolved `If` frame @@ -714,7 +717,9 @@ and must be considered separately when reasoning about rollback behavior. The execution-node builder retains a filtered iterator graph for the current source-graph identity. It builds this graph when compatibility preparation needs sibling-Collect relationships or iterator ancestry. The graph replaces Collect item inputs with its collection inputs for traversal; edge-scoped ancestry retains item axes needed for grouping. -Final For outputs do not carry an open iteration axis. +Final For outputs do not carry an open iteration axis. `If` branch gating and axis validation instead use the ancestry's +frame iterators, which follow a final For output into the For's parent frame: the For's own axis is closed there, but +the axes it runs under stay open. Iterator ancestry owns ordered ancestor and edge-axis analysis plus effective-path memoization for one graph traversal. The builder shares that object across candidate admission and materialization. Its separate per-source axis cache adds @@ -803,10 +808,51 @@ Run `C` -> `D:0` -> enqueue `D`. Run `D` -> done. performs pruning or skip propagation. - Legacy snapshots that already contain skipped prepared nodes retain their compatibility metadata and scheduler projection. Fresh execution does not create that projection. -- If an `Iterate.item` stream feeding `If.condition` closes empty, no `If` execution is created. The branch sources are - retired empty, the downstream `Collect` produces an empty collection, and the state completes without an error. - Retirement may add the branch source IDs to `executed_history`; neither branch invocation runs and no branch result - is created. +- An `If` runs once per frame of the innermost iterator in its iterator ancestry, whichever branch is selected. When a + condition, literal or computed in the root frame, selects a root value while the other branch carries an open + `Iterate`, the `If` keeps that per-item cardinality and returns the root value in each item frame; the unselected + per-item computation stays inactive. A branch `For`'s final output carries no axis into that ancestry, so selecting + the root value over such a branch runs the `If` once in the root frame, while selecting it runs the `If` per item of + the iterator the `For` runs under. +- The iterators whose frames are open at an `If`, including one reached through a `For` final output (which closes only + the `For`'s own axis), run regardless of the condition and are never gated by it. +- Before invocation, graph validation raises `UnsupportedIfIterationAxesError` for these frame sets: + - Two unrelated iterators open at the `If` (neither stays open inside the other, judged by iterator ancestry, where a + `Collect.item` input closes an axis), unless the input that decides the `If`'s frames carries both: the condition, + or the known selected branch of a literal condition. A cross product computed by one selected node therefore keeps + its frames, while one `Iterate` per branch, or a root value selected over a cross product, is rejected. A literal + condition whose selected branch carries one of the two also runs when the other comes only from the unselected + branch and either the unselected one reaches the `If` only through a `For` final output, which adds no frames, or + the selected axis, or an iterator enclosing it, reopens the unselected one. An `Iterate` reopens every axis of its + `Collect`'s item (the one the `Collect` closes and the ones it keeps), and transitively what those reopen, when + its `collection` comes directly from a `Collect` with a single `item` input and nothing appended through its + `collection`: each item becomes one frame, so the selected frames pair with those axes by index. Any other reopen + (a node between the `Collect` and the `Iterate`, or extra collected items) can change the length, and selecting + the upstream axis over its reopening pairs it with per-group frames; both would drop or invent frames. + - A branch whose open axes, or the axes a branch `Collect` keeps, are deeper than a per-item condition's, unless + both branches carry the same deeper axes. + - A root-frame condition choosing between branches with different deeper axes. When both branches share nested + deeper axes, a root-frame condition runs the `If` once per frame of the innermost one. + Selected-branch frame cardinality is not supported. +- A completed `Collect` on a branch closes its item axis; the `If` does not wait for an `Iterate` that only that + `Collect` consumes. When the `Collect` keeps an outer frame open and is selected, the `If` runs once per preserved + group; when it is rejected, it never runs and the `If` runs in the frames of its remaining iterators. A `Collect` of a + `For` final output groups in the `For`'s parent frame, so it keeps that frame open. +- A `For` whose loop body belongs to an `If` branch (its body feeds the `ForReturn`, which reaches the `For` only + through the loop linkage) is gated with that branch: for a frame that rejects the branch, neither the `For` nor its + body runs. A gated `For` brings in a `For` nested in its body together with that body. Two `For`s stay ungated. A + `For` nested in the body of an ungated `For` runs with that loop. A `For` whose final output feeds the `If` directly + is attached only after it finalizes, and the `If` is not prepared before that, so gating it on the `If` would + deadlock. Known limitation: with a per-item condition, such a `For` and its body run even for items that reject the + branch, and an error they raise is recorded. +- An `If` has no execution in a frame where the innermost iterator in its ancestry has no items. When an `Iterate` in + its iterator ancestry, whether it feeds `If.condition` or a branch, has no items at all, no `If` execution is created, + the branch sources are retired empty, the downstream `Collect` produces an empty collection, and the state completes + without an error. The same holds for a literal condition whose selected branch runs under an empty `Iterate` reached + only through a `For` final output. Retirement may add the branch source IDs to `executed_history`; neither branch + invocation runs and no branch result is created. Known limitation: with a wired condition, an `If` whose branch + reaches an empty `Iterate` only through a `For` final output stalls, whether the condition is root-frame or per item + of that `Iterate`. - The SQLite queue/processor path has evidence for cancellation before and after `If` resolution and retry from each boundary for both condition polarities. A canceled attempt keeps its activation ledger and cannot resume; retry starts a fresh state and must emit a fresh matching activation token before selected-only completion. diff --git a/invokeai/app/services/shared/graph_if_activation.py b/invokeai/app/services/shared/graph_if_activation.py index e9916db102a..1de7345e1d5 100644 --- a/invokeai/app/services/shared/graph_if_activation.py +++ b/invokeai/app/services/shared/graph_if_activation.py @@ -2,7 +2,7 @@ from invokeai.app.invocations.logic import IfInvocation from invokeai.app.services.shared.execution_engine.scheduler import ActivationDependency -from invokeai.app.services.shared.graph_validation import nx +from invokeai.app.services.shared.graph_validation import ForInvocation, IterateInvocation, nx if TYPE_CHECKING: from invokeai.app.services.shared.graph import GraphExecutionState @@ -33,16 +33,54 @@ def _branch_sources(self, if_node_id: str, branch_field: str, source_graph: "nx. branch_sources = set(direct_sources) for source_node_id in direct_sources: branch_sources.update(nx.ancestors(source_graph, source_node_id)) + # An iterator still open at the If gives the If its frames, whichever branch is selected, so it and its + # ancestors run regardless of the condition. + if any( + isinstance(self._state.graph.nodes.get(node_id), (ForInvocation, IterateInvocation)) + for node_id in branch_sources + ): + for iterator_id in self._state._runtime()._iterator_ancestry.get_frame_iterators(if_node_id): + branch_sources.discard(iterator_id) + branch_sources.difference_update(nx.ancestors(source_graph, iterator_id)) + # A For's loop body feeds its ForReturn, which reaches the For only through the loop linkage. The body belongs + # to the branch whenever its For does, and leaves it with the For. Fors are visited outermost first, so a gated + # For brings in a For nested in its body together with that body. Two Fors stay ungated. One whose final + # output enters the If directly is attached only after it finalizes, and the If is not prepared before that, + # so the For cannot wait for the If's verdict. One nested in an ungated For's body is expanded by that loop, + # which needs it to run. (A For enclosing the If is a frame iterator of the If and was removed above.) + for_by_return_id: dict[str, str] = {} + if any(isinstance(self._state.graph.nodes.get(node_id), ForInvocation) for node_id in branch_sources): + return_id_by_for = { + for_id: return_id for return_id, for_id in self._state._get_for_source_by_return_id().items() + } + body_ids_by_for = { + for_id: nx.ancestors(source_graph, return_id) & nx.descendants(source_graph, for_id) + for for_id, return_id in return_id_by_for.items() + } + for for_id in (node_id for node_id in nx.topological_sort(source_graph) if node_id in return_id_by_for): + if ( + for_id not in branch_sources + or for_id in direct_sources + or any( + for_id in body_ids_by_for[enclosing_id] and enclosing_id not in for_by_return_id.values() + for enclosing_id in body_ids_by_for + ) + ): + continue + for_by_return_id[return_id_by_for[for_id]] = for_id + branch_sources.update(body_ids_by_for[for_id]) + branch_sources.add(return_id_by_for[for_id]) changed = True while changed: changed = False for source_node_id in tuple(branch_sources): + linked_for_id = for_by_return_id.get(source_node_id) if all( edge.destination.node_id in branch_sources or (edge.destination.node_id == if_node_id and edge.destination.field == branch_field) for edge in self._state.graph._get_output_edges(source_node_id) - ): + ) and (linked_for_id is None or linked_for_id in branch_sources): continue branch_sources.remove(source_node_id) changed = True diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index ecee72c78ef..c9a75e9438b 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Callable, Optional +from invokeai.app.invocations.fields import OutputScope from invokeai.app.invocations.logic import IfInvocation from invokeai.app.services.shared.execution_engine.primitives import ActivationGate from invokeai.app.services.shared.execution_engine.scheduler import ActivationDependency @@ -13,6 +14,7 @@ CollectInvocation, ForInvocation, IterateInvocation, + get_output_field_scope, ) if TYPE_CHECKING: @@ -26,11 +28,13 @@ class UnsupportedIfIterationAxesError(ValueError): def validate_if_iteration_axes(state: "GraphExecutionState") -> None: """Reject unrelated condition and branch axes using graph topology alone.""" iterator_ancestry = _IteratorAncestry(state) + collect_axes_cache: dict[str, Optional[tuple[str, ...]]] = {} for node_id, node in state.graph.nodes.items(): if not isinstance(node, IfInvocation): continue condition_edges = state.graph._get_input_edges(node_id, "condition") if not condition_edges: + _validate_frame_axes(state, iterator_ancestry, node_id, condition_edges, collect_axes_cache) continue condition_iterators = iterator_ancestry.get_iterator_ancestors_for_edge(condition_edges[0]) branch_edges = { @@ -68,6 +72,229 @@ def validate_if_iteration_axes(state: "GraphExecutionState") -> None: f"unrelated open iterator axes: condition {sorted(condition_iterators)}, " f"branch {sorted(branch_iterators)}" ) + _validate_frame_axes(state, iterator_ancestry, node_id, condition_edges, collect_axes_cache) + + +def _validate_frame_axes( + state: "GraphExecutionState", + iterator_ancestry: _IteratorAncestry, + node_id: str, + condition_edges: list[Edge], + collect_axes_cache: dict[str, Optional[tuple[str, ...]]], +) -> None: + """Reject frame sets an If cannot run in. + + The If's frames come from the iterators open at it, including one reached through a For's final output. Two of + them that are unrelated (neither stays open inside the other) form a cross product, which only the input deciding + the If's frames can carry: the condition, or the known selected branch of a literal condition. A literal condition + whose selected branch carries one of the two axes also runs when the other comes only from the unselected branch + and either: + - the unselected one reaches the If only through a For's final output, so it adds no frames and the If runs in + the selected branch's frames; or + - the selected axis, or an iterator enclosing it, reopens the unselected one one-to-one (see + `_get_reopened_source_axes`), so the selected frames pair with it by index. Selecting the upstream axis over its + reopening is not supported. + A root-frame condition cannot select between branches with different deeper axes. + """ + if condition_edges: + _validate_collect_branch_axes(state, iterator_ancestry, node_id, condition_edges[0]) + frame_iterators = iterator_ancestry.get_frame_iterators(node_id) + if condition_edges: + deciding_edges = condition_edges + else: + node = state.graph.get_node(node_id) + assert isinstance(node, IfInvocation) + deciding_edges = state.graph._get_input_edges(node_id, "true_input" if node.condition else "false_input") + deciding_axes = {axis_id for edge in deciding_edges for axis_id in iterator_ancestry.get_frame_axes_for_edge(edge)} + for index, axis_id in enumerate(frame_iterators): + for other_axis_id in frame_iterators[index + 1 :]: + if axis_id in iterator_ancestry.get_frame_iterators( + other_axis_id + ) or other_axis_id in iterator_ancestry.get_frame_iterators(axis_id): + continue + if axis_id in deciding_axes and other_axis_id in deciding_axes: + continue + if not condition_edges and (axis_id in deciding_axes) != (other_axis_id in deciding_axes): + selected_id, unselected_id = ( + (axis_id, other_axis_id) if axis_id in deciding_axes else (other_axis_id, axis_id) + ) + if unselected_id not in iterator_ancestry.get_ordered_iterator_ancestors(node_id): + continue + if unselected_id in _get_reopened_source_axes( + state, iterator_ancestry, selected_id, collect_axes_cache + ): + continue + raise UnsupportedIfIterationAxesError( + f"If node '{node_id}' cannot run in the frames of unrelated open iterator axes: " + f"{sorted((axis_id, other_axis_id))}" + ) + if len(frame_iterators) < 2 or not condition_edges or iterator_ancestry.get_frame_axes_for_edge(condition_edges[0]): + return + branch_axes = { + branch_field: { + axis_id + for edge in state.graph._get_input_edges(node_id, branch_field) + for axis_id in iterator_ancestry.get_frame_axes_for_edge(edge) + } + for branch_field in ("true_input", "false_input") + } + if branch_axes["true_input"] != branch_axes["false_input"]: + raise UnsupportedIfIterationAxesError( + f"If node '{node_id}' cannot select between branches with different deeper open iterator axes against a " + f"root-frame condition: true {sorted(branch_axes['true_input'])}, " + f"false {sorted(branch_axes['false_input'])}" + ) + + +def _get_reopened_source_axes( + state: "GraphExecutionState", + iterator_ancestry: _IteratorAncestry, + axis_id: str, + collect_axes_cache: dict[str, Optional[tuple[str, ...]]], +) -> set[str]: + """Return the axes whose frames an axis, or an iterator enclosing it, reopens one-to-one. + + An Iterate reopens a Collect's item axes when its collection comes straight from that Collect, the Collect + gathers exactly one item per frame of them (a single item input and nothing appended through `collection`), and + those axes form one nested chain. Each item then becomes one frame of the Iterate, in the item frames' order, so a + selected axis can pair by index with those axes, and with the axes they in turn reopen. A node between the Collect + and the Iterate, extra collected items, or an item crossing unrelated axes can break that pairing, and the If + would then pair frames of unequal axes by index. + """ + sources: set[str] = set() + visited: set[str] = set() + pending = [axis_id] + while pending: + current_id = pending.pop() + for reopened_id in (current_id, *iterator_ancestry.get_frame_iterators(current_id)): + if reopened_id in visited: + continue + visited.add(reopened_id) + if not isinstance(state.graph.get_node(reopened_id), IterateInvocation): + continue + # Graph validation gives every Iterate exactly one collection input. + [collection_edge] = state.graph._get_input_edges(reopened_id, COLLECTION_FIELD) + collect_id = collection_edge.source.node_id + collect = state.graph.get_node(collect_id) + if not isinstance(collect, CollectInvocation) or collect.collection: + continue + collect_inputs = state.graph._get_input_edges(collect_id) + if len(collect_inputs) != 1 or collect_inputs[0].destination.field != ITEM_FIELD: + continue + item_axes = _get_edge_frame_axes(state, iterator_ancestry, collect_inputs[0], collect_axes_cache) + if item_axes is None or not _is_nested_chain(iterator_ancestry, item_axes): + continue + sources.update(item_axes) + pending.extend(item_axes) + return sources + + +def _validate_collect_branch_axes( + state: "GraphExecutionState", iterator_ancestry: _IteratorAncestry, node_id: str, condition_edge: Edge +) -> None: + """Reject a branch Collect that keeps axes open below a per-item condition unless the other branch shares them.""" + condition_axes = set(iterator_ancestry.get_frame_axes_for_edge(condition_edge)) + if not condition_axes: + return + branch_axes: dict[str, set[str]] = {} + for branch_field in ("true_input", "false_input"): + axes: set[str] = set() + for edge in state.graph._get_input_edges(node_id, branch_field): + if isinstance(state.graph.get_node(edge.source.node_id), CollectInvocation): + axes.update(_get_collect_kept_axes(state, iterator_ancestry, edge.source.node_id)) + else: + axes.update(iterator_ancestry.get_frame_axes_for_edge(edge)) + branch_axes[branch_field] = axes + for branch_field, other_field in (("true_input", "false_input"), ("false_input", "true_input")): + axes = branch_axes[branch_field] + if axes > condition_axes and axes != branch_axes[other_field]: + raise UnsupportedIfIterationAxesError( + f"If node '{node_id}' cannot select branch input '{branch_field}' from deeper open iterator axes " + f"kept by a Collect: condition {sorted(condition_axes)}, branch {sorted(axes)}" + ) + + +def _get_collect_kept_axes( + state: "GraphExecutionState", iterator_ancestry: _IteratorAncestry, collect_id: str +) -> set[str]: + """Return the axes a Collect's output stays in: its item's axes except the innermost one, which it closes.""" + kept: set[str] = set() + for edge in state.graph._get_input_edges(collect_id): + kept.update(_without_closed_axis(state, edge, iterator_ancestry.get_frame_axes_for_edge(edge))) + return kept + + +def _without_closed_axis(state: "GraphExecutionState", edge: Edge, axes: tuple[str, ...]) -> tuple[str, ...]: + """Return the axes a Collect input keeps: an item closes its innermost axis unless it is a For's final output.""" + source = state.graph.get_node(edge.source.node_id) + is_for_final = isinstance(source, ForInvocation) and ( + get_output_field_scope(source, edge.source.field) == OutputScope.Final + ) + if edge.destination.field == ITEM_FIELD and not is_for_final: + return axes[:-1] + return axes + + +def _is_nested_chain(iterator_ancestry: _IteratorAncestry, axes: tuple[str, ...]) -> bool: + """Return whether each axis encloses the next one.""" + return all( + outer_id in iterator_ancestry.get_frame_iterators(inner_id) + for outer_id, inner_id in zip(axes, axes[1:], strict=False) + ) + + +def _get_edge_frame_axes( + state: "GraphExecutionState", + iterator_ancestry: _IteratorAncestry, + edge: Edge, + collect_axes_cache: dict[str, Optional[tuple[str, ...]]], +) -> Optional[tuple[str, ...]]: + """Return the frames open on an edge, outer to inner, following a Collect source into the axes it keeps. + + Iterator ancestry stops at a Collect, so a Collect's output would otherwise look like it runs in no frames. A + Collect whose kept axes are not one nested chain, or that collects such a Collect, has no frame order a reopen could + pair with, so it yields None. Results are cached per Collect for one validation run and resolved with an explicit + stack, so a lattice or a long chain of Collects costs one pass per node. + """ + if not isinstance(state.graph.get_node(edge.source.node_id), CollectInvocation): + return iterator_ancestry.get_frame_axes_for_edge(edge) + pending = [edge.source.node_id] + while pending: + collect_id = pending[-1] + if collect_id in collect_axes_cache: + pending.pop() + continue + input_edges = state.graph._get_input_edges(collect_id) + unresolved = [ + input_edge.source.node_id + for input_edge in input_edges + if isinstance(state.graph.get_node(input_edge.source.node_id), CollectInvocation) + and input_edge.source.node_id not in collect_axes_cache + ] + if unresolved: + pending.extend(unresolved) + continue + pending.pop() + kept: Optional[set[str]] = set() + for input_edge in input_edges: + source_id = input_edge.source.node_id + if isinstance(state.graph.get_node(source_id), CollectInvocation): + input_axes = collect_axes_cache[source_id] + else: + input_axes = iterator_ancestry.get_frame_axes_for_edge(input_edge) + if input_axes is None: + kept = None + break + kept.update(_without_closed_axis(state, input_edge, input_axes)) + ordered = None + if kept is not None: + ordered = tuple( + sorted(kept, key=lambda axis_id: (len(iterator_ancestry.get_frame_iterators(axis_id)), axis_id)) + ) + if not _is_nested_chain(iterator_ancestry, ordered): + ordered = None + collect_axes_cache[collect_id] = ordered + return collect_axes_cache[edge.source.node_id] def _get_effective_iteration_paths_for_edge( @@ -130,7 +357,12 @@ def _get_effective_iteration_paths_for_edge( cache, ) group_depth = len(iterator_ancestry.get_iterator_axes_for_edge(input_edge)) - if input_edge.destination.field == ITEM_FIELD: + input_source = state.graph.get_node(input_edge.source.node_id) + # A For's final output already excludes the For's own axis, so the Collect closes nothing more. + is_for_final = isinstance(input_source, ForInvocation) and ( + get_output_field_scope(input_source, input_edge.source.field) == OutputScope.Final + ) + if input_edge.destination.field == ITEM_FIELD and not is_for_final: group_depth = max(group_depth - 1, 0) collapsed_paths.update(path[:group_depth] for path in input_paths) cache[cache_key] = collapsed_paths or ({()} if not input_edges else set()) @@ -217,6 +449,17 @@ def _get_if_iteration_paths_with_ancestry( else: branch_paths = {condition_path or ()} + if ( + not branch_paths + and condition_path is not None + and any( + edge.source.node_id not in state.source_prepared_mapping and edge.source.node_id not in state.executed + for edge in selected_edges + ) + ): + # The selected branch is not materialized yet; the If still runs in the condition's frame. + branch_paths = {condition_path} + for branch_path in branch_paths: if condition_path is None: paths.add(branch_path) @@ -225,30 +468,31 @@ def _get_if_iteration_paths_with_ancestry( elif branch_path[: len(condition_path)] == condition_path: paths.add(branch_path) - if not condition_edges: - frame_iterators = iterator_ancestry.get_ordered_iterator_ancestors(node_id) - if frame_iterators: - frame_paths: set[tuple[int, ...]] = set() - for iterator_id in reversed(frame_iterators): - if iterator_id not in state.source_prepared_mapping: - if iterator_id not in state.executed: - return [] - continue - prepared_ids = get_prepared_nodes_for_source(iterator_id) - iterator_edges = state.graph._get_output_edges(iterator_id, ITEM_FIELD) - if prepared_ids and iterator_edges: - frame_paths = { - get_prepared_edge_iteration_path(iterator_edges[0], prepared_id) for prepared_id in prepared_ids - } - break - if not frame_paths: - return [] - paths = { - frame_path if frame_path[: len(branch_path)] == branch_path else branch_path - for branch_path in paths - for frame_path in frame_paths - if frame_path[: len(branch_path)] == branch_path or branch_path[: len(frame_path)] == frame_path - } + # The If runs once per frame of the innermost iterator open at it; validation keeps those iterators nested. An + # iterator that has not run leaves the frames unknown, and one with no items leaves the If no frames. An iterator + # reaching the If only through a For's final output does not add frames here. + frame_iterators = iterator_ancestry.get_ordered_iterator_ancestors(node_id) + if frame_iterators: + innermost_id = frame_iterators[-1] + if innermost_id not in state.executed and innermost_id not in state.source_prepared_mapping: + return [] + prepared_ids = ( + get_prepared_nodes_for_source(innermost_id) if innermost_id in state.source_prepared_mapping else set() + ) + iterator_edges = state.graph._get_output_edges(innermost_id, ITEM_FIELD) + if not prepared_ids: + return [] + frame_paths = ( + {get_prepared_edge_iteration_path(iterator_edges[0], prepared_id) for prepared_id in prepared_ids} + if iterator_edges + else {state._get_iteration_path(prepared_id) for prepared_id in prepared_ids} + ) + paths = { + frame_path if frame_path[: len(branch_path)] == branch_path else branch_path + for branch_path in paths + for frame_path in frame_paths + if frame_path[: len(branch_path)] == branch_path or branch_path[: len(frame_path)] == frame_path + } if not paths: if not has_selected_edges: @@ -392,7 +636,7 @@ def _is_source_inactive( frames = {frame for dependency in dependencies for frame in state._prepared_if_exec_frames(dependency.owner_id)} if not frames: - return all(_is_empty_if_condition(state, dependency) for dependency in dependencies) + return all(_has_empty_if_frame(state, dependency.owner_id) for dependency in dependencies) source_axes = set(state._runtime()._get_ordered_iterator_sources(source_node_id)) for dependency in dependencies: @@ -407,20 +651,16 @@ def _is_source_inactive( ) -def _is_empty_if_condition(state: "GraphExecutionState", dependency: ActivationDependency) -> bool: - """Recognize a branch whose If never materialized because its Iterate was empty.""" +def _has_empty_if_frame(state: "GraphExecutionState", if_node_id: str) -> bool: + """Recognize an If that has no executions because an Iterate open at it had no items at all.""" - owner = state.graph.nodes.get(dependency.owner_id) - if not isinstance(owner, IfInvocation): + if not isinstance(state.graph.nodes.get(if_node_id), IfInvocation): return False - condition_edges = state.graph._get_input_edges(dependency.owner_id, "condition") - if len(condition_edges) != 1: - return False - condition_source_id = condition_edges[0].source.node_id - return ( - isinstance(state.graph.get_node(condition_source_id), IterateInvocation) - and condition_source_id in state.executed - and not state.source_prepared_mapping.get(condition_source_id) + return any( + isinstance(state.graph.get_node(iterator_id), IterateInvocation) + and iterator_id in state.executed + and not state.source_prepared_mapping.get(iterator_id) + for iterator_id in state._runtime()._iterator_ancestry.get_frame_iterators(if_node_id) ) diff --git a/invokeai/app/services/shared/graph_iterator_ancestry.py b/invokeai/app/services/shared/graph_iterator_ancestry.py index 3ecab6a1b2c..e28570dbe43 100644 --- a/invokeai/app/services/shared/graph_iterator_ancestry.py +++ b/invokeai/app/services/shared/graph_iterator_ancestry.py @@ -24,6 +24,7 @@ def __init__( self._iterator_graph = iterator_graph self._has_iterator_ancestor: dict[str, bool] | None = None self._ordered_iterator_ancestors: dict[str, tuple[str, ...]] = {} + self._frame_iterators: dict[str, tuple[str, ...]] = {} self.effective_iteration_paths_by_edge: dict[tuple[str, str], set[tuple[int, ...]]] = {} def clear_effective_iteration_paths(self) -> None: @@ -108,3 +109,61 @@ def get_iterator_axes_for_edge(self, edge: "Edge") -> tuple[str, ...]: ): ancestors.append(source_node_id) return tuple(ancestors) + + def get_frame_iterators(self, node_id: str) -> tuple[str, ...]: + """Return the iterators whose frames are open at a node, outer to inner in input-port order. + + Unlike ``get_ordered_iterator_ancestors``, a For's final output is followed into the For's parent frame: the + For's own axis is closed there, but the axes the For runs under stay open. + """ + cached = self._frame_iterators.get(node_id) + if cached is not None: + return cached + self.has_iterator_ancestor(node_id) + assert self._iterator_graph is not None + + ordered: list[str] = [] + visited: set[tuple[str, bool]] = set() + # (node, exiting, opens): ``opens`` is False when the node is reached only through a For's final output. + pending: list[tuple[str, bool, bool]] = [(node_id, False, False)] + while pending: + current_node_id, exiting, opens = pending.pop() + if exiting: + if opens and current_node_id not in ordered: + ordered.append(current_node_id) + continue + if (current_node_id, opens) in visited: + continue + visited.add((current_node_id, opens)) + pending.append((current_node_id, True, opens)) + current_node = self._state.graph.get_node(current_node_id) + field_order = {field: index for index, field in enumerate(type(current_node).model_fields)} + input_edges = sorted( + self._state.graph._get_input_edges(current_node_id), + key=lambda edge: (field_order.get(edge.destination.field, len(field_order)), edge.source.node_id), + ) + for edge in reversed(input_edges): + source_node_id = edge.source.node_id + source_node = self._state.graph.get_node(source_node_id) + if self._iterator_graph.has_edge(source_node_id, current_node_id): + pending.append((source_node_id, False, isinstance(source_node, (ForInvocation, IterateInvocation)))) + elif ( + isinstance(source_node, ForInvocation) + and get_output_field_scope(source_node, edge.source.field) == OutputScope.Final + ): + pending.append((source_node_id, False, False)) + + result = tuple(ordered) + self._frame_iterators[node_id] = result + return result + + def get_frame_axes_for_edge(self, edge: "Edge") -> tuple[str, ...]: + """Return the iterator frames open on an edge, following For final outputs into their parent frame.""" + axes = list(self.get_frame_iterators(edge.source.node_id)) + source_node = self._state.graph.get_node(edge.source.node_id) + if isinstance(source_node, (ForInvocation, IterateInvocation)) and not ( + isinstance(source_node, ForInvocation) + and get_output_field_scope(source_node, edge.source.field) == OutputScope.Final + ): + axes.append(edge.source.node_id) + return tuple(axes) diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 0f6ca9cbece..5d8c8c12676 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -1688,16 +1688,20 @@ def _has_admitted_source_mapping( if isinstance(self._state.graph.get_node(node_id), IfInvocation): condition_edges = self._state.graph._get_input_edges(node_id, "condition") if not condition_edges: + iteration_paths = graph_if_runtime._get_if_iteration_paths_with_ancestry( + self._state, + node_id, + set(), + self._get_prepared_nodes_for_source, + self._get_prepared_edge_iteration_path, + iterator_ancestry, + ) + if not iteration_paths: + # An empty Iterate giving the If its frames leaves it no executions; preparing it closes it empty. + return graph_if_runtime._has_empty_if_frame(self._state, node_id) return any( self._state._is_source_activation_admitted(node_id, iteration_path) - for iteration_path in graph_if_runtime._get_if_iteration_paths_with_ancestry( - self._state, - node_id, - set(), - self._get_prepared_nodes_for_source, - self._get_prepared_edge_iteration_path, - iterator_ancestry, - ) + for iteration_path in iteration_paths ) mappings = self._get_if_condition_iteration_mappings(node_id, graph) input_edges = condition_edges @@ -1709,7 +1713,7 @@ def _has_admitted_source_mapping( node_id, self.iterator_graph(graph) ): if not any(mappings): - iterator_nodes = self.get_node_iterators(node_id, self.iterator_graph(graph)) + iterator_nodes = self._get_if_frame_iterator_nodes(node_id, self.iterator_graph(graph)) if all( iterator_node_id in self._state.executed and not self._state.source_prepared_mapping.get(iterator_node_id) @@ -1726,6 +1730,20 @@ def _has_admitted_source_mapping( for iteration_mapping in mappings ) + def _get_if_frame_iterator_nodes(self, node_id: str, iterator_graph: "nx.DiGraph") -> list[str]: + """Return the If's branch iterators, reduced to the innermost one when they are nested. + + Each frame of the innermost nested axis also fixes the outer indexes, so that axis alone maps the If's frames. + """ + iterator_nodes = self.get_node_iterators(node_id, iterator_graph) + iterator_ancestry = self._get_iterator_ancestry(iterator_graph) + if all( + outer_id in iterator_ancestry.get_ordered_iterator_ancestors(inner_id) + for outer_id, inner_id in itertools.pairwise(iterator_nodes) + ): + return iterator_nodes[-1:] + return iterator_nodes + def _get_if_condition_iteration_mappings( self, node_id: str, graph: "nx.DiGraph" ) -> Iterable[list[tuple[str, str]]]: @@ -1756,9 +1774,10 @@ def _get_if_condition_iteration_mappings( # This specialization derives one branch axis against a root-frame condition. Mixed or nested condition # axes use the general parent-mapping path. + iterator_ancestry = self._get_iterator_ancestry(iterator_graph) + iterator_nodes = self._get_if_frame_iterator_nodes(node_id, iterator_graph) if len(iterator_nodes) != 1: return iter([[]]) - iterator_ancestry = self._get_iterator_ancestry(iterator_graph) for branch_field in ("true_input", "false_input"): branch_edges = self._state.graph._get_input_edges(node_id, branch_field) branch_iterator_ids = { @@ -1775,7 +1794,9 @@ def _get_if_condition_iteration_mappings( ] if condition_edges: - condition_mappings = self._get_parent_iteration_mappings(node_id, graph, input_edges=condition_edges) + condition_mappings = self._get_parent_iteration_mappings( + node_id, graph, input_edges=condition_edges, iterator_node_ids=iterator_nodes + ) return iter( condition_mapping + list(zip(iterator_nodes, prepared_iterators, strict=True)) for condition_mapping, prepared_iterators in zip( @@ -1801,15 +1822,30 @@ def _is_if_condition_ready(self, node_id: str) -> bool: return True def _get_if_readiness_ancestors(self, node_id: str, graph: "nx.DiGraph") -> set[str]: - """Get iterator ancestors that must finish before preparing an If execution.""" + """Get iterator ancestors that must finish before preparing an If execution. + + A wired condition gates the If through its ancestry. A literal condition's known selected branch supplies the + value; an iterator only on the unselected branch is gated by the If and must not hold it back. A wired If also + waits for the iterators in its own ancestry: they give it its frames and are never gated by it, and preparing + the If before they finish would find no frames and close it empty. + """ condition_edges = self._state.graph._get_input_edges(node_id, "condition") - if not condition_edges: - return set(nx.ancestors(graph, node_id)) + if condition_edges: + gating_source_ids = [ + *(edge.source.node_id for edge in condition_edges), + *self.get_node_iterators(node_id, self.iterator_graph(graph)), + ] + else: + node = self._state.graph.get_node(node_id) + assert isinstance(node, IfInvocation) + selected_field = "true_input" if node.condition else "false_input" + gating_source_ids = [ + edge.source.node_id for edge in self._state.graph._get_input_edges(node_id, selected_field) + ] ancestors: set[str] = set() - for edge in condition_edges: - condition_source_id = edge.source.node_id - ancestors.update(nx.ancestors(graph, condition_source_id)) - ancestors.add(condition_source_id) + for source_node_id in gating_source_ids: + ancestors.update(nx.ancestors(graph, source_node_id)) + ancestors.add(source_node_id) return ancestors def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, Optional[str]]]: diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index 1e40a62a898..38a42f86ed2 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -2,6 +2,7 @@ import copy import inspect +import itertools import json import random import sys @@ -3259,6 +3260,1520 @@ def test_nested_iterate_keeps_values_in_outer_frame_with_unrelated_if( assert state.is_complete() +def _collected_iterate_branch_graph(condition: bool) -> Graph: + """src -> iterate -> body -> collect -> if.true_input, the shape reported in #9484.""" + graph = Graph() + graph.add_node(FixedCollectionTestInvocation(id="src", value=True)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("src", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +def _doubly_collected_nested_iterate_branch_graph(condition: bool) -> Graph: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=["x", "y"])) + graph.add_node(IterateInvocation(id="outer_iterate")) + graph.add_node(FixedCollectionTestInvocation(id="inner_source")) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(CollectInvocation(id="outer_collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("outer_source", "collection", "outer_iterate", "collection")) + graph.add_edge(create_edge("outer_iterate", "item", "inner_source", "value")) + graph.add_edge(create_edge("inner_source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "outer_collect", "item")) + graph.add_edge(create_edge("outer_collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +def _for_output_iterate_branch_graph(condition: bool) -> Graph: + graph = Graph() + graph.add_node(ForInvocation(id="for", collection=["a", "b"])) + graph.add_node(ForReturnInvocation(id="return")) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("for", "item", "return", "output")) + graph.add_edge(create_loop_linkage("for", "return")) + graph.add_edge(create_edge("for", "output_collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +@pytest.mark.parametrize( + "build_graph", + [_collected_iterate_branch_graph, _doubly_collected_nested_iterate_branch_graph, _for_output_iterate_branch_graph], + ids=["one-collect", "two-collects", "for-output"], +) +def test_literal_if_rejecting_closed_iterate_branch_completes(build_graph: Callable[[bool], Graph]) -> None: + """A literal If does not wait for an Iterate that closes on its rejected branch, which never runs (#9484).""" + graph = build_graph(False) + graph.validate_self() + + trace, state = _run(GraphExecutionState(graph=graph)) + + # This shape is outside the generic scheduler's supported set. + assert isinstance(state._execution_scheduler, _ExecutionScheduler) + assert state.is_complete(), f"execution stalled after {trace}" + sink_exec_ids = state._prepared_registry().get_prepared_ids("sink") + assert [state.results[exec_id].value for exec_id in sink_exec_ids] == [["fallback"]] + assert not {"for", "return", "iterate", "body", "collect"} & set(trace) + + +def _per_item_for_output_collect_branch_graph() -> Graph: + """A Collect groups a For's final output in the For's parent frame, so the outer item frame stays open.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=[["a", "b"], ["c"]])) + graph.add_node(IterateInvocation(id="outer_iterate")) + graph.add_node(ForInvocation(id="for")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(ForReturnInvocation(id="return")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=True)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("outer_source", "collection", "outer_iterate", "collection")) + graph.add_edge(create_edge("outer_iterate", "item", "for", "collection")) + graph.add_edge(create_edge("for", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "return", "output")) + graph.add_edge(create_loop_linkage("for", "return")) + graph.add_edge(create_edge("for", "output_collection", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +def test_literal_if_selecting_collect_of_for_output_runs_per_outer_item() -> None: + """A Collect of a For's final output stays in the outer item frame, so the If runs once per outer item.""" + graph = _per_item_for_output_collect_branch_graph() + graph.validate_self() + + trace, state = _run(GraphExecutionState(graph=graph)) + + # This shape is outside the generic scheduler's supported set. + assert isinstance(state._execution_scheduler, _ExecutionScheduler) + assert state.is_complete(), f"execution stalled after {trace}" + assert { + state._get_iteration_path(exec_id): state.results[exec_id].value + for exec_id in state._prepared_registry().get_prepared_ids("sink") + } == {(0,): [["a", "b"]], (1,): [["c"]]} + + +@pytest.mark.parametrize("fail_unselected", [False, True], ids=["succeeds", "unselected-fails"]) +def test_per_item_if_does_not_run_for_output_branch_for_unselected_item(fail_unselected: bool) -> None: + """A For whose final output is collected into one If branch runs, with its body, only for selecting items.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=["", "x"])) + graph.add_node(IterateInvocation(id="outer_iterate")) + graph.add_node(FixedCollectionTestInvocation(id="items")) + graph.add_node(ForInvocation(id="for")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(ForReturnInvocation(id="return")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("outer_source", "collection", "outer_iterate", "collection")) + graph.add_edge(create_edge("outer_iterate", "item", "items", "value")) + graph.add_edge(create_edge("items", "condition", "if", "condition")) + graph.add_edge(create_edge("items", "collection", "for", "collection")) + graph.add_edge(create_edge("for", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "return", "output")) + graph.add_edge(create_loop_linkage("for", "return")) + graph.add_edge(create_edge("for", "output_collection", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + graph.validate_self() + unselected_sources = {"for", "body", "return", "collect"} + + state = GraphExecutionState(graph=graph) + trace: list[tuple[str, tuple[int, ...]]] = [] + while (node := state.next()) is not None: + source_id = state.prepared_source_mapping[node.id] + iteration_path = state._get_iteration_path(node.id) + trace.append((source_id, iteration_path)) + if fail_unselected and source_id in unselected_sources and iteration_path[:1] == (0,): + state.set_node_error(node.id, f"{source_id} ran for an unselected item") + continue + state.complete(node.id, node.invoke(Mock())) + + # This shape is outside the generic scheduler's supported set. + assert isinstance(state._execution_scheduler, _ExecutionScheduler) + assert not state.has_error(), state.errors + assert state.is_complete(), f"execution stalled after {trace}" + assert not [entry for entry in trace if entry[0] in unselected_sources and entry[1][:1] == (0,)] + assert {source_id for source_id, path in trace if path[:1] == (1,)} >= {"for", "body", "collect"} + assert { + state._get_iteration_path(exec_id): state.results[exec_id].value + for exec_id in state._prepared_registry().get_prepared_ids("sink") + } == {(0,): "fallback", (1,): [[0, 1]]} + + +def _graph_from(nodes: list[Any], edges: list[Any]) -> Graph: + graph = Graph() + for node in nodes: + graph.add_node(node) + for edge in edges: + graph.add_edge(edge if isinstance(edge, Edge) else create_edge(*edge)) + graph.validate_self() + return graph + + +def _run_to_completion(graph: Graph) -> tuple[list[str], GraphExecutionState]: + trace, state = _run(GraphExecutionState(graph=graph)) + # These shapes are outside the generic scheduler's supported set. + assert isinstance(state._execution_scheduler, _ExecutionScheduler) + assert not state.has_error(), state.errors + assert state.is_complete(), f"execution stalled after {trace}" + return trace, state + + +def _source_values(state: GraphExecutionState, source_id: str) -> dict[tuple[int, ...], Any]: + values = {} + for exec_id in state._prepared_registry().get_prepared_ids(source_id): + output = state.results[exec_id] + values[state._get_iteration_path(exec_id)] = getattr(output, "collection", getattr(output, "value", None)) + return values + + +@pytest.mark.parametrize("computed_condition", [False, True], ids=["literal", "computed"]) +def test_if_in_for_body_rejecting_nested_for_branch_completes(computed_condition: bool) -> None: + """A For nested in an outer For's body runs with its loop even when it feeds only a rejected If branch.""" + nodes = [ + CollectionConcatInvocation(id="outer_src", first=[["a", "b"], ["c"]]), + ForInvocation(id="for"), + ForInvocation(id="ifor"), + MarkedAnyTypeTestInvocation(id="ibody", marker="i:"), + ForReturnInvocation(id="iret"), + AnyTypeTestInvocation(id="mid"), + AnyTypeTestInvocation(id="fallback", value="fb"), + IfInvocation(id="if", condition=False), + ForReturnInvocation(id="ret"), + AnyTypeTestInvocation(id="sink"), + ] + edges = [ + ("outer_src", "collection", "for", "collection"), + ("for", "item", "ifor", "collection"), + ("ifor", "item", "ibody", "value"), + ("ibody", "value", "iret", "output"), + create_loop_linkage("ifor", "iret"), + ("ifor", "output_collection", "mid", "value"), + ("mid", "value", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "ret", "output"), + create_loop_linkage("for", "ret"), + ("for", "output_collection", "sink", "value"), + ] + if computed_condition: + nodes.append(FixedCollectionTestInvocation(id="condition", value=False)) + edges.append(("condition", "condition", "if", "condition")) + + _trace, state = _run_to_completion(_graph_from(nodes, edges)) + + assert _source_values(state, "sink") == {(): ["fb", "fb"]} + + +@pytest.mark.parametrize( + ("condition", "computed_condition"), + [(True, True), (False, False), (False, True)], + ids=["computed-True", "literal-False", "computed-False"], +) +def test_root_condition_if_with_for_output_branch_under_iterate_completes( + condition: bool, computed_condition: bool +) -> None: + """The Iterate a branch For runs under is not gated by the If, and the rejected For never runs.""" + nodes = [ + CollectionConcatInvocation(id="src", first=[["a", "b"], ["c"]]), + IterateInvocation(id="iterate"), + ForInvocation(id="for"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + ForReturnInvocation(id="ret"), + AnyTypeTestInvocation(id="mid"), + AnyTypeTestInvocation(id="fallback", value="fb"), + IfInvocation(id="if", condition=condition), + CollectInvocation(id="collect"), + ] + edges = [ + ("src", "collection", "iterate", "collection"), + ("iterate", "item", "for", "collection"), + ("for", "item", "body", "value"), + ("body", "value", "ret", "output"), + create_loop_linkage("for", "ret"), + ("for", "output_collection", "mid", "value"), + ("mid", "value", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ] + if computed_condition: + nodes.append(FixedCollectionTestInvocation(id="condition", value=condition)) + edges.append(("condition", "condition", "if", "condition")) + + trace, state = _run_to_completion(_graph_from(nodes, edges)) + + # Selecting the root value keeps main's single root-frame If; selected-branch cardinality is decided separately. + expected = [["b:a", "b:b"], ["b:c"]] if condition else ["fb"] + assert _source_values(state, "collect") == {(): expected} + if not condition: + assert not {"for", "body", "ret", "mid"} & set(trace) + + +@pytest.mark.parametrize("condition", [True, False]) +def test_literal_if_with_empty_nested_iterate_on_a_branch_has_no_frames(condition: bool) -> None: + """An open Iterate with no items gives the If no frame, whichever branch the literal condition selects.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="src", first=[[], []]), + IterateInvocation(id="outer"), + CollectionConcatInvocation(id="inner_src", first=[]), + IterateInvocation(id="inner"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=condition), + CollectInvocation(id="collect"), + ], + [ + ("src", "collection", "outer", "collection"), + ("outer", "item", "inner_src", "first"), + ("inner_src", "collection", "inner", "collection"), + ("inner", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("outer", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + + trace, state = _run_to_completion(graph) + + assert _source_values(state, "collect") == {(0,): [], (1,): []} + assert "if" not in trace + + +def _unrelated_open_branches_graph(condition: bool, computed_condition: bool) -> Graph: + nodes = [ + CollectionConcatInvocation(id="true_src", first=[1, 2]), + IterateInvocation(id="true_iterate"), + CollectionConcatInvocation(id="false_src", first=[10, 20, 30]), + IterateInvocation(id="false_iterate"), + IfInvocation(id="if", condition=condition), + CollectInvocation(id="collect"), + ] + edges = [ + ("true_src", "collection", "true_iterate", "collection"), + ("false_src", "collection", "false_iterate", "collection"), + ("true_iterate", "item", "if", "true_input"), + ("false_iterate", "item", "if", "false_input"), + ("if", "value", "collect", "item"), + ] + if computed_condition: + nodes.append(FixedCollectionTestInvocation(id="condition", value=condition)) + edges.append(("condition", "condition", "if", "condition")) + return _graph_from(nodes, edges) + + +def _root_condition_with_deeper_axis_on_one_branch_graph(outer_values: list[list[str]]) -> Graph: + return _graph_from( + [ + CollectionConcatInvocation(id="src", first=outer_values), + IterateInvocation(id="outer"), + CollectionConcatInvocation(id="inner_src", first=[]), + IterateInvocation(id="inner"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + FixedCollectionTestInvocation(id="condition", value=False), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ], + [ + ("src", "collection", "outer", "collection"), + ("outer", "item", "inner_src", "first"), + ("inner_src", "collection", "inner", "collection"), + ("inner", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("outer", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("condition", "condition", "if", "condition"), + ("if", "value", "collect", "item"), + ], + ) + + +def _per_item_condition_with_deeper_collect_branch_graph() -> Graph: + return _graph_from( + [ + CollectionConcatInvocation(id="src", first=["x", "y"]), + IterateInvocation(id="outer"), + FixedCollectionTestInvocation(id="middle_src"), + IterateInvocation(id="middle"), + FixedCollectionTestInvocation(id="inner_src"), + IterateInvocation(id="inner"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + CollectInvocation(id="branch_collect"), + AnyTypeTestInvocation(id="fallback", value="fb"), + FixedCollectionTestInvocation(id="condition"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ], + [ + ("src", "collection", "outer", "collection"), + ("outer", "item", "middle_src", "value"), + ("middle_src", "collection", "middle", "collection"), + ("middle", "item", "inner_src", "value"), + ("inner_src", "collection", "inner", "collection"), + ("inner", "item", "body", "value"), + ("body", "value", "branch_collect", "item"), + ("branch_collect", "collection", "if", "true_input"), + ("outer", "item", "fallback", "value"), + ("fallback", "value", "if", "false_input"), + ("outer", "item", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("if", "value", "collect", "item"), + ], + ) + + +def _per_item_condition_with_deeper_for_output_collect_branch_graph() -> Graph: + return _graph_from( + [ + CollectionConcatInvocation(id="src", first=["", "x"]), + IterateInvocation(id="outer"), + FixedCollectionTestInvocation(id="condition"), + FixedCollectionTestInvocation(id="inner_src"), + IterateInvocation(id="inner"), + FixedCollectionTestInvocation(id="items"), + ForInvocation(id="for"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + ForReturnInvocation(id="ret"), + CollectInvocation(id="branch_collect"), + MarkedAnyTypeTestInvocation(id="fallback", marker="f:"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ], + [ + ("src", "collection", "outer", "collection"), + ("outer", "item", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("outer", "item", "inner_src", "value"), + ("inner_src", "collection", "inner", "collection"), + ("inner", "item", "items", "value"), + ("items", "collection", "for", "collection"), + ("for", "item", "body", "value"), + ("body", "value", "ret", "output"), + create_loop_linkage("for", "ret"), + ("for", "output_collection", "branch_collect", "item"), + ("branch_collect", "collection", "if", "true_input"), + ("outer", "item", "fallback", "value"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + + +def _cross_product_branch_graph(condition: bool, condition_from_branch: bool = False) -> Graph: + nodes = [ + CollectionConcatInvocation(id="left_src", first=[1, 2]), + IterateInvocation(id="left"), + CollectionConcatInvocation(id="right_src", first=[10, 20, 30]), + IterateInvocation(id="right"), + AddInvocation(id="add"), + AnyTypeTestInvocation(id="fallback", value=0), + IfInvocation(id="if", condition=condition), + CollectInvocation(id="collect"), + ] + edges = [ + ("left_src", "collection", "left", "collection"), + ("right_src", "collection", "right", "collection"), + ("left", "item", "add", "a"), + ("right", "item", "add", "b"), + ("add", "value", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ] + if condition_from_branch: + nodes.append(FixedCollectionTestInvocation(id="condition")) + edges += [("add", "value", "condition", "value"), ("condition", "condition", "if", "condition")] + return _graph_from(nodes, edges) + + +def _collect_then_reopen_graph( + condition: bool, extra_items: tuple[str, ...] = (), appended_items: tuple[str, ...] = () +) -> Graph: + nodes = [ + CollectionConcatInvocation(id="src", first=["a1", "a2"]), + IterateInvocation(id="first"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + CollectInvocation(id="closing_collect"), + IterateInvocation(id="second"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=condition), + CollectInvocation(id="collect"), + ] + edges = [ + ("src", "collection", "first", "collection"), + ("first", "item", "body", "value"), + ("body", "value", "closing_collect", "item"), + ] + if appended_items: + nodes.append(CollectionConcatInvocation(id="appended", first=list(appended_items))) + edges.append(("appended", "collection", "closing_collect", "collection")) + if extra_items: + nodes.append(CollectionConcatInvocation(id="extend", second=list(extra_items))) + edges += [ + ("closing_collect", "collection", "extend", "first"), + ("extend", "collection", "second", "collection"), + ] + else: + edges.append(("closing_collect", "collection", "second", "collection")) + edges += [ + ("second", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("first", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ] + return _graph_from(nodes, edges) + + +def _collect_with_literal_items_then_reopen_graph() -> Graph: + graph = _collect_then_reopen_graph(True) + graph.get_node("closing_collect").collection = ["x"] + return graph + + +def _collect_of_collect_reopen_graph( + groups: list[list[str]], unselected_axis: str, extra_collect: bool = False, append_unrelated_axis: bool = False +) -> Graph: + """An Iterate reopens a Collect of per-group Collects, so its frames are the groups of the outer axis.""" + nodes: list[Any] = [ + CollectionConcatInvocation(id="groups_src", first=groups), + IterateInvocation(id="group"), + CollectionConcatInvocation(id="members_src"), + IterateInvocation(id="member"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + CollectInvocation(id="group_collect"), + CollectInvocation(id="closing_collect"), + IterateInvocation(id="reopened"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ] + edges: list[Any] = [ + ("groups_src", "collection", "group", "collection"), + ("group", "item", "members_src", "first"), + ("members_src", "collection", "member", "collection"), + ("member", "item", "body", "value"), + ("body", "value", "group_collect", "item"), + ("group_collect", "collection", "closing_collect", "item"), + ] + if append_unrelated_axis: + nodes += [ + CollectionConcatInvocation(id="appended_src", first=[["v1"], ["v2"], ["v3"]]), + IterateInvocation(id="appended"), + CollectionConcatInvocation(id="appended_items"), + ] + edges += [ + ("appended_src", "collection", "appended", "collection"), + ("appended", "item", "appended_items", "first"), + ("appended_items", "collection", "group_collect", "collection"), + ] + if extra_collect: + nodes.append(CollectInvocation(id="outer_collect")) + edges += [ + ("closing_collect", "collection", "outer_collect", "item"), + ("outer_collect", "collection", "reopened", "collection"), + ] + else: + edges.append(("closing_collect", "collection", "reopened", "collection")) + edges += [ + ("reopened", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + (unselected_axis, "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ] + return _graph_from(nodes, edges) + + +def _per_item_condition_with_collect_of_collects_branch_graph() -> Graph: + """A per-item condition on the outer axis chooses between a Collect over a deeper axis and a Collect of one.""" + nodes: list[Any] = [ + CollectionConcatInvocation(id="src", first=[[[["a"], ["b"]], [["c"]]], [[["x"]], []], [[["d", "e"]]]]), + ] + edges: list[Any] = [] + previous = ("src", "collection") + for axis_id in ("outer", "middle", "deep", "inner"): + nodes.append(IterateInvocation(id=axis_id)) + edges.append((*previous, axis_id, "collection")) + if axis_id != "inner": + nodes.append(CollectionConcatInvocation(id=f"{axis_id}_items")) + edges.append((axis_id, "item", f"{axis_id}_items", "first")) + previous = (f"{axis_id}_items", "collection") + nodes += [ + MarkedAnyTypeTestInvocation(id="inner_body", marker="b:"), + CollectInvocation(id="inner_collect"), + CollectInvocation(id="deep_collect"), + MarkedAnyTypeTestInvocation(id="deep_body", marker="n:"), + CollectInvocation(id="deep_item_collect"), + FixedCollectionTestInvocation(id="condition"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ] + edges += [ + ("inner", "item", "inner_body", "value"), + ("inner_body", "value", "inner_collect", "item"), + ("inner_collect", "collection", "deep_collect", "item"), + ("deep", "item", "deep_body", "value"), + ("deep_body", "value", "deep_item_collect", "item"), + ("outer", "index", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("deep_item_collect", "collection", "if", "true_input"), + ("deep_collect", "collection", "if", "false_input"), + ("if", "value", "collect", "item"), + ] + return _graph_from(nodes, edges) + + +def _reopen_of_a_cross_product_graph(swap_add_inputs: bool, condition: bool = False) -> Graph: + """A literal If chooses between an Iterate reopening a Collect over two crossed axes and the crossed inner axis.""" + first_axis, second_axis = ("right", "left") if swap_add_inputs else ("left", "right") + return _graph_from( + [ + CollectionConcatInvocation(id="left_src", first=[1, 2]), + IterateInvocation(id="left"), + CollectionConcatInvocation(id="right_src", first=[10, 20, 30]), + IterateInvocation(id="right"), + AddInvocation(id="add"), + CollectInvocation(id="closing_collect"), + IterateInvocation(id="reopened"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=condition), + CollectInvocation(id="collect"), + ], + [ + ("left_src", "collection", "left", "collection"), + ("right_src", "collection", "right", "collection"), + (first_axis, "item", "add", "a"), + (second_axis, "item", "add", "b"), + ("add", "value", "closing_collect", "item"), + ("closing_collect", "collection", "reopened", "collection"), + ("reopened", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + (second_axis, "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + + +def _reopen_of_a_different_closed_axis_graph() -> Graph: + return _graph_from( + [ + CollectionConcatInvocation(id="closed_src", first=["z1", "z2", "z3"]), + IterateInvocation(id="closed"), + CollectInvocation(id="closing_collect"), + IterateInvocation(id="reopened"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + CollectionConcatInvocation(id="other_src", first=["u1", "u2"]), + IterateInvocation(id="other"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("closed_src", "collection", "closed", "collection"), + ("closed", "item", "closing_collect", "item"), + ("closing_collect", "collection", "reopened", "collection"), + ("reopened", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("other_src", "collection", "other", "collection"), + ("other", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + + +def _for_output_connected_to_selected_axis_graph(selected_is_downstream: bool) -> Graph: + """A For's final output on the unselected branch, under an axis connected to the selected one by a reopen.""" + upstream_ids = ("upstream_src", "upstream") + downstream_ids = ("downstream_src", "downstream") + for_parent = upstream_ids[1] if selected_is_downstream else downstream_ids[1] + selected_axis = downstream_ids[1] if selected_is_downstream else upstream_ids[1] + return _graph_from( + [ + CollectionConcatInvocation(id="upstream_src", first=[["a0", "a1"], ["a2"]]), + IterateInvocation(id="upstream"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + CollectInvocation(id="closing_collect"), + CollectionConcatInvocation(id="downstream_src", second=[["x"]]), + IterateInvocation(id="downstream"), + FixedCollectionTestInvocation(id="for_items"), + ForInvocation(id="for"), + ForReturnInvocation(id="return"), + MarkedAnyTypeTestInvocation(id="selected_value", marker="s:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("upstream_src", "collection", "upstream", "collection"), + ("upstream", "item", "body", "value"), + ("body", "value", "closing_collect", "item"), + ("closing_collect", "collection", "downstream_src", "first"), + ("downstream_src", "collection", "downstream", "collection"), + (for_parent, "item", "for_items", "value"), + ("for_items", "collection", "for", "collection"), + ("for", "item", "return", "output"), + create_loop_linkage("for", "return"), + ("for", "output_collection", "if", "false_input"), + (selected_axis, "item", "selected_value", "value"), + ("selected_value", "value", "if", "true_input"), + ("if", "value", "collect", "item"), + ], + ) + + +@pytest.mark.parametrize( + ("selected_is_downstream", "expected"), + [ + (True, ["s:b:['a0', 'a1']", "s:b:['a2']", "s:['x']"]), + (False, ["s:['a0', 'a1']", "s:['a2']"]), + ], + ids=["selects-reopened-axis", "selects-closed-axis"], +) +def test_literal_if_ignores_unselected_for_output_under_an_axis_connected_to_the_selected_one( + selected_is_downstream: bool, expected: list[str] +) -> None: + """An unselected For final output adds no frames even when its axis feeds, or is fed by, the selected axis.""" + _trace, state = _run_to_completion(_for_output_connected_to_selected_axis_graph(selected_is_downstream)) + + assert _source_values(state, "collect") == {(): expected} + + +@pytest.mark.parametrize("condition_from_branch", [False, True], ids=["literal", "condition-from-branch"]) +def test_if_over_cross_product_carried_by_its_deciding_input_completes(condition_from_branch: bool) -> None: + """Unrelated axes combined by one input are that input's own frames, so the If keeps the cross product.""" + _trace, state = _run_to_completion(_cross_product_branch_graph(True, condition_from_branch)) + + assert _source_values(state, "collect") == {(0,): [11, 21, 31], (1,): [12, 22, 32]} + + +def test_literal_if_selecting_axis_reopened_directly_from_a_collect_pairs_frames() -> None: + """An Iterate fed straight from the Collect that closed the other axis has its length, so frames pair by index.""" + _trace, state = _run_to_completion(_collect_then_reopen_graph(True)) + + assert _source_values(state, "collect") == {(0,): ["t:b:a1"], (1,): ["t:b:a2"]} + + +def test_literal_if_selecting_axis_reopened_through_two_collects_pairs_frames() -> None: + """A chain of direct reopens keeps the first axis's length, so the last reopened axis pairs with it.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="src", first=["a1", "a2", "a3"]), + IterateInvocation(id="first"), + MarkedAnyTypeTestInvocation(id="first_body", marker="b:"), + CollectInvocation(id="first_collect"), + IterateInvocation(id="middle"), + MarkedAnyTypeTestInvocation(id="middle_body", marker="c:"), + CollectInvocation(id="middle_collect"), + IterateInvocation(id="last"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("src", "collection", "first", "collection"), + ("first", "item", "first_body", "value"), + ("first_body", "value", "first_collect", "item"), + ("first_collect", "collection", "middle", "collection"), + ("middle", "item", "middle_body", "value"), + ("middle_body", "value", "middle_collect", "item"), + ("middle_collect", "collection", "last", "collection"), + ("last", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("first", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + _trace, state = _run_to_completion(graph) + + assert _source_values(state, "collect") == {(0,): ["t:c:b:a1"], (1,): ["t:c:b:a2"], (2,): ["t:c:b:a3"]} + + +def test_literal_if_selecting_axis_reopened_from_nested_collects_pairs_with_the_outer_axis() -> None: + """Each Collect in a chain closes one more axis, so the last one gathers one item per outermost frame.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="outer_src", first=[[["a"], ["b"]], [["c", "d"]]]), + IterateInvocation(id="outer"), + CollectionConcatInvocation(id="middle_src"), + IterateInvocation(id="middle"), + CollectionConcatInvocation(id="inner_src"), + IterateInvocation(id="inner"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + CollectInvocation(id="inner_collect"), + CollectInvocation(id="middle_collect"), + CollectInvocation(id="outer_collect"), + IterateInvocation(id="reopened"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("outer_src", "collection", "outer", "collection"), + ("outer", "item", "middle_src", "first"), + ("middle_src", "collection", "middle", "collection"), + ("middle", "item", "inner_src", "first"), + ("inner_src", "collection", "inner", "collection"), + ("inner", "item", "body", "value"), + ("body", "value", "inner_collect", "item"), + ("inner_collect", "collection", "middle_collect", "item"), + ("middle_collect", "collection", "outer_collect", "item"), + ("outer_collect", "collection", "reopened", "collection"), + ("reopened", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("outer", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + _trace, state = _run_to_completion(graph) + + assert _source_values(state, "collect") == { + (0,): ["t:[['b:a'], ['b:b']]"], + (1,): ["t:[['b:c', 'b:d']]"], + } + + +def test_literal_if_selecting_axis_nested_under_a_reopened_axis_pairs_frames() -> None: + """An Iterate nested under a direct reopen runs in the reopened frames, which pair with the closed axis.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="src", first=[["x1", "x2"], ["y1"]]), + IterateInvocation(id="first"), + AnyTypeTestInvocation(id="body"), + CollectInvocation(id="closing_collect"), + IterateInvocation(id="reopened"), + CollectionConcatInvocation(id="nested_src"), + IterateInvocation(id="nested"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("src", "collection", "first", "collection"), + ("first", "item", "body", "value"), + ("body", "value", "closing_collect", "item"), + ("closing_collect", "collection", "reopened", "collection"), + ("reopened", "item", "nested_src", "first"), + ("nested_src", "collection", "nested", "collection"), + ("nested", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("first", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + _trace, state = _run_to_completion(graph) + + assert _source_values(state, "collect") == {(0, 0): ["t:x1"], (0, 1): ["t:x2"], (1, 0): ["t:y1"]} + + +def test_literal_if_selecting_axis_reopened_within_kept_outer_frames_pairs_both_axes() -> None: + """A Collect closing an inner axis keeps the outer one; its reopening pairs with the inner axis per outer frame.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="outer_src", first=[["a1", "a2"], ["b1"]]), + IterateInvocation(id="outer"), + CollectionConcatInvocation(id="inner_src"), + IterateInvocation(id="inner"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + CollectInvocation(id="closing_collect"), + IterateInvocation(id="reopened"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("outer_src", "collection", "outer", "collection"), + ("outer", "item", "inner_src", "first"), + ("inner_src", "collection", "inner", "collection"), + ("inner", "item", "body", "value"), + ("body", "value", "closing_collect", "item"), + ("closing_collect", "collection", "reopened", "collection"), + ("reopened", "item", "true_value", "value"), + ("true_value", "value", "if", "true_input"), + ("inner", "item", "false_value", "value"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + _trace, state = _run_to_completion(graph) + + assert _source_values(state, "collect") == {(0, 0): ["t:b:a1"], (0, 1): ["t:b:a2"], (1, 0): ["t:b:b1"]} + + +@pytest.mark.parametrize( + ("condition", "expected"), + [(False, [["a0", "a1"], ["a2"]]), (True, [["b0"], ["b1"], ["b2"]])], + ids=["selects-a", "selects-b"], +) +def test_literal_if_ignores_independent_axis_reaching_it_through_unselected_for_output( + condition: bool, expected: list[list[str]] +) -> None: + """An unselected For final output adds no frames, so the If runs once per item of the selected branch's axis.""" + nodes: list[Any] = [IfInvocation(id="if", condition=condition), CollectInvocation(id="collect")] + edges: list[Any] = [("if", "value", "collect", "item")] + for prefix, values, branch in ( + ("a", [["a0", "a1"], ["a2"]], "false_input"), + ("b", [["b0"], ["b1"], ["b2"]], "true_input"), + ): + nodes += [ + CollectionConcatInvocation(id=f"{prefix}_src", first=values), + IterateInvocation(id=f"{prefix}_iter"), + ForInvocation(id=f"{prefix}_for"), + ForReturnInvocation(id=f"{prefix}_return"), + ] + edges += [ + (f"{prefix}_src", "collection", f"{prefix}_iter", "collection"), + (f"{prefix}_iter", "item", f"{prefix}_for", "collection"), + (f"{prefix}_for", "item", f"{prefix}_return", "output"), + create_loop_linkage(f"{prefix}_for", f"{prefix}_return"), + (f"{prefix}_for", "output_collection", "if", branch), + ] + _trace, state = _run_to_completion(_graph_from(nodes, edges)) + + assert _source_values(state, "collect") == {(): expected} + + +@pytest.mark.parametrize("condition", [True, False]) +def test_root_condition_if_broadcasts_over_nested_axes_both_branches_share(condition: bool) -> None: + """A root-frame condition selects between branches that run in the same nested frames, once per inner item.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="src", first=[["p", "q"], ["r"]]), + IterateInvocation(id="outer"), + CollectionConcatInvocation(id="inner_src", first=[]), + IterateInvocation(id="inner"), + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + FixedCollectionTestInvocation(id="condition", value=condition), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ], + [ + ("src", "collection", "outer", "collection"), + ("outer", "item", "inner_src", "first"), + ("inner_src", "collection", "inner", "collection"), + ("inner", "item", "true_value", "value"), + ("inner", "item", "false_value", "value"), + ("true_value", "value", "if", "true_input"), + ("false_value", "value", "if", "false_input"), + ("condition", "condition", "if", "condition"), + ("if", "value", "collect", "item"), + ], + ) + + _trace, state = _run_to_completion(graph) + + marker = "t:" if condition else "f:" + assert _source_values(state, "collect") == {(0,): [f"{marker}p", f"{marker}q"], (1,): [f"{marker}r"]} + + +def _shared_nested_axes_graph(data: list[Any], levels: int, condition_on_outer: bool) -> Graph: + """Both branches read the innermost of ``levels`` nested Iterates; the condition is root or per outer item.""" + axis_ids = ["outer", "middle", "inner"] if levels == 3 else ["outer", "inner"] + nodes: list[Any] = [CollectionConcatInvocation(id="src", first=data), IterateInvocation(id=axis_ids[0])] + edges: list[Any] = [("src", "collection", axis_ids[0], "collection")] + for parent_id, axis_id in itertools.pairwise(axis_ids): + nodes += [CollectionConcatInvocation(id=f"{axis_id}_src", first=[]), IterateInvocation(id=axis_id)] + edges += [ + (parent_id, "item", f"{axis_id}_src", "first"), + (f"{axis_id}_src", "collection", axis_id, "collection"), + ] + nodes += [ + MarkedAnyTypeTestInvocation(id="true_value", marker="t:"), + MarkedAnyTypeTestInvocation(id="false_value", marker="f:"), + FixedCollectionTestInvocation(id="condition", value=True), + IfInvocation(id="if"), + ] + edges += [ + (axis_ids[-1], "item", "true_value", "value"), + (axis_ids[-1], "item", "false_value", "value"), + ("true_value", "value", "if", "true_input"), + ("false_value", "value", "if", "false_input"), + ("condition", "condition", "if", "condition"), + ] + if condition_on_outer: + edges.append((axis_ids[0], "item", "condition", "value")) + previous = ("if", "value") + for index in range(levels): + nodes.append(CollectInvocation(id=f"collect{index}")) + edges.append((*previous, f"collect{index}", "item")) + previous = (f"collect{index}", "collection") + return _graph_from(nodes, edges) + + +def test_per_item_condition_if_over_three_shared_nested_axes_completes() -> None: + """A per-outer-item condition broadcasts over two deeper shared axes, and no Collect closes before the If runs.""" + graph = _shared_nested_axes_graph([[["p", "q"], ["r"]], [["s"]]], levels=3, condition_on_outer=True) + + trace, state = _run_to_completion(graph) + + assert _source_values(state, "collect2") == {(): [[["t:p", "t:q"], ["t:r"]], [["t:s"]]]} + assert trace.index("collect0") > max(index for index, source_id in enumerate(trace) if source_id == "if") + + +@pytest.mark.parametrize( + ("data", "levels", "expected"), + [ + pytest.param([[], []], 2, {(): [[], []]}, id="two-levels"), + pytest.param([[[], []], [[]]], 3, {(): [[[], []], [[]]]}, id="three-levels"), + ], +) +def test_root_condition_if_over_shared_nested_axes_with_empty_innermost_iterate_has_no_frames( + data: list[Any], levels: int, expected: dict[tuple[int, ...], Any] +) -> None: + """An innermost shared Iterate empty in every frame leaves a root-conditioned If no frames, as a literal one has.""" + trace, state = _run_to_completion(_shared_nested_axes_graph(data, levels=levels, condition_on_outer=False)) + + assert _source_values(state, f"collect{levels - 1}") == expected + assert "if" not in trace + + +def test_literal_if_selecting_for_output_under_empty_iterate_has_no_frames() -> None: + """An Iterate with no items that a selected branch For runs under leaves the If no executions.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="src", first=[]), + IterateInvocation(id="iterate"), + ForInvocation(id="for"), + MarkedAnyTypeTestInvocation(id="body", marker="b:"), + ForReturnInvocation(id="ret"), + AnyTypeTestInvocation(id="mid"), + AnyTypeTestInvocation(id="fallback", value="fb"), + IfInvocation(id="if", condition=True), + CollectInvocation(id="collect"), + ], + [ + ("src", "collection", "iterate", "collection"), + ("iterate", "item", "for", "collection"), + ("for", "item", "body", "value"), + ("body", "value", "ret", "output"), + create_loop_linkage("for", "ret"), + ("for", "output_collection", "mid", "value"), + ("mid", "value", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ], + ) + + trace, state = _run_to_completion(graph) + + assert _source_values(state, "collect") == {(): []} + assert "if" not in trace and "fallback" not in trace + + +@pytest.mark.parametrize("fail_unselected", [False, True], ids=["succeeds", "unselected-fails"]) +def test_per_item_if_gates_branch_for_together_with_the_for_nested_in_its_body(fail_unselected: bool) -> None: + """A gated For brings the For nested in its body along, so a false item runs neither loop.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="outer_source", first=["", "x"]), + IterateInvocation(id="outer_iterate"), + FixedCollectionTestInvocation(id="items"), + ForInvocation(id="for"), + FixedCollectionTestInvocation(id="mk"), + ForInvocation(id="ifor"), + AnyTypeTestInvocation(id="ibody"), + ForReturnInvocation(id="ireturn"), + ForReturnInvocation(id="return"), + CollectInvocation(id="branch_collect"), + AnyTypeTestInvocation(id="fallback", value="fallback"), + IfInvocation(id="if"), + AnyTypeTestInvocation(id="sink"), + ], + [ + ("outer_source", "collection", "outer_iterate", "collection"), + ("outer_iterate", "item", "items", "value"), + ("items", "condition", "if", "condition"), + ("items", "collection", "for", "collection"), + ("for", "item", "mk", "value"), + ("mk", "collection", "ifor", "collection"), + ("ifor", "item", "ibody", "value"), + ("ibody", "value", "ireturn", "output"), + create_loop_linkage("ifor", "ireturn"), + ("ifor", "output_collection", "return", "output"), + create_loop_linkage("for", "return"), + ("for", "output_collection", "branch_collect", "item"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "sink", "value"), + ], + ) + gated_sources = {"for", "mk", "ifor", "ibody", "ireturn", "return", "branch_collect"} + + state = GraphExecutionState(graph=graph) + trace: list[tuple[str, tuple[int, ...]]] = [] + while (node := state.next()) is not None: + source_id = state.prepared_source_mapping[node.id] + iteration_path = state._get_iteration_path(node.id) + trace.append((source_id, iteration_path)) + if fail_unselected and source_id in gated_sources and iteration_path[:1] == (0,): + state.set_node_error(node.id, f"{source_id} ran for an unselected item") + continue + state.complete(node.id, node.invoke(Mock())) + + assert not state.has_error(), state.errors + assert state.is_complete(), f"execution stalled after {trace}" + assert not [entry for entry in trace if entry[0] in gated_sources and entry[1][:1] == (0,)] + assert _source_values(state, "sink") == {(0,): "fallback", (1,): [[[0, 1], [0, 1]]]} + + +@pytest.mark.parametrize( + ("build_graph", "message"), + [ + pytest.param( + lambda: _unrelated_open_branches_graph(True, False), + "unrelated open iterator axes", + id="unrelated-literal-T", + ), + pytest.param( + lambda: _unrelated_open_branches_graph(False, False), + "unrelated open iterator axes", + id="unrelated-literal-F", + ), + pytest.param( + lambda: _unrelated_open_branches_graph(True, True), + "unrelated open iterator axes", + id="unrelated-computed-T", + ), + pytest.param( + lambda: _unrelated_open_branches_graph(False, True), + "unrelated open iterator axes", + id="unrelated-computed-F", + ), + pytest.param( + lambda: _root_condition_with_deeper_axis_on_one_branch_graph([[], ["p", "q"]]), + "different deeper open iterator axes", + id="root-condition-partly-empty-deeper-axis", + ), + pytest.param( + lambda: _root_condition_with_deeper_axis_on_one_branch_graph([["p"], ["q", "r"]]), + "different deeper open iterator axes", + id="root-condition-deeper-axis", + ), + pytest.param( + _per_item_condition_with_deeper_collect_branch_graph, + "deeper open iterator axes kept by a Collect", + id="per-item-condition-deeper-collect", + ), + pytest.param( + _per_item_condition_with_deeper_for_output_collect_branch_graph, + "deeper open iterator axes kept by a Collect", + id="per-item-condition-deeper-for-output-collect", + ), + pytest.param( + lambda: _cross_product_branch_graph(condition=False), + "unrelated open iterator axes", + id="literal-root-value-over-cross-product", + ), + pytest.param( + _per_item_condition_with_collect_of_collects_branch_graph, + "deeper open iterator axes kept by a Collect", + id="per-item-condition-deeper-collect-beside-collect-of-collects", + ), + pytest.param( + lambda: _collect_then_reopen_graph(condition=False, extra_items=("x",)), + "unrelated open iterator axes", + id="literal-selects-axis-a-collect-closed-before-the-other", + ), + pytest.param( + lambda: _collect_then_reopen_graph(condition=True, extra_items=("x",)), + "unrelated open iterator axes", + id="literal-selects-axis-reopened-through-a-node-after-the-collect", + ), + pytest.param( + lambda: _collect_then_reopen_graph(condition=True, appended_items=("x",)), + "unrelated open iterator axes", + id="literal-selects-axis-reopened-from-a-collect-that-appends-items", + ), + pytest.param( + _collect_with_literal_items_then_reopen_graph, + "unrelated open iterator axes", + id="literal-selects-axis-reopened-from-a-collect-with-literal-items", + ), + pytest.param( + lambda: _collect_then_reopen_graph(condition=False), + "unrelated open iterator axes", + id="literal-selects-axis-an-iterate-reopens", + ), + pytest.param( + lambda: _reopen_of_a_cross_product_graph(swap_add_inputs=False), + "unrelated open iterator axes", + id="literal-selects-closed-axis-of-a-reopened-cross-product", + ), + pytest.param( + lambda: _reopen_of_a_cross_product_graph(swap_add_inputs=True), + "unrelated open iterator axes", + id="literal-selects-closed-axis-of-a-reopened-swapped-cross-product", + ), + pytest.param( + lambda: _reopen_of_a_cross_product_graph(swap_add_inputs=True, condition=True), + "unrelated open iterator axes", + id="literal-selects-axis-reopening-a-cross-product", + ), + pytest.param( + lambda: _collect_of_collect_reopen_graph([["a"], [], ["b", "c"]], "member"), + "unrelated open iterator axes", + id="literal-selects-collect-of-collects-reopening-over-the-inner-axis", + ), + pytest.param( + lambda: _collect_of_collect_reopen_graph([["a"], ["b", "c"]], "group", extra_collect=True), + "unrelated open iterator axes", + id="literal-selects-reopening-of-more-collects-than-axes", + ), + *( + pytest.param( + lambda unselected_axis=unselected_axis: _collect_of_collect_reopen_graph( + [["a"], ["b", "c"]], unselected_axis, extra_collect=True, append_unrelated_axis=True + ), + "unrelated open iterator axes", + id=f"literal-selects-reopening-of-collects-keeping-unrelated-axes-over-{unselected_axis}", + ) + for unselected_axis in ("group", "appended") + ), + pytest.param( + _reopen_of_a_different_closed_axis_graph, + "unrelated open iterator axes", + id="literal-selects-axis-reopening-a-different-closed-axis", + ), + ], +) +def test_if_without_one_nested_chain_of_frames_is_rejected_before_any_node_runs( + build_graph: Callable[[], Graph], message: str +) -> None: + """An If runs once per frame of a single chain of nested open axes; other frame sets fail before running.""" + state = GraphExecutionState(graph=build_graph()) + + with pytest.raises(graph_if_runtime.UnsupportedIfIterationAxesError, match=message): + _run(state) + assert state.executed_history == [] + + +@pytest.mark.parametrize("fail_unselected", [False, True], ids=["succeeds", "unselected-fails"]) +def test_per_item_if_keeps_shared_for_running_but_gates_its_branch_collect(fail_unselected: bool) -> None: + """A For whose final output also feeds a node outside the If runs for every item; its branch Collect does not.""" + graph = _graph_from( + [ + CollectionConcatInvocation(id="outer_source", first=["", "x"]), + IterateInvocation(id="outer_iterate"), + FixedCollectionTestInvocation(id="items"), + ForInvocation(id="for"), + AnyTypeTestInvocation(id="body"), + ForReturnInvocation(id="return"), + CollectInvocation(id="collect"), + CollectInvocation(id="other_collect"), + AnyTypeTestInvocation(id="side_sink"), + AnyTypeTestInvocation(id="fallback", value="fallback"), + IfInvocation(id="if"), + AnyTypeTestInvocation(id="sink"), + ], + [ + ("outer_source", "collection", "outer_iterate", "collection"), + ("outer_iterate", "item", "items", "value"), + ("items", "condition", "if", "condition"), + ("items", "collection", "for", "collection"), + ("for", "item", "body", "value"), + ("body", "value", "return", "output"), + create_loop_linkage("for", "return"), + ("for", "output_collection", "collect", "item"), + ("collect", "collection", "if", "true_input"), + ("for", "output_collection", "other_collect", "item"), + ("other_collect", "collection", "side_sink", "value"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "sink", "value"), + ], + ) + + state = GraphExecutionState(graph=graph) + trace: list[tuple[str, tuple[int, ...]]] = [] + while (node := state.next()) is not None: + source_id = state.prepared_source_mapping[node.id] + iteration_path = state._get_iteration_path(node.id) + trace.append((source_id, iteration_path)) + if fail_unselected and source_id == "collect" and iteration_path[:1] == (0,): + state.set_node_error(node.id, "collect ran for an unselected item") + continue + state.complete(node.id, node.invoke(Mock())) + + assert not state.has_error(), state.errors + assert state.is_complete(), f"execution stalled after {trace}" + assert ("collect", (0,)) not in trace + assert _source_values(state, "sink") == {(0,): "fallback", (1,): [[0, 1]]} + assert _source_values(state, "side_sink") == {(0,): [[0, 1]], (1,): [[0, 1]]} + + +def _shared_outer_iterate_branch_graph(condition: bool) -> Graph: + """The outer frame reaches the If only through the deferred branch, so the If cannot see it yet.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=["x", "y"])) + graph.add_node(IterateInvocation(id="outer_iterate")) + graph.add_node(AnyTypeTestInvocation(id="side")) + graph.add_node(FixedCollectionTestInvocation(id="inner_source")) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("outer_source", "collection", "outer_iterate", "collection")) + graph.add_edge(create_edge("outer_iterate", "item", "side", "value")) + graph.add_edge(create_edge("outer_iterate", "item", "inner_source", "value")) + graph.add_edge(create_edge("inner_source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "body", "value")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +def _sibling_iterates_branch_graph(condition: bool) -> Graph: + """The Collect closes only one of two sibling frames, so the If inherits the other.""" + graph = Graph() + graph.add_node(FixedCollectionTestInvocation(id="left_source")) + graph.add_node(FixedCollectionTestInvocation(id="right_source")) + graph.add_node(IterateInvocation(id="left_iterate")) + graph.add_node(IterateInvocation(id="right_iterate")) + graph.add_node(AddInvocation(id="body")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value=["fallback"])) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_edge(create_edge("left_source", "collection", "left_iterate", "collection")) + graph.add_edge(create_edge("right_source", "collection", "right_iterate", "collection")) + graph.add_edge(create_edge("left_iterate", "item", "body", "a")) + graph.add_edge(create_edge("right_iterate", "item", "body", "b")) + graph.add_edge(create_edge("body", "value", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + return graph + + +@pytest.mark.parametrize( + "build_graph", + [_shared_outer_iterate_branch_graph, _sibling_iterates_branch_graph], + ids=["outer-frame-through-branch", "sibling-frame-left-open"], +) +def test_literal_if_rejecting_collect_branch_that_keeps_a_frame_open_runs_once( + build_graph: Callable[[bool], Graph], +) -> None: + """A rejected Collect branch never runs, so the frame it would leave open gives the If no executions. + + The If does not wait for the Iterate inside that branch and runs once in the root frame. + """ + graph = build_graph(False) + graph.validate_self() + + trace, state = _run(GraphExecutionState(graph=graph)) + + # This shape is outside the generic scheduler's supported set. + assert isinstance(state._execution_scheduler, _ExecutionScheduler) + assert state.is_complete(), f"execution stalled after {trace}" + assert [ + (state._get_iteration_path(exec_id), state.results[exec_id].value) + for exec_id in state._prepared_registry().get_prepared_ids("sink") + ] == [((), ["fallback"])] + assert "body" not in trace + + +@pytest.mark.parametrize( + ("condition", "expected_paths", "expected_values"), + [ + pytest.param(True, {(0,), (1,)}, [0, 1], id="selects-item-branch"), + pytest.param(False, {(0,), (1,)}, ["root-value", "root-value"], id="selects-root-branch"), + ], +) +def test_root_condition_if_keeps_branch_iterate_cardinality( + condition: bool, expected_paths: set[tuple[int, ...]], expected_values: list[int | str] +) -> None: + """A root-frame condition is read by one If execution per item of the Iterate open on a branch.""" + graph = Graph() + graph.add_node(FixedCollectionTestInvocation(id="source", value=True)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(FixedCollectionTestInvocation(id="condition", value=condition)) + graph.add_node(AnyTypeTestInvocation(id="root_branch", value="root-value")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("condition", "condition", "if", "condition")) + graph.add_edge(create_edge("iterate", "item", "if", "true_input")) + graph.add_edge(create_edge("root_branch", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph(GraphExecutionState(graph=graph)) + + # This shape is outside the generic scheduler's supported set. + assert isinstance(state._execution_scheduler, _ExecutionScheduler) + assert state.is_complete(), f"execution stalled after {trace}" + if_paths = {state._get_iteration_path(exec_id) for exec_id in state._prepared_registry().get_prepared_ids("if")} + assert if_paths == expected_paths + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == expected_values + + +def test_literal_if_selecting_empty_iterate_branch_produces_no_values() -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="root_branch", value="root-value")) + graph.add_node(IfInvocation(id="if", condition=True)) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "true_input")) + graph.add_edge(create_edge("root_branch", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph(GraphExecutionState(graph=graph)) + + # This shape is outside the generic scheduler's supported set. + assert isinstance(state._execution_scheduler, _ExecutionScheduler) + assert state.is_complete(), f"execution stalled after {trace}" + assert "root_branch" not in trace + assert not state._prepared_registry().get_prepared_ids("if") + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == [] + + +def test_node_after_literal_if_selecting_root_branch_keeps_iterate_cardinality() -> None: + """A literal If selecting a root value still runs once per item of the Iterate open on its other branch.""" + graph = Graph() + graph.add_node(FixedCollectionTestInvocation(id="source", value=True)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="root_branch", value="root-value")) + graph.add_node(IfInvocation(id="if", condition=False)) + graph.add_node(AnyTypeTestInvocation(id="sink")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "true_input")) + graph.add_edge(create_edge("root_branch", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "sink", "value")) + graph.add_edge(create_edge("sink", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph(GraphExecutionState(graph=graph)) + + # This shape is outside the generic scheduler's supported set. + assert isinstance(state._execution_scheduler, _ExecutionScheduler) + assert state.is_complete(), f"execution stalled after {trace}" + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == ["root-value", "root-value"] + + +def test_computed_condition_selecting_empty_iterate_branch_produces_no_values() -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(FixedCollectionTestInvocation(id="condition", value=True)) + graph.add_node(AnyTypeTestInvocation(id="root_branch", value="root-value")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("condition", "condition", "if", "condition")) + graph.add_edge(create_edge("iterate", "item", "if", "true_input")) + graph.add_edge(create_edge("root_branch", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph(GraphExecutionState(graph=graph)) + + # This shape is outside the generic scheduler's supported set. + assert isinstance(state._execution_scheduler, _ExecutionScheduler) + assert state.is_complete(), f"execution stalled after {trace}" + assert "root_branch" not in trace + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == [] + + +@pytest.mark.parametrize( + ("conditions", "expected_values"), + [ + pytest.param([True, False], ["inner:True", "outer-false"], id="first-nested"), + pytest.param([False, True], ["outer-false", "inner:True"], id="second-nested"), + ], +) +def test_per_item_if_selecting_per_item_nested_if_keeps_every_item( + conditions: list[bool], expected_values: list[str] +) -> None: + """The outer If runs in every item frame even when an item selects a nested If that is not materialized yet.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="conditions", first=conditions)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(FixedCollectionTestInvocation(id="inner_condition")) + graph.add_node(MarkedAnyTypeTestInvocation(id="inner_true", marker="inner:")) + graph.add_node(AnyTypeTestInvocation(id="inner_false", value="inner-false")) + graph.add_node(IfInvocation(id="inner_if")) + graph.add_node(AnyTypeTestInvocation(id="outer_false", value="outer-false")) + graph.add_node(IfInvocation(id="outer_if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("conditions", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "outer_if", "condition")) + graph.add_edge(create_edge("iterate", "item", "inner_condition", "value")) + graph.add_edge(create_edge("inner_condition", "condition", "inner_if", "condition")) + graph.add_edge(create_edge("iterate", "item", "inner_true", "value")) + graph.add_edge(create_edge("inner_true", "value", "inner_if", "true_input")) + graph.add_edge(create_edge("inner_false", "value", "inner_if", "false_input")) + graph.add_edge(create_edge("inner_if", "value", "outer_if", "true_input")) + graph.add_edge(create_edge("outer_false", "value", "outer_if", "false_input")) + graph.add_edge(create_edge("outer_if", "value", "collect", "item")) + graph.validate_self() + + trace, state = _run_graph(GraphExecutionState(graph=graph)) + + # This shape is outside the generic scheduler's supported set. + assert isinstance(state._execution_scheduler, _ExecutionScheduler) + assert state.is_complete(), f"execution stalled after {trace}" + collect_id = next(iter(state._prepared_registry().get_prepared_ids("collect"))) + assert state.results[collect_id].collection == expected_values + + @pytest.mark.parametrize("force_compatibility_scheduler", [False, True]) def test_if_collect_collection_inner_iterate_keeps_values_in_outer_frame( force_compatibility_scheduler: bool, @@ -9726,6 +11241,33 @@ def test_literal_if_branch_feeds_per_item_if_without_stranding_unselected_root( assert "per_item_unselected" not in trace +def test_per_item_if_condition_from_iterate_index_selects_literal_values() -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[10, 20])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="false_value", value="no")) + graph.add_node(AnyTypeTestInvocation(id="true_value", value="yes")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="output")) + for source, source_field, destination, destination_field in ( + ("source", "collection", "iterate", "collection"), + ("iterate", "index", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("false_value", "value", "if", "false_input"), + ("true_value", "value", "if", "true_input"), + ("if", "value", "output", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + trace, state = _run_graph(GraphExecutionState(graph=graph)) + + assert state.is_complete(), f"graph stalled after {trace}" + output_exec_id = next(iter(state.source_prepared_mapping["output"])) + assert state.results[output_exec_id].collection == ["no", "yes"] + assert trace.count("if") == 2 + + @pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) def test_deep_branch_root_false_with_two_outer_frames_and_nested_collect( force_compatibility_scheduler: bool,