def expand_unit_rules(rules):
  """Removes unit rules, i.e. X -> Y where X and Y are non-terminals."""
  # List of 2-tuple of non-terminal strings.
  unit_rules = set()
  # List of Rules.
  other_rules = []

  for rule in rules:
    tokens = rule.rhs.split(" ")
    if len(tokens) == 1 and tokens[0].startswith(NON_TERMINAL_PREFIX):
      unit_rhs = tokens[0][len(NON_TERMINAL_PREFIX):]
      unit_rules.add((rule.lhs, unit_rhs))
    else:
      other_rules.append(rule)

  # Identify any chains of unit rules.
  # For example, if we have X -> Y and Y -> Z, then we need to consider X -> Z.
  unit_lhs_to_rhs_set = collections.defaultdict(set)
  for unit_lhs, unit_rhs in unit_rules:
    unit_lhs_to_rhs_set[unit_lhs].add(unit_rhs)

  derived_unit_rules = set()
  for unit_lhs_start in unit_lhs_to_rhs_set.keys():
    stack = [unit_lhs_start]
    visited_lhs = []

    while stack:
      unit_lhs = stack.pop()
      # Check for cycles so that we don't loop indefinitely.
      if unit_lhs in visited_lhs:
        raise ValueError("There exists an unallowed cycle in unit rules: %s" %
                         visited_lhs)
      visited_lhs.append(unit_lhs)
      for unit_rhs in unit_lhs_to_rhs_set.get(unit_lhs, {}):
        stack.append(unit_rhs)
      # Add derived unit rule following transitive chain.
      derived_unit_rules.add((unit_lhs_start, unit_rhs))
  unit_rules |= derived_unit_rules

  # Add derived rules based on unit rules.
  # For example, if we have X -> Y and Y -> foo, then we add X -> foo.
  new_rules = []
  for unit_lhs, unit_rhs in unit_rules:
    for rule in other_rules:
      if rule.lhs == unit_rhs:
        new_rules.append(Rule(unit_lhs, rule.rhs))

  return other_rules + new_rules
Removes unit rules, i.e. X -> Y where X and Y are non-terminals.
language/compgen/nqg/tasks/spider/sql_parser.py