def _convert_to_qcfg(nested_rule):
  """Convert nested JointRule to QCFG source and target."""
  sources = []
  targets = []
  rule = nested_rule[0]
  idx_to_source = {}
  idx_to_target = {}
  for nt_idx, child_rule in enumerate(nested_rule[1:]):
    source, target = _convert_to_qcfg(child_rule)
    idx_to_source[nt_idx + 1] = source
    idx_to_target[nt_idx + 1] = target
  for symbol in rule.source:
    if qcfg_rule.is_nt_fast(symbol):
      index = qcfg_rule.get_nt_index(symbol)
      sources.extend(idx_to_source[index])
    else:
      sources.append(symbol)
  for symbol in rule.target:
    if qcfg_rule.is_nt_fast(symbol):
      index = qcfg_rule.get_nt_index(symbol)
      targets.extend(idx_to_target[index])
    else:
      targets.append(symbol)
  return sources, targets
Convert nested JointRule to QCFG source and target.
language/compgen/csl/augment/qcfg_sampler.py