def compute_nth_child_edges(
    graph, max_child_count
):
  """Computes CHILD_INDEX_* edges from a graph.

  We assume that the graph was generated by `generic_ast_graphs.ast_to_graph`
  and thus that sequence edges are represented with "{field}_out_first/last" and
  sequence helper items. Note that the produced edges connect the parent node
  to each of the child helper nodes, since those exist between the parent node
  and the child AST nodes. This is done so that the edges for the field name and
  the index align with each other (i.e. connect the same pair of nodes), to make
  it easier to learn useful embeddings.

  Args:
    graph: Automaton graph to use.
    max_child_count: Maximum number of children that get explicit nth-child
      edges.

  Returns:
    List of CHILD_INDEX_* edges.
  """
  result = []
  for source_node_id, node_info in graph.items():
    for out_type, destinations in node_info.out_edges.items():
      if out_type.endswith("_out_first"):
        # This is a sequence node, so we should add nth-child edges.
        in_tagged_node, = destinations
        i = 0
        while (i < max_child_count and
               not in_tagged_node.in_edge.endswith("_missing")):
          result.append(
              (source_node_id, in_tagged_node.node_id, f"CHILD_INDEX_{i}"))
          # Follow the chain of next pointers.
          in_tagged_node, = graph[in_tagged_node.node_id].out_edges["next_out"]
          i += 1

  return result