def collapse_branches()

in lucid/misc/graph_analysis/parse_overlay.py [0:0]


def collapse_branches(overlay):
  """Detect and collapse brances of nodes in an overlay."""
  structure_map = {}

  for node in overlay.nodes:
    if len(node.inputs) <= 1: continue
    gcd = node.gcd
    if all(inp == gcd or inp.inputs == set([gcd]) for inp in node.inputs):
      branches = [inp if inp != gcd else None
                  for inp in overlay.sorted(node.inputs)]
      structure_map[node] = OverlayStructure("HeadBranch", {"branches" : branches, "head": node})

  for node in overlay.nodes:
    if len(node.consumers) <= 1: continue
    if all(len(out.consumers) == 0 for out in node.consumers):
      branches = overlay.sorted(node.consumers)
      max_node = overlay.sorted(branches)[-1]
      structure_map[max_node] = OverlayStructure("TailBranch", {"branches" : branches, "tail": node})

  return overlay.collapse_structures(structure_map)