def _substitute(rhs_a, rhs_b, nt=qcfg_rule.NT_1):
  """Replace nt in rhs_a with rhs_b, re-indexing non-terminals if needed."""
  output = []
  for token in rhs_a:
    if token == nt and nt == qcfg_rule.NT_2:
      # Our goal is to replace NT_2 in rhs_a with rhs_b, but we need to
      # do some re-indexing to avoid collisions.
      # First, we re-index NT_1 in rhs_b to NT_2.
      # Based on the logic in `apply`, if rhs_a has arity 2, then rhs_b
      # will have arity < 2, i.e. will not contain NT_2.
      rhs_b = rule_utils.rhs_replace(rhs_b, [qcfg_rule.NT_1], qcfg_rule.NT_2)
      # We can now safely replace NT_2 in rhs_a with rhs_b, which should
      # contain only NT_2.
      output.extend(rhs_b)
    elif token == nt:
      # Replace NT_1 in rhs_a with rhs_b.
      # Based on the logic in `apply`, no collisions on non-terminal indexes
      # should occur, since we should either be in the case:
      # 1. rhs_a only has NT_1, and rhs_b has NT_1 and NT_2
      # 2. rhs_a has NT_1 and NT_2, but rhs_b only has NT_1
      output.extend(rhs_b)
    else:
      output.append(token)
  return output
Replace nt in rhs_a with rhs_b, re-indexing non-terminals if needed.
language/compgen/nqg/model/induction/derivation_utils.py