def TopologicalSort(nodes, node_filter=None):
  """Topological sort.

  We use a BFS-like walk which ensures that sibling are always grouped
  together in the output. This is more convenient for some later analyses.

  Args:
    nodes: [Node, ...] Nodes to sort.
    node_filter: a filter Node->boolean to restrict the graph. A node passes the
      filter on a return value of True. Only the subgraph reachable from a root
      passing the filter is considered.

  Returns:
    A list of Nodes in topological order. Note that node indicies are
    unchanged; the original list nodes is not modified.
  """
  if node_filter is None:
    node_filter = lambda _: True
  sorted_nodes = []
  sources = []
  remaining_in_edges = {}
  for n in nodes:
    if n.Predecessors():
      remaining_in_edges[n] = len(n.Predecessors())
    elif node_filter(n):
      sources.append(n)
  while sources:
    n = sources.pop(0)
    assert node_filter(n)
    sorted_nodes.append(n)
    # We sort by index to get consistent sorts across runs/machines.
    for c in n.SortedSuccessors():
      assert remaining_in_edges[c] > 0
      if not node_filter(c):
        continue
      remaining_in_edges[c] -= 1
      if not remaining_in_edges[c]:
        sources.append(c)
  return sorted_nodes