Revert "Revert "Updated SchemaNodeIdentifier namespace handling.""

[yangtools.git] / yang / yang-parser-impl / src / main / java / org / opendaylight / yangtools / yang / parser / stmt / rfc6020 / AugmentStatementImpl.java

diff --git a/yang/yang-parser-impl/src/main/java/org/opendaylight/yangtools/yang/parser/stmt/rfc6020/AugmentStatementImpl.java b/yang/yang-parser-impl/src/main/java/org/opendaylight/yangtools/yang/parser/stmt/rfc6020/AugmentStatementImpl.java
index 12e1a06bb184d16c9b2d1fb7c984b40ab76142ae..4d660e91b2e02a4a57768bb7b90964e02de5181e 100644 (file)
--- a/yang/yang-parser-impl/src/main/java/org/opendaylight/yangtools/yang/parser/stmt/rfc6020/AugmentStatementImpl.java
+++ b/yang/yang-parser-impl/src/main/java/org/opendaylight/yangtools/yang/parser/stmt/rfc6020/AugmentStatementImpl.java
@@ -20,8 +20,10 @@ import org.opendaylight.yangtools.yang.parser.spi.meta.AbstractDeclaredStatement
 import org.opendaylight.yangtools.yang.parser.spi.meta.AbstractStatementSupport;
 import org.opendaylight.yangtools.yang.parser.spi.meta.InferenceException;
 import org.opendaylight.yangtools.yang.parser.spi.meta.ModelActionBuilder;
+import org.opendaylight.yangtools.yang.parser.spi.meta.ModelActionBuilder.Prerequisite;
 import org.opendaylight.yangtools.yang.parser.spi.meta.ModelProcessingPhase;
 import org.opendaylight.yangtools.yang.parser.spi.meta.StmtContext;
+import org.opendaylight.yangtools.yang.parser.spi.meta.StmtContext.Mutable;
 import org.opendaylight.yangtools.yang.parser.spi.meta.StmtContextUtils;
 import org.opendaylight.yangtools.yang.parser.spi.source.SourceException;
 import org.opendaylight.yangtools.yang.parser.spi.source.StmtOrderingNamespace;
@@ -77,32 +79,21 @@ public class AugmentStatementImpl extends AbstractDeclaredStatement<SchemaNodeId
             }
 
             final ModelActionBuilder augmentAction = augmentNode
-                    .newInferenceAction(ModelProcessingPhase.FULL_DECLARATION);
+                    .newInferenceAction(ModelProcessingPhase.EFFECTIVE_MODEL);
             final ModelActionBuilder.Prerequisite<StmtContext<SchemaNodeIdentifier, AugmentStatement, EffectiveStatement<SchemaNodeIdentifier, AugmentStatement>>> sourceCtxPrereq = augmentAction
-                    .requiresCtx(augmentNode,
-                            ModelProcessingPhase.FULL_DECLARATION);
-
+                    .requiresCtx(augmentNode, ModelProcessingPhase.EFFECTIVE_MODEL);
+            final Prerequisite<Mutable<?, ?, EffectiveStatement<?, ?>>> target = augmentAction.mutatesEffectiveCtx(getSearchRoot(augmentNode), SchemaNodeIdentifierBuildNamespace.class, augmentNode.getStatementArgument());
             augmentAction.apply(new ModelActionBuilder.InferenceAction() {
 
                 @Override
                 public void apply() throws InferenceException {
-                    final StatementContextBase<?, ?, ?> augmentTargetCtx = AugmentUtils
-                            .getAugmentTargetCtx(augmentNode);
-
-                    if (augmentTargetCtx == null) {
-                        throw new InferenceException(
-                                "Augment target not found: "
-                                        + augmentNode.getStatementArgument(),
-                                augmentNode.getStatementSourceReference());
-                    }
+                    final StatementContextBase<?, ?, ?> augmentTargetCtx = (StatementContextBase<?, ?, ?>) target.get();
 
                     if (!AugmentUtils.isSupportedAugmentTarget(augmentTargetCtx) || StmtContextUtils.isInExtensionBody(augmentTargetCtx)) {
                         augmentNode.setIsSupportedToBuildEffective(false);
                         return;
                     }
-
                     final StatementContextBase<?, ?, ?> augmentSourceCtx = (StatementContextBase<?, ?, ?>) augmentNode;
-
                     try {
                         AugmentUtils.copyFromSourceToTarget(augmentSourceCtx,
                                 augmentTargetCtx);
@@ -140,6 +131,15 @@ public class AugmentStatementImpl extends AbstractDeclaredStatement<SchemaNodeId
                 }
             });
         }
+
+        private Mutable<?, ?, ?> getSearchRoot(Mutable<?, ?, ?> augmentContext) {
+            Mutable<?, ?, ?> parent = augmentContext.getParentContext();
+            // Augment is in uses - we need to augment instantiated nodes in parent.
+            if(Rfc6020Mapping.USES.equals(parent.getPublicDefinition())) {
+                return parent.getParentContext();
+            }
+            return parent;
+        }
     }
 
     @Nonnull