Improve KeyStatement implementations

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

diff --git a/yang/yang-parser-rfc7950/src/main/java/org/opendaylight/yangtools/yang/parser/rfc7950/stmt/key/AbstractKeyEffectiveStatement.java b/yang/yang-parser-rfc7950/src/main/java/org/opendaylight/yangtools/yang/parser/rfc7950/stmt/key/AbstractKeyEffectiveStatement.java
index 5b292581887c8c56cf41676373e1d7039d6a38a4..3fb868353740d149739b23ec8dd5bf498e986554 100644 (file)
--- a/yang/yang-parser-rfc7950/src/main/java/org/opendaylight/yangtools/yang/parser/rfc7950/stmt/key/AbstractKeyEffectiveStatement.java
+++ b/yang/yang-parser-rfc7950/src/main/java/org/opendaylight/yangtools/yang/parser/rfc7950/stmt/key/AbstractKeyEffectiveStatement.java
@@ -7,9 +7,6 @@
  */
 package org.opendaylight.yangtools.yang.parser.rfc7950.stmt.key;
 
-import static com.google.common.base.Verify.verify;
-
-import com.google.common.collect.ImmutableSet;
 import java.util.Collection;
 import org.opendaylight.yangtools.yang.model.api.stmt.KeyEffectiveStatement;
 import org.opendaylight.yangtools.yang.model.api.stmt.KeyStatement;
@@ -25,17 +22,12 @@ abstract class AbstractKeyEffectiveStatement
 
         Foreign(final KeyStatement declared, final Collection<SchemaNodeIdentifier> argument) {
             super(declared);
-            this.argument = argument.size() == 1 ? argument.iterator().next() : argument;
+            this.argument = KeyStatementSupport.maskCollection(argument);
         }
 
         @Override
-        @SuppressWarnings("unchecked")
         public final Collection<SchemaNodeIdentifier> argument() {
-            if (argument instanceof Collection) {
-                return (Collection<SchemaNodeIdentifier>) argument;
-            }
-            verify(argument instanceof SchemaNodeIdentifier, "Unexpected argument %s", argument);
-            return ImmutableSet.of((SchemaNodeIdentifier) argument);
+            return KeyStatementSupport.unmaskCollection(argument);
         }
     }