Binding generator v2 - augments fix #2

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

diff --git a/binding2/mdsal-binding2-generator-impl/src/main/java/org/opendaylight/mdsal/binding/javav2/generator/impl/ModuleContext.java b/binding2/mdsal-binding2-generator-impl/src/main/java/org/opendaylight/mdsal/binding/javav2/generator/impl/ModuleContext.java
index 4a80b28c80db90a6bcf410617319bad0aa9842f1..b240b1183b071b0487d18d6a4f7939bce5bf97e6 100644 (file)
--- a/binding2/mdsal-binding2-generator-impl/src/main/java/org/opendaylight/mdsal/binding/javav2/generator/impl/ModuleContext.java
+++ b/binding2/mdsal-binding2-generator-impl/src/main/java/org/opendaylight/mdsal/binding/javav2/generator/impl/ModuleContext.java
@@ -50,11 +50,13 @@ public final class ModuleContext {
     private final Set<GeneratedTypeBuilder> topLevelNodes = new HashSet<>();
     private final List<GeneratedTypeBuilder> augmentations = new ArrayList<>();
     private final BiMap<Type,AugmentationSchema> typeToAugmentation = HashBiMap.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();
     private final BiMap<Type,ChoiceCaseNode> caseTypeToSchema = HashBiMap.create();
     private final Map<SchemaPath, Type> innerTypes = new HashMap<>();
 
+
     List<Type> getGeneratedTypes() {
         final List<Type> result = new ArrayList<>();
 
@@ -163,11 +165,19 @@ public final class ModuleContext {
         return Maps.unmodifiableBiMap(this.typeToAugmentation);
     }
 
+    public BiMap<SchemaPath, Type> getTargetToAugmentation() {
+        return Maps.unmodifiableBiMap(this.targetToAugmentation);
+    }
+
     public void addTypeToAugmentation(final GeneratedTypeBuilder builder, final AugmentationSchema schema) {
         this.typeToAugmentation.put(builder, schema);
         this.typeToSchema.put(builder, schema);
     }
 
+    public void addTargetToAugmentation(final GeneratedTypeBuilder builder, final SchemaPath augmentTarget) {
+        this.targetToAugmentation.put(augmentTarget, builder);
+    }
+
     public void addChoiceToCaseMapping(final Type choiceType, final Type caseType, final ChoiceCaseNode schema) {
         this.choiceToCases.put(choiceType, caseType);
         this.caseTypeToSchema.put(caseType, schema);