Retain DeclarationReference in DeclaredStatements

[yangtools.git] / yang / yang-parser-rfc7950 / src / main / java / org / opendaylight / yangtools / yang / parser / rfc7950 / stmt / augment / AbstractAugmentStatementSupport.java

diff --git a/yang/yang-parser-rfc7950/src/main/java/org/opendaylight/yangtools/yang/parser/rfc7950/stmt/augment/AbstractAugmentStatementSupport.java b/yang/yang-parser-rfc7950/src/main/java/org/opendaylight/yangtools/yang/parser/rfc7950/stmt/augment/AbstractAugmentStatementSupport.java
index 797cd04bbc64b28674b428c4b821c0a4207ea9b0..7cb61b4f06ad7456346a9ce14b01e9717ec1a088 100644 (file)
--- a/yang/yang-parser-rfc7950/src/main/java/org/opendaylight/yangtools/yang/parser/rfc7950/stmt/augment/AbstractAugmentStatementSupport.java
+++ b/yang/yang-parser-rfc7950/src/main/java/org/opendaylight/yangtools/yang/parser/rfc7950/stmt/augment/AbstractAugmentStatementSupport.java
@@ -18,6 +18,7 @@ import org.opendaylight.yangtools.yang.common.Empty;
 import org.opendaylight.yangtools.yang.model.api.AugmentationSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.Status;
 import org.opendaylight.yangtools.yang.model.api.YangStmtMapping;
+import org.opendaylight.yangtools.yang.model.api.meta.DeclarationReference;
 import org.opendaylight.yangtools.yang.model.api.meta.DeclaredStatement;
 import org.opendaylight.yangtools.yang.model.api.meta.EffectiveStatement;
 import org.opendaylight.yangtools.yang.model.api.meta.StatementDefinition;
@@ -28,7 +29,10 @@ import org.opendaylight.yangtools.yang.model.api.stmt.SchemaNodeIdentifier.Absol
 import org.opendaylight.yangtools.yang.model.api.stmt.SchemaNodeIdentifier.Descendant;
 import org.opendaylight.yangtools.yang.model.api.stmt.StatusEffectiveStatement;
 import org.opendaylight.yangtools.yang.model.api.stmt.WhenEffectiveStatement;
+import org.opendaylight.yangtools.yang.model.parser.api.YangParserConfiguration;
+import org.opendaylight.yangtools.yang.model.ri.stmt.DeclaredStatementDecorators;
 import org.opendaylight.yangtools.yang.model.ri.stmt.DeclaredStatements;
+import org.opendaylight.yangtools.yang.model.ri.stmt.EffectiveStatements;
 import org.opendaylight.yangtools.yang.model.spi.meta.EffectiveStatementMixins.EffectiveStatementWithFlags.FlagsBuilder;
 import org.opendaylight.yangtools.yang.model.spi.meta.SubstatementIndexingException;
 import org.opendaylight.yangtools.yang.parser.rfc7950.stmt.ArgumentUtils;
@@ -49,13 +53,13 @@ abstract class AbstractAugmentStatementSupport
     private static final Pattern PATH_REL_PATTERN1 = Pattern.compile("\\.\\.?\\s*/(.+)");
     private static final Pattern PATH_REL_PATTERN2 = Pattern.compile("//.*");
 
-    AbstractAugmentStatementSupport() {
+    AbstractAugmentStatementSupport(final YangParserConfiguration config) {
         super(YangStmtMapping.AUGMENT, StatementPolicy.copyDeclared(
             (copy, current, substatements) ->
                 copy.getArgument().equals(current.getArgument())
                 && copy.moduleName().getModule().equals(current.moduleName().getModule())
                 && Objects.equals(copy.original(), current.original())
-            ));
+            ), config);
     }
 
     @Override
@@ -112,9 +116,9 @@ abstract class AbstractAugmentStatementSupport
     }
 
     @Override
-    protected final AugmentStatement createEmptyDeclared(
-            final StmtContext<SchemaNodeIdentifier, AugmentStatement, ?> ctx) {
-        return DeclaredStatements.createAugment(ctx.getRawArgument(), ctx.getArgument());
+    protected final AugmentStatement attachDeclarationReference(final AugmentStatement stmt,
+            final DeclarationReference reference) {
+        return DeclaredStatementDecorators.decorateAugment(stmt, reference);
     }
 
     @Override
@@ -140,8 +144,8 @@ abstract class AbstractAugmentStatementSupport
                 .toFlags();
 
         try {
-            return new AugmentEffectiveStatementImpl(stmt.declared(), stmt.getArgument(), flags,
-                stmt.moduleName().getModule(), substatements, (AugmentationSchemaNode) stmt.original());
+            return EffectiveStatements.createAugment(stmt.declared(), stmt.getArgument(), flags,
+                stmt.moduleName().getModule(), substatements, stmt.original(AugmentationSchemaNode.class));
         } catch (SubstatementIndexingException e) {
             throw new SourceException(e.getMessage(), stmt, e);
         }