Add DataSchemaContextNode/SchemaInferenceStack integration

[yangtools.git] / data / yang-data-util / src / main / java / org / opendaylight / yangtools / yang / data / util / ChoiceNodeContextNode.java

diff --git a/data/yang-data-util/src/main/java/org/opendaylight/yangtools/yang/data/util/ChoiceNodeContextNode.java b/data/yang-data-util/src/main/java/org/opendaylight/yangtools/yang/data/util/ChoiceNodeContextNode.java
index a93d1f1883b8b3900dea5fe6351b434687ae3517..c2b3d694ea43fea7aad8d5e6f0cd81b90dc8209b 100644 (file)
--- a/data/yang-data-util/src/main/java/org/opendaylight/yangtools/yang/data/util/ChoiceNodeContextNode.java
+++ b/data/yang-data-util/src/main/java/org/opendaylight/yangtools/yang/data/util/ChoiceNodeContextNode.java
@@ -7,21 +7,28 @@
  */
 package org.opendaylight.yangtools.yang.data.util;
 
+import static com.google.common.base.Verify.verifyNotNull;
+
 import com.google.common.collect.ImmutableMap;
 import java.util.Set;
+import org.eclipse.jdt.annotation.NonNull;
+import org.eclipse.jdt.annotation.Nullable;
 import org.opendaylight.yangtools.yang.common.QName;
 import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier.NodeIdentifier;
 import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier.PathArgument;
 import org.opendaylight.yangtools.yang.model.api.CaseSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.ChoiceSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.DataSchemaNode;
+import org.opendaylight.yangtools.yang.model.util.SchemaInferenceStack;
 
 final class ChoiceNodeContextNode extends AbstractMixinContextNode<NodeIdentifier> {
-    private final ImmutableMap<QName, DataSchemaContextNode<?>> byQName;
     private final ImmutableMap<PathArgument, DataSchemaContextNode<?>> byArg;
+    private final ImmutableMap<QName, DataSchemaContextNode<?>> byQName;
+    private final ImmutableMap<DataSchemaContextNode<?>, QName> childToCase;
 
     ChoiceNodeContextNode(final ChoiceSchemaNode schema) {
         super(NodeIdentifier.create(schema.getQName()), schema);
+        ImmutableMap.Builder<DataSchemaContextNode<?>, QName> childToCaseBuilder = ImmutableMap.builder();
         ImmutableMap.Builder<QName, DataSchemaContextNode<?>> byQNameBuilder = ImmutableMap.builder();
         ImmutableMap.Builder<PathArgument, DataSchemaContextNode<?>> byArgBuilder = ImmutableMap.builder();
 
@@ -29,11 +36,14 @@ final class ChoiceNodeContextNode extends AbstractMixinContextNode<NodeIdentifie
             for (DataSchemaNode cazeChild : caze.getChildNodes()) {
                 DataSchemaContextNode<?> childOp = DataSchemaContextNode.of(cazeChild);
                 byArgBuilder.put(childOp.getIdentifier(), childOp);
+                childToCaseBuilder.put(childOp, caze.getQName());
                 for (QName qname : childOp.getQNameIdentifiers()) {
                     byQNameBuilder.put(qname, childOp);
                 }
             }
         }
+
+        childToCase = childToCaseBuilder.build();
         byQName = byQNameBuilder.build();
         byArg = byArgBuilder.build();
     }
@@ -52,4 +62,29 @@ final class ChoiceNodeContextNode extends AbstractMixinContextNode<NodeIdentifie
     protected Set<QName> getQNameIdentifiers() {
         return byQName.keySet();
     }
-}
+
+    @Override
+    protected DataSchemaContextNode<?> enterChild(final QName child, final SchemaInferenceStack stack) {
+        return pushToStack(getChild(child), stack);
+    }
+
+    @Override
+    protected DataSchemaContextNode<?> enterChild(final PathArgument child, final SchemaInferenceStack stack) {
+        return pushToStack(getChild(child), stack);
+    }
+
+    @Override
+    protected void pushToStack(final @NonNull SchemaInferenceStack stack) {
+        stack.enterChoice(getIdentifier().getNodeType());
+    }
+
+    private @Nullable DataSchemaContextNode<?> pushToStack(final @Nullable DataSchemaContextNode<?> child,
+            final @NonNull SchemaInferenceStack stack) {
+        if (child != null) {
+            final var caseName = verifyNotNull(childToCase.get(child), "No case statement for %s in %s", child, this);
+            stack.enterSchemaTree(caseName);
+            child.pushToStack(stack);
+        }
+        return child;
+    }
+}
\ No newline at end of file