Repository: dask/dask
File: dask/order.py (Python)

Code before:
                key = candidates.pop()
                if key in runnable_hull or key in result:
                    continue
                if key in leaf_nodes:
                    add_to_result(key)
                    continue
                path = [key]
                branches = deque([path])
                while branches:
                    path = branches.popleft()
                    while True:
                        current = path[-1]
                        runnable_hull.add(current)
                        deps_downstream = dependents[current]
                        deps_upstream = dependencies[current]  # type: ignore
                        if not deps_downstream:
                            # FIXME: The fact that it is possible for
                            # num_needed[current] == 0 means we're doing some
                            # work twice
                            if num_needed[current] <= 1:
                                for k in path:
                                    add_to_result(k)
                            else:
                                runnable_hull.discard(current)
                        elif len(path) == 1 or len(deps_upstream) == 1:
                            if len(deps_downstream) > 1:
                                for d in sorted(deps_downstream, key=sort_key):
                                    # This ensures we're only considering splitters
                                    # that are genuinely splitting and not
                                    # interleaving
                                    if len(dependencies[d]) == 1:  # type: ignore
                                        branch = path.copy()
                                        branch.append(d)
                                        branches.append(branch)
                                break
                            runnable_hull.update(deps_downstream)
                            path.extend(deps_downstream)
                            continue
                        elif current in known_runnable_paths:
                            known_runnable_paths[current].append(path)
                            if (
                                len(known_runnable_paths[current])
                                >= num_needed[current]
                            ):
                                pruned_branches: deque[list[Key]] = deque()
                                for path in known_runnable_paths_pop(current):
                                    if path[-2] not in result:
                                        pruned_branches.append(path)
                                if len(pruned_branches) < num_needed[current]:
                                    known_runnable_paths[current] = list(
                                        pruned_branches

Review comment (suggestion):
the loop invariant this should preserve is that runnable_hull and known_runnable_paths are disjoint

Code after:
                key = candidates.pop()
                if key in runnable_hull or key in result:
                    continue
                if key in leaf_nodes:
                    add_to_result(key)
                    continue
                path = [key]
                branches = deque([path])
                while branches:
                    path = branches.popleft()
                    while True:
                        # Loop invariant. Too expensive to compute at runtime
                        # assert not set(known_runnable_paths).intersection(runnable_hull)
                        current = path[-1]
                        runnable_hull.add(current)
                        deps_downstream = dependents[current]
                        deps_upstream = dependencies[current]  # type: ignore
                        if not deps_downstream:
                            # FIXME: The fact that it is possible for
                            # num_needed[current] == 0 means we're doing some
                            # work twice
                            if num_needed[current] <= 1:
                                for k in path:
                                    add_to_result(k)
                            else:
                                runnable_hull.discard(current)
                        elif len(path) == 1 or len(deps_upstream) == 1:
                            if len(deps_downstream) > 1:
                                for d in sorted(deps_downstream, key=sort_key):
                                    # This ensures we're only considering splitters
                                    # that are genuinely splitting and not
                                    # interleaving
                                    if len(dependencies[d]) == 1:  # type: ignore
                                        branch = path.copy()
                                        branch.append(d)
                                        branches.append(branch)
                                break
                            path.extend(deps_downstream)
                            continue
                        elif current in known_runnable_paths:
                            known_runnable_paths[current].append(path)
                            runnable_hull.discard(current)
                            if (
                                len(known_runnable_paths[current])
                                >= num_needed[current]
                            ):
                                pruned_branches: deque[list[Key]] = deque()
                                for path in known_runnable_paths_pop(current):
                                    if path[-2] not in result:
                                        pruned_branches.append(path)
                                if len(pruned_branches) < num_needed[current]: