+ static final class ChoiceBuilder extends CompositeRuntimeTypeBuilder<ChoiceEffectiveStatement, ChoiceRuntimeType> {
+ private ImmutableList<CaseRuntimeType> augmentedCases;
+
+ ChoiceBuilder(final ChoiceEffectiveStatement statement) {
+ super(statement);
+ }
+
+ @Override
+ CompositeRuntimeTypeBuilder<ChoiceEffectiveStatement, ChoiceRuntimeType> fillTypes(
+ final ChildLookup lookup,
+ final AbstractCompositeGenerator<ChoiceEffectiveStatement, ChoiceRuntimeType> generator) {
+ fillAugmentedCases(lookup, generator.augments());
+ return super.fillTypes(lookup, generator);
+ }
+
+ @Override
+ boolean isAugmentedChild(final ChildLookup lookup, final QName qname) {
+ for (var augmented : augmentedCases) {
+ if (qname.equals(augmented.statement().argument())) {
+ return true;
+ }
+ }
+ return super.isAugmentedChild(lookup, qname);
+ }
+
+ @Override
+ ChoiceRuntimeType build(final GeneratedType type, final ChoiceEffectiveStatement statement,
+ final List<RuntimeType> children, final List<AugmentRuntimeType> augments) {
+ verify(augments.isEmpty(), "Unexpected augments %s", augments);
+ children.addAll(augmentedCases);
+ return new DefaultChoiceRuntimeType(type, statement, children);
+ }
+
+ private void fillAugmentedCases(final ChildLookup lookup, final List<AbstractAugmentGenerator> augments) {
+ final var builder = ImmutableList.<CaseRuntimeType>builder();
+ for (var augment : augments) {
+ builder.addAll(augment.augmentedCasesIn(lookup, statement()));
+ }
+ augmentedCases = builder.build();
+ }
+ }
+