def _apply(rule_a, rule_b):
  """Applies rule_b to rule_a, returning set of derived rules."""
  outputs = []

  if rule_a.arity == 2:
    new_arity = 1 + rule_b.arity
    if new_arity <= 2:
      # Cannot apply an arity 2 rule to an arity 2 rule because this would lead
      # to a rule with 3 different non-terminal indexes, which is disallowed
      # by our QCFG conventions.
      source_0 = _substitute(rule_a.source, rule_b.source)
      target_0 = _substitute(rule_a.target, rule_b.target)
      outputs.append((source_0, target_0, new_arity))
      # Rule can potentially be applied to either non-terminal in rule_a.
      source_1 = _substitute(rule_a.source, rule_b.source, nt=qcfg_rule.NT_2)
      target_1 = _substitute(rule_a.target, rule_b.target, nt=qcfg_rule.NT_2)
      outputs.append((source_1, target_1, new_arity))

  elif rule_a.arity == 1:
    new_arity = rule_b.arity
    source = _substitute(rule_a.source, rule_b.source)
    target = _substitute(rule_a.target, rule_b.target)
    outputs.append((source, target, new_arity))

  output_rules = set()
  for source, target, arity in outputs:
    source, target = rule_utils.canonicalize_nts(source, target, arity)
    output_rules.add(qcfg_rule.QCFGRule(tuple(source), tuple(target), arity))

  return output_rules
Applies rule_b to rule_a, returning set of derived rules.
language/compgen/nqg/model/induction/derivation_utils.py