Binding codec v2 - fix augmentation #3

[mdsal.git] / binding2 / mdsal-binding2-generator-api / src / main / java / org / opendaylight / mdsal / binding / javav2 / generator / context / ModuleContext.java

diff --git a/binding2/mdsal-binding2-generator-api/src/main/java/org/opendaylight/mdsal/binding/javav2/generator/context/ModuleContext.java b/binding2/mdsal-binding2-generator-api/src/main/java/org/opendaylight/mdsal/binding/javav2/generator/context/ModuleContext.java
index 34f0114cc8c542042d89a7c5381b9e0503fb5c1a..2f3cef5cc25ad088353f40001a21faef6483fbc4 100644 (file)
--- a/binding2/mdsal-binding2-generator-api/src/main/java/org/opendaylight/mdsal/binding/javav2/generator/context/ModuleContext.java
+++ b/binding2/mdsal-binding2-generator-api/src/main/java/org/opendaylight/mdsal/binding/javav2/generator/context/ModuleContext.java
@@ -52,7 +52,7 @@ public final class ModuleContext {
     private final Map<QName,GeneratedTOBuilder> identities = new HashMap<>();
     private final Set<GeneratedTypeBuilder> topLevelNodes = new HashSet<>();
     private final List<GeneratedTypeBuilder> augmentations = new ArrayList<>();
-    private final BiMap<Type, AugmentationSchemaNode> typeToAugmentation = HashBiMap.create();
+    private final Multimap<Type,AugmentationSchemaNode> typeToAugmentations = HashMultimap.create();
     private final BiMap<SchemaPath,Type> targetToAugmentation = HashBiMap.create();
     private final Map<Type,Object> typeToSchema = new HashMap<>();
     private final Multimap<Type, Type> choiceToCases = HashMultimap.create();
@@ -169,8 +169,8 @@ public final class ModuleContext {
         return Collections.unmodifiableList(this.augmentations);
     }
 
-    public BiMap<Type, AugmentationSchemaNode> getTypeToAugmentation() {
-        return Maps.unmodifiableBiMap(this.typeToAugmentation);
+    public Multimap<Type, AugmentationSchemaNode> getTypeToAugmentations() {
+        return Multimaps.unmodifiableMultimap(this.typeToAugmentations);
     }
 
     public BiMap<SchemaPath, Type> getTargetToAugmentation() {
@@ -178,10 +178,14 @@ public final class ModuleContext {
     }
 
     public void addTypeToAugmentation(final GeneratedTypeBuilder builder, final AugmentationSchemaNode schema) {
-        this.typeToAugmentation.put(builder, schema);
         this.typeToSchema.put(builder, schema);
     }
 
+    public void addTypeToAugmentations(final GeneratedTypeBuilder builder,
+            final List<AugmentationSchemaNode> schemaList) {
+        schemaList.forEach(augmentNode -> this.typeToAugmentations.put(builder, augmentNode));
+    }
+
     public void addTargetToAugmentation(final GeneratedTypeBuilder builder, final SchemaPath augmentTarget) {
         this.targetToAugmentation.put(augmentTarget, builder);
     }