BUG-1196: fixed bug in choice case codec.
[yangtools.git] / code-generator / binding-generator-impl / src / main / java / org / opendaylight / yangtools / sal / binding / generator / impl / LazyGeneratedCodecRegistry.java
index a8488bd34a5e74896cc3dcf79f06bf5bcbd487ad..c251ba18203add8324274bf6a6a8e63563eb263e 100644 (file)
@@ -9,191 +9,222 @@ package org.opendaylight.yangtools.sal.binding.generator.impl;
 
 import java.lang.ref.WeakReference;
 import java.lang.reflect.Field;
-import java.lang.reflect.ParameterizedType;
-import java.util.*;
+import java.util.AbstractMap.SimpleEntry;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
 import java.util.Map.Entry;
-import java.util.concurrent.Callable;
+import java.util.Set;
+import java.util.WeakHashMap;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
-import org.opendaylight.yangtools.sal.binding.generator.util.ClassLoaderUtils;
-import org.opendaylight.yangtools.yang.data.impl.codec.*;
 import org.opendaylight.yangtools.binding.generator.util.ReferencedTypeImpl;
 import org.opendaylight.yangtools.binding.generator.util.Types;
 import org.opendaylight.yangtools.concepts.Delegator;
 import org.opendaylight.yangtools.concepts.Identifiable;
+import org.opendaylight.yangtools.sal.binding.generator.api.ClassLoadingStrategy;
+import org.opendaylight.yangtools.sal.binding.generator.util.CodeGenerationException;
 import org.opendaylight.yangtools.sal.binding.model.api.ConcreteType;
 import org.opendaylight.yangtools.sal.binding.model.api.Type;
 import org.opendaylight.yangtools.sal.binding.model.api.type.builder.GeneratedTOBuilder;
-import org.opendaylight.yangtools.sal.binding.model.api.type.builder.GeneratedTypeBuilder;
-import org.opendaylight.yangtools.yang.binding.*;
+import org.opendaylight.yangtools.yang.binding.Augmentable;
+import org.opendaylight.yangtools.yang.binding.Augmentation;
+import org.opendaylight.yangtools.yang.binding.BaseIdentity;
+import org.opendaylight.yangtools.yang.binding.BindingCodec;
+import org.opendaylight.yangtools.yang.binding.DataContainer;
+import org.opendaylight.yangtools.yang.binding.DataObject;
+import org.opendaylight.yangtools.yang.binding.Identifier;
+import org.opendaylight.yangtools.yang.binding.InstanceIdentifier;
 import org.opendaylight.yangtools.yang.binding.util.BindingReflections;
 import org.opendaylight.yangtools.yang.common.QName;
 import org.opendaylight.yangtools.yang.data.api.CompositeNode;
 import org.opendaylight.yangtools.yang.data.api.Node;
 import org.opendaylight.yangtools.yang.data.impl.CompositeNodeTOImpl;
-import org.opendaylight.yangtools.yang.model.api.*;
-import org.opendaylight.yangtools.yang.model.util.SchemaContextUtil;
+import org.opendaylight.yangtools.yang.data.impl.codec.AugmentationCodec;
+import org.opendaylight.yangtools.yang.data.impl.codec.ChoiceCaseCodec;
+import org.opendaylight.yangtools.yang.data.impl.codec.ChoiceCodec;
+import org.opendaylight.yangtools.yang.data.impl.codec.CodecRegistry;
+import org.opendaylight.yangtools.yang.data.impl.codec.DataContainerCodec;
+import org.opendaylight.yangtools.yang.data.impl.codec.DomCodec;
+import org.opendaylight.yangtools.yang.data.impl.codec.IdentifierCodec;
+import org.opendaylight.yangtools.yang.data.impl.codec.IdentityCodec;
+import org.opendaylight.yangtools.yang.data.impl.codec.InstanceIdentifierCodec;
+import org.opendaylight.yangtools.yang.data.impl.codec.ValueWithQName;
+import org.opendaylight.yangtools.yang.model.api.AugmentationSchema;
+import org.opendaylight.yangtools.yang.model.api.AugmentationTarget;
+import org.opendaylight.yangtools.yang.model.api.ChoiceCaseNode;
+import org.opendaylight.yangtools.yang.model.api.ChoiceNode;
+import org.opendaylight.yangtools.yang.model.api.DataNodeContainer;
+import org.opendaylight.yangtools.yang.model.api.DataSchemaNode;
+import org.opendaylight.yangtools.yang.model.api.LeafListSchemaNode;
+import org.opendaylight.yangtools.yang.model.api.LeafSchemaNode;
+import org.opendaylight.yangtools.yang.model.api.Module;
+import org.opendaylight.yangtools.yang.model.api.SchemaContext;
+import org.opendaylight.yangtools.yang.model.api.SchemaContextListener;
+import org.opendaylight.yangtools.yang.model.api.SchemaPath;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.base.Optional;
 import com.google.common.base.Preconditions;
-
-public class LazyGeneratedCodecRegistry implements //
+import com.google.common.collect.BiMap;
+import com.google.common.collect.HashBiMap;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Multimaps;
+
+class LazyGeneratedCodecRegistry implements //
         CodecRegistry, //
-        SchemaServiceListener, //
+        SchemaContextListener, //
         GeneratorListener {
 
-    private final static Logger LOG = LoggerFactory.getLogger(LazyGeneratedCodecRegistry.class);
-    private final static LateMixinCodec NOT_READY_CODEC = new LateMixinCodec();
-
-    private final InstanceIdentifierCodec instanceIdentifierCodec = new InstanceIdentifierCodecImpl(this);
-    private final IdentityCompositeCodec identityRefCodec = new IdentityCompositeCodec();
-
-    private TransformerGenerator generator;
+    private static final Logger LOG = LoggerFactory.getLogger(LazyGeneratedCodecRegistry.class);
 
     // Concrete class to codecs
-    private static final Map<Class<?>, DataContainerCodec<?>> containerCodecs = Collections.synchronizedMap(new WeakHashMap<Class<?>, DataContainerCodec<?>>());
-    private static final Map<Class<?>, IdentifierCodec<?>> identifierCodecs = Collections.synchronizedMap(new WeakHashMap<Class<?>, IdentifierCodec<?>>());
-    private static final Map<Class<?>, ChoiceCodecImpl<?>> choiceCodecs = Collections.synchronizedMap(new WeakHashMap<Class<?>, ChoiceCodecImpl<?>>());
-    private static final Map<Class<?>, ChoiceCaseCodecImpl<?>> caseCodecs = Collections.synchronizedMap(new WeakHashMap<Class<?>, ChoiceCaseCodecImpl<?>>());
-    private static final Map<Class<?>, AugmentableCompositeCodec> augmentableCodecs = Collections.synchronizedMap(new WeakHashMap<Class<?>, AugmentableCompositeCodec>());
-    private static final Map<Class<?>, AugmentationCodec<?>> augmentationCodecs = Collections.synchronizedMap(new WeakHashMap<Class<?>, AugmentationCodec<?>>());
-    private static final Map<Class<?>, QName> identityQNames = Collections.synchronizedMap(new WeakHashMap<Class<?>, QName>());
+    private static final Map<Class<?>, DataContainerCodec<?>> containerCodecs = Collections
+            .synchronizedMap(new WeakHashMap<Class<?>, DataContainerCodec<?>>());
+    private static final Map<Class<?>, IdentifierCodec<?>> identifierCodecs = Collections
+            .synchronizedMap(new WeakHashMap<Class<?>, IdentifierCodec<?>>());
+    private static final Map<Class<?>, PublicChoiceCodecImpl<?>> choiceCodecs = Collections
+            .synchronizedMap(new WeakHashMap<Class<?>, PublicChoiceCodecImpl<?>>());
+    private static final Map<Class<?>, ChoiceCaseCodecImpl<?>> caseCodecs = Collections
+            .synchronizedMap(new WeakHashMap<Class<?>, ChoiceCaseCodecImpl<?>>());
+    private static final Map<Class<?>, AugmentableDispatchCodec> augmentableCodecs = Collections
+            .synchronizedMap(new WeakHashMap<Class<?>, AugmentableDispatchCodec>());
+    private static final Map<Class<?>, AugmentationCodecWrapper<?>> augmentationCodecs = Collections
+            .synchronizedMap(new WeakHashMap<Class<?>, AugmentationCodecWrapper<?>>());
+    private static final Map<Class<?>, QName> identityQNames = Collections
+            .synchronizedMap(new WeakHashMap<Class<?>, QName>());
     private static final Map<QName, Type> qnamesToIdentityMap = new ConcurrentHashMap<>();
     /** Binding type to encountered classes mapping **/
     @SuppressWarnings("rawtypes")
     private static final Map<Type, WeakReference<Class>> typeToClass = new ConcurrentHashMap<>();
 
-    @SuppressWarnings("rawtypes")
-    private static final ConcurrentMap<Type, ChoiceCaseCodecImpl> typeToCaseCodecs = new ConcurrentHashMap<>();
-
-    private CaseClassMapFacade classToCaseRawCodec = new CaseClassMapFacade();
+    private static final ConcurrentMap<Type, ChoiceCaseNode> caseTypeToCaseSchema = new ConcurrentHashMap<>();
 
-    private static final Map<SchemaPath, GeneratedTypeBuilder> pathToType = new ConcurrentHashMap<>();
+    private static final Map<SchemaPath, Type> pathToType = new ConcurrentHashMap<>();
     private static final Map<List<QName>, Type> pathToInstantiatedType = new ConcurrentHashMap<>();
     private static final Map<Type, QName> typeToQname = new ConcurrentHashMap<>();
+    private static final BiMap<Type, AugmentationSchema> typeToAugment = HashBiMap
+            .create(new ConcurrentHashMap<Type, AugmentationSchema>());
 
+    private static final Multimap<Type, Type> augmentableToAugmentations = Multimaps.synchronizedMultimap(HashMultimap
+            .<Type, Type> create());
+    private static final Multimap<Type, Type> choiceToCases = Multimaps.synchronizedMultimap(HashMultimap
+            .<Type, Type> create());
+
+    private final InstanceIdentifierCodec instanceIdentifierCodec = new InstanceIdentifierCodecImpl(this);
+    private final IdentityCompositeCodec identityRefCodec = new IdentityCompositeCodec();
+    private final ClassLoadingStrategy classLoadingStrategy;
+    private final AbstractTransformerGenerator generator;
     private final SchemaLock lock;
 
+    // FIXME: how is this protected?
     private SchemaContext currentSchema;
 
-    LazyGeneratedCodecRegistry(SchemaLock lock) {
+    LazyGeneratedCodecRegistry(final SchemaLock lock, final AbstractTransformerGenerator generator,
+            final ClassLoadingStrategy classLoadingStrategy) {
         this.lock = Preconditions.checkNotNull(lock);
+        this.classLoadingStrategy = Preconditions.checkNotNull(classLoadingStrategy);
+        this.generator = Preconditions.checkNotNull(generator);
     }
 
     public SchemaLock getLock() {
         return lock;
     }
 
-    public TransformerGenerator getGenerator() {
-        return generator;
-    }
-
-    public void setGenerator(TransformerGenerator generator) {
-        this.generator = generator;
-    }
-
     @Override
     public InstanceIdentifierCodec getInstanceIdentifierCodec() {
         return instanceIdentifierCodec;
     }
 
+    @SuppressWarnings("unchecked")
     @Override
-    public <T extends Augmentation<?>> AugmentationCodec<T> getCodecForAugmentation(Class<T> object) {
-        AugmentationCodec<T> codec = null;
+    public <T extends Augmentation<?>> AugmentationCodecWrapper<T> getCodecForAugmentation(final Class<T> augClass) {
+        AugmentationCodecWrapper<T> codec = null;
         @SuppressWarnings("rawtypes")
-        AugmentationCodec potentialCodec = augmentationCodecs.get(object);
+        AugmentationCodecWrapper potentialCodec = augmentationCodecs.get(augClass);
         if (potentialCodec != null) {
             codec = potentialCodec;
-        } else
-            try {
-                lock.waitForSchema(object);
-                Class<? extends BindingCodec<Map<QName, Object>, Object>> augmentRawCodec = generator
-                        .augmentationTransformerFor(object);
-                BindingCodec<Map<QName, Object>, Object> rawCodec = augmentRawCodec.newInstance();
-                codec = new AugmentationCodecWrapper<T>(rawCodec);
-                augmentationCodecs.put(augmentRawCodec, codec);
-            } catch (InstantiationException e) {
-                LOG.error("Can not instantiate raw augmentation codec {}", object.getSimpleName(), e);
-            } catch (IllegalAccessException e) {
-                LOG.debug("BUG: Constructor for {} is not accessible.", object.getSimpleName(), e);
-            }
-        Class<? extends Augmentable<?>> objectSupertype = getAugmentableArgumentFrom(object);
-        if (objectSupertype != null) {
-            getAugmentableCodec(objectSupertype).addAugmentationCodec(object, codec);
         } else {
-            LOG.warn("Could not find augmentation target for augmentation {}", object);
+            lock.waitForSchema(augClass);
+            Class<? extends BindingCodec<Map<QName, Object>, Object>> augmentRawCodec = generator
+                    .augmentationTransformerFor(augClass);
+
+            BindingCodec<Map<QName, Object>, Object> rawCodec = newInstanceOf(augmentRawCodec);
+            codec = new AugmentationCodecWrapper<T>(rawCodec, augClass);
+            augmentationCodecs.put(augClass, codec);
+        }
+
+        final Class<? extends Augmentable<?>> objectSupertype;
+        try {
+            objectSupertype = BindingReflections.findAugmentationTarget(augClass);
+        } catch (Exception e) {
+            LOG.warn("Failed to find target for augmentation {}, ignoring it", augClass, e);
+            return codec;
         }
+
+        if (objectSupertype == null) {
+            LOG.warn("Augmentation target for {} not found, ignoring it", augClass);
+            return codec;
+        }
+
+        getAugmentableCodec(objectSupertype).addImplementation(codec);
         return codec;
     }
 
+    @SuppressWarnings("unchecked")
     @Override
-    public QName getQNameForAugmentation(Class<?> cls) {
+    public QName getQNameForAugmentation(final Class<?> cls) {
         Preconditions.checkArgument(Augmentation.class.isAssignableFrom(cls));
-        return getCodecForAugmentation((Class<? extends Augmentation>)cls).getAugmentationQName();
-    }
-
-    private static Class<? extends Augmentable<?>> getAugmentableArgumentFrom(
-            final Class<? extends Augmentation<?>> augmentation) {
-        try {
-            Class<? extends Augmentable<?>> ret = ClassLoaderUtils.withClassLoader(augmentation.getClassLoader(),
-                    new Callable<Class<? extends Augmentable<?>>>() {
-                        @Override
-                        @SuppressWarnings("unchecked")
-                        public Class<? extends Augmentable<?>> call() throws Exception {
-                            for (java.lang.reflect.Type supertype : augmentation.getGenericInterfaces()) {
-                                if (supertype instanceof ParameterizedType
-                                        && Augmentation.class.equals(((ParameterizedType) supertype).getRawType())) {
-                                    ParameterizedType augmentationGeneric = (ParameterizedType) supertype;
-                                    return (Class<? extends Augmentable<?>>) augmentationGeneric
-                                            .getActualTypeArguments()[0];
-                                }
-                            }
-                            return null;
-                        }
-                    });
-            return ret;
-        } catch (Exception e) {
-            LOG.debug("Could not find augmentable for {} using {}", augmentation, augmentation.getClassLoader(), e);
-            return null;
-        }
+        return getCodecForAugmentation((Class<? extends Augmentation<?>>) cls).getAugmentationQName();
     }
 
     @Override
-    public Class<?> getClassForPath(List<QName> names) {
-        DataSchemaNode node = getSchemaNode(names);
-        SchemaPath path = node.getPath();
-        Type type = pathToType.get(path);
-        if (type != null) {
-            type = new ReferencedTypeImpl(type.getPackageName(), type.getName());
+    public Class<?> getClassForPath(final List<QName> names) {
+        final DataSchemaNode node = getSchemaNode(names);
+        final SchemaPath path = node.getPath();
+        final Type t = pathToType.get(path);
+
+        final Type type;
+        if (t != null) {
+            type = new ReferencedTypeImpl(t.getPackageName(), t.getName());
         } else {
             type = pathToInstantiatedType.get(names);
+            Preconditions.checkState(type != null, "Failed to lookup instantiated type for path %s", path);
         }
+
         @SuppressWarnings("rawtypes")
-        WeakReference<Class> weakRef = typeToClass.get(type);
-        if (weakRef == null) {
-            LOG.error("Could not find loaded class for path: {} and type: {}", path, type.getFullyQualifiedName());
-        }
+        final WeakReference<Class> weakRef = typeToClass.get(type);
+        Preconditions.checkState(weakRef != null, "Could not find loaded class for path: %s and type: %s", path,
+                type.getFullyQualifiedName());
         return weakRef.get();
     }
 
     @Override
-    public void putPathToClass(List<QName> names, Class<?> cls) {
-        Type reference = Types.typeForClass(cls);
+    public void putPathToClass(final List<QName> names, final Class<?> cls) {
+        final Type reference = Types.typeForClass(cls);
         pathToInstantiatedType.put(names, reference);
+        LOG.trace("Path {} attached to class {} reference {}", names, cls, reference);
         bindingClassEncountered(cls);
     }
 
     @Override
-    public IdentifierCodec<?> getKeyCodecForPath(List<QName> names) {
+    public IdentifierCodec<?> getKeyCodecForPath(final List<QName> names) {
         @SuppressWarnings("unchecked")
         Class<? extends Identifiable<?>> cls = (Class<? extends Identifiable<?>>) getClassForPath(names);
         return getIdentifierCodecForIdentifiable(cls);
     }
 
     @Override
-    public <T extends DataContainer> DataContainerCodec<T> getCodecForDataObject(Class<T> type) {
+    public <T extends DataContainer> DataContainerCodec<T> getCodecForDataObject(final Class<T> type) {
         @SuppressWarnings("unchecked")
         DataContainerCodec<T> ret = (DataContainerCodec<T>) containerCodecs.get(type);
         if (ret != null) {
@@ -206,8 +237,9 @@ public class LazyGeneratedCodecRegistry implements //
         return newWrapper;
     }
 
+    @Override
     @SuppressWarnings("rawtypes")
-    public void bindingClassEncountered(Class cls) {
+    public void bindingClassEncountered(final Class cls) {
 
         ConcreteType typeRef = Types.typeForClass(cls);
         if (typeToClass.containsKey(typeRef)) {
@@ -225,20 +257,23 @@ public class LazyGeneratedCodecRegistry implements //
     }
 
     @Override
-    public void onClassProcessed(Class<?> cls) {
+    public void onClassProcessed(final Class<?> cls) {
         ConcreteType typeRef = Types.typeForClass(cls);
         if (typeToClass.containsKey(typeRef)) {
             return;
         }
         LOG.trace("Binding Class {} encountered.", cls);
-        WeakReference<Class> weakRef = new WeakReference<>((Class) cls);
+        @SuppressWarnings("rawtypes")
+        WeakReference<Class> weakRef = new WeakReference<Class>(cls);
         typeToClass.put(typeRef, weakRef);
     }
 
-    private DataSchemaNode getSchemaNode(List<QName> path) {
+    private DataSchemaNode getSchemaNode(final List<QName> path) {
         QName firstNode = path.get(0);
         DataNodeContainer previous = currentSchema.findModuleByNamespaceAndRevision(firstNode.getNamespace(),
                 firstNode.getRevision());
+        Preconditions.checkArgument(previous != null, "Failed to find module %s for path %s", firstNode, path);
+
         Iterator<QName> iterator = path.iterator();
         while (iterator.hasNext()) {
             QName arg = iterator.next();
@@ -256,7 +291,7 @@ public class LazyGeneratedCodecRegistry implements //
         return (DataSchemaNode) previous;
     }
 
-    private DataSchemaNode searchInChoices(DataNodeContainer node, QName arg) {
+    private DataSchemaNode searchInChoices(final DataNodeContainer node, final QName arg) {
         Set<DataSchemaNode> children = node.getChildNodes();
         for (DataSchemaNode child : children) {
             if (child instanceof ChoiceNode) {
@@ -270,7 +305,7 @@ public class LazyGeneratedCodecRegistry implements //
         return null;
     }
 
-    private DataSchemaNode searchInCases(ChoiceNode choiceNode, QName arg) {
+    private DataSchemaNode searchInCases(final ChoiceNode choiceNode, final QName arg) {
         Set<ChoiceCaseNode> cases = choiceNode.getCases();
         for (ChoiceCaseNode caseNode : cases) {
             DataSchemaNode node = caseNode.getDataChildByName(arg);
@@ -281,20 +316,24 @@ public class LazyGeneratedCodecRegistry implements //
         return null;
     }
 
-    private <T> T newInstanceOf(Class<?> newType) {
+    private static <T> T newInstanceOf(final Class<?> cls) {
         try {
             @SuppressWarnings("unchecked")
-            T ret = (T) newType.newInstance();
+            T ret = (T) cls.newInstance();
             return ret;
         } catch (InstantiationException e) {
-            throw new IllegalStateException(e);
+            LOG.error("Failed to instantiate codec {}", cls.getSimpleName(), e);
+            throw new IllegalStateException(String.format("Failed to instantiate codec %s", cls), e);
         } catch (IllegalAccessException e) {
-            throw new IllegalStateException(e);
+            LOG.debug(
+                    "Run-time consistency issue: constructor for {} is not available. This indicates either a code generation bug or a misconfiguration of JVM.",
+                    cls.getSimpleName(), e);
+            throw new IllegalStateException(String.format("Cannot access contructor of %s", cls), e);
         }
     }
 
     @Override
-    public <T extends Identifiable<?>> IdentifierCodec<?> getIdentifierCodecForIdentifiable(Class<T> type) {
+    public <T extends Identifiable<?>> IdentifierCodec<?> getIdentifierCodecForIdentifiable(final Class<T> type) {
         IdentifierCodec<?> obj = identifierCodecs.get(type);
         if (obj != null) {
             return obj;
@@ -313,20 +352,21 @@ public class LazyGeneratedCodecRegistry implements //
         return identityRefCodec;
     }
 
+    @SuppressWarnings("unchecked")
     @Override
-    public <T extends BaseIdentity> IdentityCodec<T> getCodecForIdentity(Class<T> codec) {
+    public <T extends BaseIdentity> IdentityCodec<T> getCodecForIdentity(final Class<T> codec) {
         bindingClassEncountered(codec);
         return identityRefCodec;
     }
 
     @Override
-    public void onCodecCreated(Class<?> cls) {
+    public void onCodecCreated(final Class<?> cls) {
         CodecMapping.setIdentifierCodec(cls, instanceIdentifierCodec);
         CodecMapping.setIdentityRefCodec(cls, identityRefCodec);
     }
 
     @Override
-    public <T extends Identifier<?>> IdentifierCodec<T> getCodecForIdentifier(Class<T> object) {
+    public <T extends Identifier<?>> IdentifierCodec<T> getCodecForIdentifier(final Class<T> object) {
         @SuppressWarnings("unchecked")
         IdentifierCodec<T> obj = (IdentifierCodec<T>) identifierCodecs.get(object);
         if (obj != null) {
@@ -342,23 +382,21 @@ public class LazyGeneratedCodecRegistry implements //
     }
 
     @SuppressWarnings("rawtypes")
-    public ChoiceCaseCodecImpl getCaseCodecFor(Class caseClass) {
+    public ChoiceCaseCodecImpl getCaseCodecFor(final Class caseClass) {
         ChoiceCaseCodecImpl<?> potential = caseCodecs.get(caseClass);
         if (potential != null) {
             return potential;
         }
         ConcreteType typeref = Types.typeForClass(caseClass);
-        ChoiceCaseCodecImpl caseCodec = typeToCaseCodecs.get(typeref);
+        ChoiceCaseNode caseSchema = caseTypeToCaseSchema.get(typeref);
 
-        Preconditions.checkState(caseCodec != null, "Case Codec was not created proactivelly for %s", caseClass.getName());
-        Preconditions.checkState(caseCodec.getSchema() != null, "Case schema is not available for %s", caseClass.getName());
-        @SuppressWarnings("unchecked")
-        Class<? extends BindingCodec> newCodec = generator.caseCodecFor(caseClass, caseCodec.getSchema());
+        Preconditions.checkState(caseSchema != null, "Case schema is not available for %s", caseClass.getName());
+        Class<? extends BindingCodec> newCodec = generator.caseCodecFor(caseClass, caseSchema);
         BindingCodec newInstance = newInstanceOf(newCodec);
-        caseCodec.setDelegate(newInstance);
+        ChoiceCaseCodecImpl caseCodec = new ChoiceCaseCodecImpl(caseClass, caseSchema, newInstance);
         caseCodecs.put(caseClass, caseCodec);
 
-        for (Entry<Class<?>, ChoiceCodecImpl<?>> choice : choiceCodecs.entrySet()) {
+        for (Entry<Class<?>, PublicChoiceCodecImpl<?>> choice : choiceCodecs.entrySet()) {
             if (choice.getKey().isAssignableFrom(caseClass)) {
                 choice.getValue().cases.put(caseClass, caseCodec);
             }
@@ -366,106 +404,81 @@ public class LazyGeneratedCodecRegistry implements //
         return caseCodec;
     }
 
-    public void onModuleContextAdded(SchemaContext schemaContext, Module module, ModuleContext context) {
+    public void onModuleContextAdded(final SchemaContext schemaContext, final Module module, final ModuleContext context) {
         pathToType.putAll(context.getChildNodes());
+
+        BiMap<Type, AugmentationSchema> bimap = context.getTypeToAugmentation();
+        for (Map.Entry<Type, AugmentationSchema> entry : bimap.entrySet()) {
+            Type key = entry.getKey();
+            AugmentationSchema value = entry.getValue();
+            Set<DataSchemaNode> augmentedNodes = value.getChildNodes();
+            if (augmentedNodes != null && !(augmentedNodes.isEmpty())) {
+                typeToAugment.put(key, value);
+            }
+        }
+
         qnamesToIdentityMap.putAll(context.getIdentities());
         for (Entry<QName, GeneratedTOBuilder> identity : context.getIdentities().entrySet()) {
             typeToQname.put(
                     new ReferencedTypeImpl(identity.getValue().getPackageName(), identity.getValue().getName()),
                     identity.getKey());
         }
-        captureCases(context.getCases(), schemaContext);
-    }
-
-    private void captureCases(Map<SchemaPath, GeneratedTypeBuilder> cases, SchemaContext module) {
-        for (Entry<SchemaPath, GeneratedTypeBuilder> caseNode : cases.entrySet()) {
-            ReferencedTypeImpl typeref = new ReferencedTypeImpl(caseNode.getValue().getPackageName(), caseNode
-                    .getValue().getName());
-
-            pathToType.put(caseNode.getKey(), caseNode.getValue());
 
-            ChoiceCaseNode node = (ChoiceCaseNode) SchemaContextUtil.findDataSchemaNode(module, caseNode.getKey());
-
-            if (node == null) {
-                LOG.error("YANGTools Bug: SchemaNode for {}, with path {} was not found in context.",
-                        typeref.getFullyQualifiedName(), caseNode.getKey());
-                @SuppressWarnings("rawtypes")
-                ChoiceCaseCodecImpl value = new ChoiceCaseCodecImpl();
-                typeToCaseCodecs.putIfAbsent(typeref, value);
-                continue;
-            }
-            @SuppressWarnings("rawtypes")
-            ChoiceCaseCodecImpl value = new ChoiceCaseCodecImpl(node);
-            typeToCaseCodecs.putIfAbsent(typeref, value);
+        synchronized (augmentableToAugmentations) {
+            augmentableToAugmentations.putAll(context.getAugmentableToAugmentations());
+        }
+        synchronized (choiceToCases) {
+            choiceToCases.putAll(context.getChoiceToCases());
+        }
+        synchronized (caseTypeToCaseSchema) {
+            caseTypeToCaseSchema.putAll(context.getCaseTypeToSchemas());
         }
     }
 
     @Override
-    public void onGlobalContextUpdated(SchemaContext context) {
+    public void onGlobalContextUpdated(final SchemaContext context) {
         currentSchema = context;
     }
 
     @SuppressWarnings({ "unchecked", "rawtypes" })
     @Override
-    public void onChoiceCodecCreated(Class<?> choiceClass,
-            Class<? extends BindingCodec<Map<QName, Object>, Object>> choiceCodec, ChoiceNode schema) {
+    public void onChoiceCodecCreated(final Class<?> choiceClass,
+            final Class<? extends BindingCodec<Map<QName, Object>, Object>> choiceCodec, final ChoiceNode schema) {
         ChoiceCodec<?> oldCodec = choiceCodecs.get(choiceClass);
         Preconditions.checkState(oldCodec == null);
         BindingCodec<Map<QName, Object>, Object> delegate = newInstanceOf(choiceCodec);
-        ChoiceCodecImpl<?> newCodec = new ChoiceCodecImpl(delegate);
+        PublicChoiceCodecImpl<?> newCodec = new PublicChoiceCodecImpl(delegate);
+        DispatchChoiceCodecImpl dispatchCodec = new DispatchChoiceCodecImpl(choiceClass);
         choiceCodecs.put(choiceClass, newCodec);
-        CodecMapping.setClassToCaseMap(choiceCodec, (Map<Class<?>, BindingCodec<?, ?>>) classToCaseRawCodec);
-        CodecMapping.setCompositeNodeToCaseMap(choiceCodec, newCodec.getCompositeToCase());
-
-        tryToCreateCasesCodecs(schema);
-
-    }
-
-    private void tryToCreateCasesCodecs(ChoiceNode schema) {
-        for (ChoiceCaseNode caseNode : schema.getCases()) {
-            SchemaPath path = caseNode.getPath();
-            GeneratedTypeBuilder type;
-            if (path != null && (type = pathToType.get(path)) != null) {
-                ReferencedTypeImpl typeref = new ReferencedTypeImpl(type.getPackageName(), type.getName());
-                ChoiceCaseCodecImpl partialCodec = typeToCaseCodecs.get(typeref);
-                if (partialCodec.getSchema() == null) {
-                    partialCodec.setSchema(caseNode);
-                }
-
-                Class<?> caseClass = ClassLoaderUtils.tryToLoadClassWithTCCL(type.getFullyQualifiedName());
-                if (caseClass != null) {
-                    getCaseCodecFor(caseClass);
-                }
-            }
-        }
-
+        CodecMapping.setDispatchCodec(choiceCodec, dispatchCodec);
     }
 
     @Override
-    public void onValueCodecCreated(Class<?> valueClass, Class<?> valueCodec) {
+    public void onValueCodecCreated(final Class<?> valueClass, final Class<?> valueCodec) {
     }
 
     @Override
-    public void onCaseCodecCreated(Class<?> choiceClass,
-            Class<? extends BindingCodec<Map<QName, Object>, Object>> choiceCodec) {
+    public void onCaseCodecCreated(final Class<?> choiceClass,
+            final Class<? extends BindingCodec<Map<QName, Object>, Object>> choiceCodec) {
     }
 
     @Override
-    public void onDataContainerCodecCreated(Class<?> dataClass, Class<? extends BindingCodec<?, ?>> dataCodec) {
+    public void onDataContainerCodecCreated(final Class<?> dataClass,
+            final Class<? extends BindingCodec<?, ?>> dataCodec) {
         if (Augmentable.class.isAssignableFrom(dataClass)) {
-            AugmentableCompositeCodec augmentableCodec = getAugmentableCodec(dataClass);
+            AugmentableDispatchCodec augmentableCodec = getAugmentableCodec(dataClass);
             CodecMapping.setAugmentationCodec(dataCodec, augmentableCodec);
         }
-
     }
 
-    public AugmentableCompositeCodec getAugmentableCodec(Class<?> dataClass) {
-        AugmentableCompositeCodec ret = augmentableCodecs.get(dataClass);
+    public synchronized AugmentableDispatchCodec getAugmentableCodec(final Class<?> dataClass) {
+        AugmentableDispatchCodec ret = augmentableCodecs.get(dataClass);
         if (ret != null) {
             return ret;
         }
-        ret = new AugmentableCompositeCodec(dataClass);
+        ret = new AugmentableDispatchCodec(dataClass);
         augmentableCodecs.put(dataClass, ret);
+        ret.tryToLoadImplementations();
         return ret;
     }
 
@@ -479,27 +492,28 @@ public class LazyGeneratedCodecRegistry implements //
             return delegate;
         }
 
-        public IntermediateCodec(BindingCodec<Map<QName, Object>, Object> delegate) {
+        public IntermediateCodec(final BindingCodec<Map<QName, Object>, Object> delegate) {
             this.delegate = delegate;
         }
 
         @Override
-        public Node<?> serialize(ValueWithQName<T> input) {
+        public Node<?> serialize(final ValueWithQName<T> input) {
             Map<QName, Object> intermediateOutput = delegate.serialize(input);
             return IntermediateMapping.toNode(intermediateOutput);
         }
+
     }
 
     private static class IdentifierCodecImpl<T extends Identifier<?>> //
             extends IntermediateCodec<T> //
             implements IdentifierCodec<T> {
 
-        public IdentifierCodecImpl(BindingCodec<Map<QName, Object>, Object> delegate) {
+        public IdentifierCodecImpl(final BindingCodec<Map<QName, Object>, Object> delegate) {
             super(delegate);
         }
 
         @Override
-        public ValueWithQName<T> deserialize(Node<?> input) {
+        public ValueWithQName<T> deserialize(final Node<?> input) {
             QName qname = input.getNodeType();
             @SuppressWarnings("unchecked")
             T value = (T) getDelegate().deserialize((Map<QName, Object>) input);
@@ -507,7 +521,15 @@ public class LazyGeneratedCodecRegistry implements //
         }
 
         @Override
-        public CompositeNode serialize(ValueWithQName<T> input) {
+        public ValueWithQName<T> deserialize(final Node<?> input, final InstanceIdentifier<?> bindingIdentifier) {
+            QName qname = input.getNodeType();
+            @SuppressWarnings("unchecked")
+            T value = (T) getDelegate().deserialize((Map<QName, Object>) input, bindingIdentifier);
+            return new ValueWithQName<T>(qname, value);
+        }
+
+        @Override
+        public CompositeNode serialize(final ValueWithQName<T> input) {
             return (CompositeNode) super.serialize(input);
         }
     }
@@ -516,12 +538,12 @@ public class LazyGeneratedCodecRegistry implements //
             extends IntermediateCodec<T> //
             implements DataContainerCodec<T> {
 
-        public DataContainerCodecImpl(BindingCodec<Map<QName, Object>, Object> delegate) {
+        public DataContainerCodecImpl(final BindingCodec<Map<QName, Object>, Object> delegate) {
             super(delegate);
         }
 
         @Override
-        public ValueWithQName<T> deserialize(Node<?> input) {
+        public ValueWithQName<T> deserialize(final Node<?> input) {
             if (input == null) {
                 return null;
             }
@@ -532,288 +554,413 @@ public class LazyGeneratedCodecRegistry implements //
         }
 
         @Override
-        public CompositeNode serialize(ValueWithQName<T> input) {
+        public ValueWithQName<T> deserialize(final Node<?> input, final InstanceIdentifier<?> bindingIdentifier) {
+            if (input == null) {
+                return null;
+            }
+            QName qname = input.getNodeType();
+            @SuppressWarnings("unchecked")
+            T value = (T) getDelegate().deserialize((Map<QName, Object>) input, bindingIdentifier);
+            return new ValueWithQName<T>(qname, value);
+        }
+
+        @Override
+        public CompositeNode serialize(final ValueWithQName<T> input) {
             return (CompositeNode) super.serialize(input);
         }
     }
 
+    private interface LocationAwareBindingCodec<P, I> extends BindingCodec<P, I> {
+
+        boolean isApplicable(InstanceIdentifier<?> location);
+
+        public Class<?> getDataType();
+
+    }
+
     @SuppressWarnings("rawtypes")
-    private static class ChoiceCaseCodecImpl<T extends DataContainer> implements ChoiceCaseCodec<T>, //
-            Delegator<BindingCodec> {
-        private boolean augmenting;
-        private BindingCodec delegate;
+    private abstract class LocationAwareDispatchCodec<T extends LocationAwareBindingCodec> implements BindingCodec {
 
-        private Set<String> validNames;
-        private Set<QName> validQNames;
-        private ChoiceCaseNode schema;
+        private final Map<Class, T> implementations = Collections.synchronizedMap(new WeakHashMap<Class, T>());
+        private final Set<InstanceIdentifier<?>> adaptedForPaths = new HashSet<>();
 
-        public void setSchema(ChoiceCaseNode caseNode) {
-            this.schema = schema;
-            this.schema = caseNode;
-            validNames = new HashSet<>();
-            validQNames = new HashSet<>();
-            for (DataSchemaNode node : caseNode.getChildNodes()) {
-                QName qname = node.getQName();
-                validQNames.add(qname);
-                validNames.add(qname.getLocalName());
+        protected Map<Class, T> getImplementations() {
+            return implementations;
+        }
+
+        protected void addImplementation(final T implementation) {
+            implementations.put(implementation.getDataType(), implementation);
+        }
+
+        @Override
+        public final Object deserialize(final Object input) {
+            throw new UnsupportedOperationException("Invocation of deserialize without Tree location is unsupported");
+        }
+
+        @Override
+        public Object serialize(final Object input) {
+            Preconditions.checkArgument(input instanceof DataContainer);
+            Class<? extends DataContainer> inputType = ((DataContainer) input).getImplementedInterface();
+            T implementation = implementations.get(inputType);
+            if (implementation == null) {
+                implementation = tryToLoadImplementationImpl(inputType);
             }
-            augmenting = caseNode.isAugmenting();
+
+            return null;
         }
 
-        public ChoiceCaseCodecImpl() {
-            this.delegate = NOT_READY_CODEC;
+        private T tryToLoadImplementationImpl(final Class<? extends DataContainer> inputType) {
+            T implementation = tryToLoadImplementation(inputType);
+            Preconditions.checkArgument(implementation != null, "Data type %s is not supported.", inputType);
+            addImplementation(implementation);
+            return implementation;
         }
 
-        public ChoiceCaseCodecImpl(ChoiceCaseNode caseNode) {
-            this.delegate = NOT_READY_CODEC;
-            setSchema(caseNode);
+        protected final synchronized void adaptForPath(final InstanceIdentifier<?> path) {
+            if (adaptedForPaths.contains(path)) {
+                return;
+            }
+            /**
+             * We search in schema context if the use of this location aware
+             * codec (augmentable codec, case codec) makes sense on provided
+             * location (path)
+             *
+             */
+            Optional<DataNodeContainer> contextNode = BindingSchemaContextUtils.findDataNodeContainer(currentSchema,
+                    path);
+            /**
+             * If context node is present, this codec makes sense on provided
+             * location.
+             *
+             */
+            if (contextNode.isPresent()) {
+                synchronized (this) {
+                    /**
+                     *
+                     * We adapt (turn on / off) possible implementations of
+                     * child codecs (augmentations, cases) based on this
+                     * location.
+                     *
+                     *
+                     */
+
+                    adaptForPathImpl(path, contextNode.get());
+                    try {
+                        /**
+                         * We trigger serialization of instance identifier, to
+                         * make sure instance identifier codec is aware of
+                         * combination of this path / augmentation / case
+                         */
+                        instanceIdentifierCodec.serialize(path);
+                    } catch (Exception e) {
+                        LOG.warn("Exception during preparation of instance identifier codec for  path {}.", path, e);
+                    }
+                    adaptedForPaths.add(path);
+                }
+            } else {
+                LOG.debug("Context node (parent node) not found for {}", path);
+            }
         }
 
+        protected abstract T tryToLoadImplementation(Class<? extends DataContainer> inputType);
+
+        protected abstract void tryToLoadImplementations();
+
+        protected abstract void adaptForPathImpl(InstanceIdentifier<?> path, DataNodeContainer ctx);
+    }
+
+    @SuppressWarnings("rawtypes")
+    private static class ChoiceCaseCodecImpl<T extends DataContainer> implements ChoiceCaseCodec<T>, //
+            Delegator<BindingCodec>, LocationAwareBindingCodec<Node<?>, ValueWithQName<T>> {
+        private final BindingCodec delegate;
+        private final ChoiceCaseNode schema;
+        private final Map<InstanceIdentifier<?>, ChoiceCaseNode> instantiatedLocations;
+        private final Class<?> dataType;
+
         @Override
-        public ValueWithQName<T> deserialize(Node<?> input) {
+        public boolean isApplicable(final InstanceIdentifier location) {
+            return instantiatedLocations.containsKey(location);
+        }
+
+        public ChoiceCaseCodecImpl(final Class<?> caseClass, final ChoiceCaseNode caseNode, final BindingCodec newInstance) {
+            this.delegate = newInstance;
+            this.dataType = caseClass;
+            this.schema = caseNode;
+            instantiatedLocations = new HashMap<>();
+        }
+
+        @Override
+        public ValueWithQName<T> deserialize(final Node<?> input) {
             throw new UnsupportedOperationException("Direct invocation of this codec is not allowed.");
         }
 
         @Override
-        public CompositeNode serialize(ValueWithQName<T> input) {
+        public ValueWithQName<T> deserialize(final Node<?> input, final InstanceIdentifier<?> bindingIdentifier) {
+            if (input == null) {
+                return null;
+            }
+            QName qname = input.getNodeType();
+            synchronized (instantiatedLocations) {
+                ChoiceCaseNode instantiation = instantiatedLocations.get(bindingIdentifier);
+                if (instantiation != null) {
+                    qname = instantiatedLocations.get(bindingIdentifier).getQName();
+                }
+            }
+            @SuppressWarnings("unchecked")
+            T value = (T) getDelegate().deserialize(new SimpleEntry(qname, input), bindingIdentifier);
+            return new ValueWithQName<T>(qname, value);
+        }
+
+        @Override
+        public CompositeNode serialize(final ValueWithQName<T> input) {
             throw new UnsupportedOperationException("Direct invocation of this codec is not allowed.");
         }
 
+        @Override
         public BindingCodec getDelegate() {
             return delegate;
         }
 
-        public void setDelegate(BindingCodec delegate) {
-            this.delegate = delegate;
-        }
-
         public ChoiceCaseNode getSchema() {
             return schema;
         }
 
         @Override
-        public boolean isAcceptable(Node<?> input) {
-            if (input instanceof CompositeNode) {
-                if (augmenting) {
-                    return checkAugmenting((CompositeNode) input);
-                } else {
-                    return checkLocal((CompositeNode) input);
+        @Deprecated
+        public boolean isAcceptable(final Node<?> input) {
+            return checkAgainstSchema(schema, input);
+        }
+
+        private static boolean checkAgainstSchema(final ChoiceCaseNode schema, final Node<?> node) {
+            if (node instanceof CompositeNode) {
+                CompositeNode input = (CompositeNode) node;
+                for (Node<?> childNode : input.getValue()) {
+                    QName child = childNode.getNodeType();
+                    if (schema.getDataChildByName(child) != null) {
+                        return true;
+                    }
                 }
             }
             return false;
         }
 
-        private boolean checkLocal(CompositeNode input) {
-            QName parent = input.getNodeType();
-            for (Node<?> childNode : input.getChildren()) {
-                QName child = childNode.getNodeType();
-                if (!Objects.equals(parent.getNamespace(), child.getNamespace())
-                        || !Objects.equals(parent.getRevision(), child.getRevision())) {
-                    continue;
-                }
-                if (validNames.contains(child.getLocalName())) {
-                    return true;
-                }
+        @Override
+        public Class<?> getDataType() {
+            return dataType;
+        }
+
+        public void adaptForPath(final InstanceIdentifier<?> augTarget, final ChoiceCaseNode choiceCaseNode) {
+            synchronized (instantiatedLocations) {
+                instantiatedLocations.put(augTarget, choiceCaseNode);
             }
-            return false;
         }
 
-        private boolean checkAugmenting(CompositeNode input) {
-            for (Node<?> child : input.getChildren()) {
-                if (validQNames.contains(child.getNodeType())) {
-                    return true;
-                }
+        public boolean isAcceptable(final InstanceIdentifier path, final CompositeNode input) {
+            ChoiceCaseNode instantiatedSchema = null;
+            synchronized (instantiatedLocations) {
+                instantiatedSchema = instantiatedLocations.get(path);
             }
-            return false;
+            if (instantiatedSchema == null) {
+                return false;
+            }
+            return checkAgainstSchema(instantiatedSchema, input);
         }
     }
 
-    private static class ChoiceCodecImpl<T> implements ChoiceCodec<T> {
+    private static class PublicChoiceCodecImpl<T> implements ChoiceCodec<T>,
+            Delegator<BindingCodec<Map<QName, Object>, Object>> {
 
         private final BindingCodec<Map<QName, Object>, Object> delegate;
 
         @SuppressWarnings("rawtypes")
-        private final Map<Class, ChoiceCaseCodecImpl<?>> cases = Collections.synchronizedMap(new WeakHashMap<Class, ChoiceCaseCodecImpl<?>>());
-
-        private final CaseCompositeNodeMapFacade CompositeToCase;
+        private final Map<Class, ChoiceCaseCodecImpl<?>> cases = Collections
+                .synchronizedMap(new WeakHashMap<Class, ChoiceCaseCodecImpl<?>>());
 
-        public ChoiceCodecImpl(BindingCodec<Map<QName, Object>, Object> delegate) {
+        public PublicChoiceCodecImpl(final BindingCodec<Map<QName, Object>, Object> delegate) {
             this.delegate = delegate;
-            this.CompositeToCase = new CaseCompositeNodeMapFacade(cases);
         }
 
         @Override
-        public ValueWithQName<T> deserialize(Node<?> input) {
+        public ValueWithQName<T> deserialize(final Node<?> input) {
             throw new UnsupportedOperationException("Direct invocation of this codec is not allowed.");
         }
 
         @Override
-        public Node<?> serialize(ValueWithQName<T> input) {
+        public ValueWithQName<T> deserialize(final Node<?> input, final InstanceIdentifier<?> bindingIdentifier) {
             throw new UnsupportedOperationException("Direct invocation of this codec is not allowed.");
         }
 
-        public CaseCompositeNodeMapFacade getCompositeToCase() {
-            return CompositeToCase;
-        }
-
-        public Map<Class, ChoiceCaseCodecImpl<?>> getCases() {
-            return cases;
+        @Override
+        public Node<?> serialize(final ValueWithQName<T> input) {
+            throw new UnsupportedOperationException("Direct invocation of this codec is not allowed.");
         }
 
+        @Override
         public BindingCodec<Map<QName, Object>, Object> getDelegate() {
             return delegate;
         }
 
     }
 
-    @SuppressWarnings("rawtypes")
-    private class CaseClassMapFacade extends MapFacadeBase {
+    class DispatchChoiceCodecImpl extends LocationAwareDispatchCodec<ChoiceCaseCodecImpl<?>> {
+        private final Class<?> choiceType;
+        private final QName choiceName;
 
-        @Override
-        public Set<Entry<Class, BindingCodec<Object, Object>>> entrySet() {
-            return Collections.emptySet();
+        private DispatchChoiceCodecImpl(final Class<?> type) {
+            choiceType = type;
+            choiceName = BindingReflections.findQName(type);
         }
 
         @Override
-        public BindingCodec get(Object key) {
-            if (key instanceof Class) {
-                Class cls = (Class) key;
-                // bindingClassEncountered(cls);
-                ChoiceCaseCodecImpl caseCodec = getCaseCodecFor(cls);
-                return caseCodec.getDelegate();
-            }
-            return null;
-        }
-    }
-
-    @SuppressWarnings("rawtypes")
-    private static class CaseCompositeNodeMapFacade extends MapFacadeBase<CompositeNode> {
+        public Object deserialize(final Object input, @SuppressWarnings("rawtypes") final InstanceIdentifier path) {
+            adaptForPath(path);
 
-        final Map<Class, ChoiceCaseCodecImpl<?>> choiceCases;
-
-        public CaseCompositeNodeMapFacade(Map<Class, ChoiceCaseCodecImpl<?>> choiceCases) {
-            this.choiceCases = choiceCases;
-        }
-
-        @Override
-        public BindingCodec get(Object key) {
-            if (!(key instanceof CompositeNode)) {
-                return null;
-            }
-            for (Entry<Class, ChoiceCaseCodecImpl<?>> entry : choiceCases.entrySet()) {
-                ChoiceCaseCodecImpl<?> codec = entry.getValue();
-                if (codec.isAcceptable((CompositeNode) key)) {
-                    return codec.getDelegate();
+            if (input instanceof CompositeNode) {
+                List<Entry<Class, ChoiceCaseCodecImpl<?>>> codecs = new ArrayList<>(getImplementations().entrySet());
+                for (Entry<Class, ChoiceCaseCodecImpl<?>> codec : codecs) {
+                    ChoiceCaseCodecImpl<?> caseCodec = codec.getValue();
+                    if (caseCodec.isAcceptable(path, (CompositeNode) input)) {
+                        ValueWithQName<?> value = caseCodec.deserialize((CompositeNode) input, path);
+                        if (value != null) {
+                            return value.getValue();
+                        }
+                        return null;
+                    }
                 }
             }
             return null;
         }
 
-    }
-
-    /**
-     * This map is used as only facade for {@link org.opendaylight.yangtools.yang.binding.BindingCodec} in different
-     * classloaders to retrieve codec dynamicly based on provided key.
-     *
-     * @param <T>
-     *            Key type
-     */
-    @SuppressWarnings("rawtypes")
-    private static abstract class MapFacadeBase<T> implements Map<T, BindingCodec<?, ?>> {
-
-        @Override
-        public boolean containsKey(Object key) {
-            return get(key) != null;
-        }
-
-        @Override
-        public void clear() {
-            throw notModifiable();
-        }
-
-        @Override
-        public boolean equals(Object obj) {
-            return super.equals(obj);
-        }
-
-        @Override
-        public BindingCodec remove(Object key) {
-            return null;
-        }
-
+        @SuppressWarnings("unchecked")
         @Override
-        public int size() {
-            return 0;
+        public Object serialize(final Object input) {
+            Preconditions.checkArgument(input instanceof Map.Entry<?, ?>, "Input must be QName, Value");
+            @SuppressWarnings("rawtypes")
+            Object inputValue = ((Map.Entry) input).getValue();
+            Preconditions.checkArgument(inputValue instanceof DataObject);
+            Class<? extends DataContainer> inputType = ((DataObject) inputValue).getImplementedInterface();
+            ChoiceCaseCodecImpl<?> codec = tryToLoadImplementation(inputType);
+            Preconditions.checkState(codec != null, "Unable to get codec for %s", inputType);
+            if(isAugmenting(codec.getSchema())) {
+                // If choice is augmenting we use QName which defined this augmentation
+                return codec.getDelegate().serialize(new ValueWithQName<>(codec.getSchema().getQName(), inputValue));
+            }
+            return codec.getDelegate().serialize(input);
         }
 
-        @Override
-        public Collection<BindingCodec<?, ?>> values() {
-            return Collections.emptySet();
+        private boolean isAugmenting(final ChoiceCaseNode schema) {
+            if(schema.isAugmenting()) {
+                return true;
+            }
+            QName parentQName = Iterables.get(schema.getPath().getPathTowardsRoot(), 2); // choice QName
+            if(!parentQName.getNamespace().equals(schema.getQName().getNamespace())) {
+                return true;
+            }
+            return false;
         }
 
-        private UnsupportedOperationException notModifiable() {
-            return new UnsupportedOperationException("Not externally modifiable.");
+        @SuppressWarnings("rawtypes")
+        protected Optional<ChoiceCaseCodecImpl> tryToLoadImplementation(final Type potential) {
+            try {
+                @SuppressWarnings("unchecked")
+                Class<? extends DataContainer> clazz = (Class<? extends DataContainer>) classLoadingStrategy
+                        .loadClass(potential);
+                ChoiceCaseCodecImpl codec = tryToLoadImplementation(clazz);
+                addImplementation(codec);
+                return Optional.of(codec);
+            } catch (ClassNotFoundException e) {
+                LOG.warn("Failed to find class for choice {}", potential, e);
+            }
+            return Optional.absent();
         }
 
         @Override
-        public BindingCodec<Map<QName, Object>, Object> put(T key, BindingCodec<?, ?> value) {
-            throw notModifiable();
+        protected ChoiceCaseCodecImpl<?> tryToLoadImplementation(final Class<? extends DataContainer> inputType) {
+            ChoiceCaseCodecImpl<?> codec = getCaseCodecFor(inputType);
+            addImplementation(codec);
+            return codec;
         }
 
         @Override
-        public void putAll(Map<? extends T, ? extends BindingCodec<?, ?>> m) {
-            throw notModifiable();
+        protected void tryToLoadImplementations() {
+            Type type = referencedType(choiceType);
+            Collection<Type> potentialCases;
+            synchronized (choiceToCases) {
+                potentialCases = choiceToCases.get(type);
+            }
+            for (Type potential : potentialCases) {
+                try {
+                    tryToLoadImplementation(potential);
+                } catch (CodeGenerationException e) {
+                    LOG.warn("Failed to proactively generate choice code for {}", type, e);
+                }
+            }
         }
 
         @Override
-        public int hashCode() {
-            return super.hashCode();
+        protected void adaptForPathImpl(final InstanceIdentifier<?> augTarget, final DataNodeContainer ctxNode) {
+            Optional<ChoiceNode> newChoice = findInstantiatedChoice(ctxNode, choiceName);
+            tryToLoadImplementations();
+            if (newChoice.isPresent()) {
+                for (@SuppressWarnings("rawtypes")
+                Entry<Class, ChoiceCaseCodecImpl<?>> codec : getImplementations().entrySet()) {
+                    ChoiceCaseCodecImpl<?> caseCodec = codec.getValue();
+                    Optional<ChoiceCaseNode> instantiatedSchema = findInstantiatedCase(newChoice.get(),
+                            caseCodec.getSchema());
+                    if (instantiatedSchema.isPresent()) {
+                        caseCodec.adaptForPath(augTarget, instantiatedSchema.get());
+                    }
+                }
+            }
         }
 
-        @Override
-        public boolean isEmpty() {
-            return true;
-        }
+        private Optional<ChoiceNode> findInstantiatedChoice(final DataNodeContainer ctxNode, final QName choiceName) {
+            DataSchemaNode potential = ctxNode.getDataChildByName(choiceName);
+            if (potential == null) {
+                potential = ctxNode.getDataChildByName(choiceName.getLocalName());
+            }
 
-        @Override
-        public Set<T> keySet() {
-            return Collections.emptySet();
-        }
+            if (potential instanceof ChoiceNode) {
+                return Optional.of((ChoiceNode) potential);
+            }
 
-        @Override
-        public Set<Entry<T, BindingCodec<?, ?>>> entrySet() {
-            return Collections.emptySet();
+            return Optional.absent();
         }
 
-        @Override
-        public boolean containsValue(Object value) {
-            return false;
+        private Optional<ChoiceCaseNode> findInstantiatedCase(final ChoiceNode newChoice, final ChoiceCaseNode schema) {
+            ChoiceCaseNode potential = newChoice.getCaseNodeByName(schema.getQName());
+            if (potential != null) {
+                return Optional.of(potential);
+            }
+            // FIXME: Probably requires more extensive check
+            // e.g. we have one choice and two augmentations from different
+            // modules using same local name
+            // but different namespace / contents
+            return Optional.fromNullable(newChoice.getCaseNodeByName(schema.getQName().getLocalName()));
         }
     }
 
     @SuppressWarnings({ "rawtypes", "unchecked" })
-    private class AugmentableCompositeCodec implements BindingCodec {
+    class AugmentableDispatchCodec extends LocationAwareDispatchCodec<AugmentationCodecWrapper> {
 
         private final Class augmentableType;
 
-        Map<Class, AugmentationCodec<?>> localAugmentationCodecs = Collections.synchronizedMap(new WeakHashMap<Class, AugmentationCodec<?>>());
-
-        public AugmentableCompositeCodec(Class type) {
+        public AugmentableDispatchCodec(final Class type) {
             Preconditions.checkArgument(Augmentable.class.isAssignableFrom(type));
             augmentableType = type;
         }
 
         @Override
-        public Object serialize(Object input) {
+        // TODO deprecate use without iid
+        public Object serialize(final Object input) {
             if (input instanceof Augmentable<?>) {
-
                 Map<Class, Augmentation> augmentations = getAugmentations(input);
                 return serializeImpl(augmentations);
             }
             return null;
         }
 
-        private Map<Class, Augmentation> getAugmentations(Object input) {
+        private Map<Class, Augmentation> getAugmentations(final Object input) {
             Field augmentationField;
             try {
                 augmentationField = input.getClass().getDeclaredField("augmentation");
@@ -826,7 +973,8 @@ public class LazyGeneratedCodecRegistry implements //
             return Collections.emptyMap();
         }
 
-        private List serializeImpl(Map<Class, Augmentation> input) {
+        @SuppressWarnings("deprecation")
+        private List serializeImpl(final Map<Class, Augmentation> input) {
             List ret = new ArrayList<>();
             for (Entry<Class, Augmentation> entry : input.entrySet()) {
                 AugmentationCodec codec = getCodecForAugmentation(entry.getKey());
@@ -836,64 +984,161 @@ public class LazyGeneratedCodecRegistry implements //
             return ret;
         }
 
-        public synchronized <T extends Augmentation<?>> void addAugmentationCodec(Class<T> augmentationClass,
-                AugmentationCodec<T> value) {
-            localAugmentationCodecs.put(augmentationClass, value);
-        }
-
         @Override
-        public Map<Class, Augmentation> deserialize(Object input) {
+        public Map<Class, Augmentation> deserialize(final Object input, final InstanceIdentifier path) {
+            adaptForPath(path);
             Map<Class, Augmentation> ret = new HashMap<>();
+
             if (input instanceof CompositeNode) {
-                List<Entry<Class, AugmentationCodec<?>>> codecs = new ArrayList<>(localAugmentationCodecs.entrySet());
-                for (Entry<Class, AugmentationCodec<?>> codec : codecs) {
-                    ValueWithQName<?> value = codec.getValue().deserialize((CompositeNode) input);
-                    if (value != null && value.getValue() != null) {
-                        ret.put(codec.getKey(), (Augmentation) value.getValue());
+                List<Entry<Class, AugmentationCodecWrapper>> codecs = new ArrayList<>(getImplementations().entrySet());
+                for (Entry<Class, AugmentationCodecWrapper> codec : codecs) {
+                    AugmentationCodec<?> ac = codec.getValue();
+                    if (ac.isAcceptable(path)) {
+                        // We add Augmentation Identifier to path, in order to
+                        // correctly identify children.
+                        InstanceIdentifier augmentPath = path.builder().augmentation(codec.getKey()).build();
+                        ValueWithQName<?> value = codec.getValue().deserialize((CompositeNode) input, augmentPath);
+                        if (value != null && value.getValue() != null) {
+                            ret.put(codec.getKey(), (Augmentation) value.getValue());
+                        }
                     }
                 }
             }
             return ret;
         }
 
-        public Class getAugmentableType() {
-            return augmentableType;
+        protected Optional<AugmentationCodecWrapper> tryToLoadImplementation(final Type potential) {
+            try {
+                Class<? extends Augmentation<?>> clazz = (Class<? extends Augmentation<?>>) classLoadingStrategy
+                        .loadClass(potential);
+                return Optional.of(tryToLoadImplementation(clazz));
+            } catch (ClassNotFoundException e) {
+                LOG.warn("Failed to find class for augmentation of {}", potential, e);
+            }
+            return Optional.absent();
         }
-    }
-
-    @SuppressWarnings({ "rawtypes", "unchecked" })
-    private static class LateMixinCodec implements BindingCodec, Delegator<BindingCodec> {
 
-        private BindingCodec delegate;
+        @Override
+        protected AugmentationCodecWrapper tryToLoadImplementation(final Class inputType) {
+            AugmentationCodecWrapper<? extends Augmentation<?>> potentialImpl = getCodecForAugmentation(inputType);
+            addImplementation(potentialImpl);
+            return potentialImpl;
+        }
 
         @Override
-        public BindingCodec getDelegate() {
-            if (delegate == null) {
-                throw new IllegalStateException("Codec not initialized yet.");
+        protected void tryToLoadImplementations() {
+            Type type = referencedType(augmentableType);
+            Collection<Type> potentialAugmentations;
+            synchronized (augmentableToAugmentations) {
+                potentialAugmentations = new ArrayList(augmentableToAugmentations.get(type));
+            }
+            for (Type potential : potentialAugmentations) {
+                try {
+                    tryToLoadImplementation(potential);
+                } catch (CodeGenerationException e) {
+                    LOG.warn("Failed to proactively generate augment code for {}", type, e);
+                }
             }
-            return delegate;
         }
 
         @Override
-        public Object deserialize(Object input) {
-            return getDelegate().deserialize(input);
+        protected void adaptForPathImpl(final InstanceIdentifier<?> augTarget, final DataNodeContainer ctxNode) {
+            if (ctxNode instanceof AugmentationTarget) {
+                Set<AugmentationSchema> availableAugmentations = ((AugmentationTarget) ctxNode)
+                        .getAvailableAugmentations();
+                if (!availableAugmentations.isEmpty()) {
+                    updateAugmentationMapping(augTarget, availableAugmentations);
+                }
+            }
         }
 
-        @Override
-        public Object serialize(Object input) {
-            return getDelegate().serialize(input);
+        /**
+         *
+         * Adapts augmentation codec for specific provider location (target)
+         *
+         * Since augmentation are not forward-referencing and may be discovered
+         * during runtime, we need to adapt {@link AugmentableDispatchCodec},
+         * {@link AugmentationCodecWrapper} and {@link InstanceIdentifierCodec}
+         * for this newly discovered location where augmentation may be used.
+         *
+         * Adaptation consists of:
+         * <ol>
+         * <li> scan of available (valid) augmentations for
+         * current location
+         * <li>lookup for Java classes derived from this augmentations
+         * <li>generation of missing codecs
+         * <li>updating Augmentation codecs to work with new location
+         * <li>updating Instance Identifier to work with new location
+         *
+         */
+        private void updateAugmentationMapping(final InstanceIdentifier<?> augTarget,
+                final Set<AugmentationSchema> availableAugmentations) {
+            for (AugmentationSchema aug : availableAugmentations) {
+
+                Type potentialType = getTypeForAugmentation(aug);
+                if (potentialType != null) {
+                    Optional<AugmentationCodecWrapper> potentialImpl = tryToLoadImplementation(potentialType);
+                    if (potentialImpl.isPresent()) {
+                        potentialImpl.get().addApplicableFor(augTarget, aug);
+                        Class augType = potentialImpl.get().getDataType();
+                        InstanceIdentifier augPath = augTarget.augmentation(augType);
+                        try {
+
+                            org.opendaylight.yangtools.yang.data.api.InstanceIdentifier domPath = getInstanceIdentifierCodec().serialize(augPath);
+                            if(domPath == null) {
+                                LOG.error("Unable to serialize instance identifier for {}",augPath);
+                            }
+                        } catch (Exception e) {
+                            LOG.error("Unable to serialize instance identifiers for {}",augPath,e);
+                        }
+
+                    }
+                } else {
+                    // Omits warning for empty augmentations since they are not represented in data
+                    if(!aug.getChildNodes().isEmpty()) {
+                        LOG.warn("Could not find generated type for augmentation {} with children {}", aug,
+                            aug.getChildNodes());
+                    }
+                }
+            }
+        }
+
+
+
+        private Type getTypeForAugmentation(final AugmentationSchema aug) {
+            Optional<AugmentationSchema> currentAug = Optional.of(aug);
+            while (currentAug.isPresent()) {
+                Type potentialType = typeToAugment.inverse().get(currentAug.get());
+                if (potentialType != null) {
+                    return potentialType;
+                }
+                currentAug = currentAug.get().getOriginalDefinition();
+            }
+            return null;
         }
     }
 
+    @SuppressWarnings("rawtypes")
     private static class AugmentationCodecWrapper<T extends Augmentation<?>> implements AugmentationCodec<T>,
-            Delegator<BindingCodec> {
+            Delegator<BindingCodec>, LocationAwareBindingCodec<Node<?>, ValueWithQName<T>> {
 
-        private BindingCodec delegate;
-        private QName augmentationQName;
+        private final BindingCodec delegate;
+        private final QName augmentationQName;
+        private final Multimap<InstanceIdentifier<?>, QName> validAugmentationTargets;
+        private final Class<?> augmentationType;
 
-        public AugmentationCodecWrapper(BindingCodec<Map<QName, Object>, Object> rawCodec) {
+        public AugmentationCodecWrapper(final BindingCodec<Map<QName, Object>, Object> rawCodec, final Class<?> dataType) {
             this.delegate = rawCodec;
+            this.augmentationType = dataType;
             this.augmentationQName = BindingReflections.findQName(rawCodec.getClass());
+            this.validAugmentationTargets = Multimaps.synchronizedSetMultimap(HashMultimap
+                    .<InstanceIdentifier<?>, QName> create());
+        }
+
+        public void addApplicableFor(final InstanceIdentifier<?> path, final AugmentationSchema aug) {
+            for (DataSchemaNode child : aug.getChildNodes()) {
+                validAugmentationTargets.put(path, child.getQName());
+            }
         }
 
         @Override
@@ -902,7 +1147,7 @@ public class LazyGeneratedCodecRegistry implements //
         }
 
         @Override
-        public CompositeNode serialize(ValueWithQName<T> input) {
+        public CompositeNode serialize(final ValueWithQName<T> input) {
             @SuppressWarnings("unchecked")
             List<Map<QName, Object>> rawValues = (List<Map<QName, Object>>) getDelegate().serialize(input);
             List<Node<?>> serialized = new ArrayList<>(rawValues.size());
@@ -914,8 +1159,18 @@ public class LazyGeneratedCodecRegistry implements //
 
         @Override
         @SuppressWarnings("unchecked")
-        public ValueWithQName<T> deserialize(Node<?> input) {
-            Object rawCodecValue = getDelegate().deserialize((Map<QName, Object>) input);
+        public ValueWithQName<T> deserialize(final Node<?> input) {
+            Object rawCodecValue = getDelegate().deserialize(input);
+            return new ValueWithQName<T>(input.getNodeType(), (T) rawCodecValue);
+        }
+
+        @Override
+        @SuppressWarnings("unchecked")
+        public ValueWithQName<T> deserialize(final Node<?> input, final InstanceIdentifier<?> bindingIdentifier) {
+            // if (!isAcceptable(bindingIdentifier)) {
+            // return null;
+            // }
+            Object rawCodecValue = getDelegate().deserialize(input, bindingIdentifier);
             return new ValueWithQName<T>(input.getNodeType(), (T) rawCodecValue);
         }
 
@@ -923,18 +1178,37 @@ public class LazyGeneratedCodecRegistry implements //
         public QName getAugmentationQName() {
             return augmentationQName;
         }
+
+        @Override
+        public boolean isAcceptable(final InstanceIdentifier<?> path) {
+            if (path == null) {
+                return false;
+            }
+            return validAugmentationTargets.containsKey(path);
+        }
+
+        @Override
+        public boolean isApplicable(final InstanceIdentifier location) {
+            return isAcceptable(location);
+        }
+
+        @Override
+        public Class<?> getDataType() {
+            return augmentationType;
+        }
     }
 
+    @SuppressWarnings("rawtypes")
     private class IdentityCompositeCodec implements IdentityCodec {
 
         @Override
-        public Object deserialize(Object input) {
+        public Object deserialize(final Object input) {
             Preconditions.checkArgument(input instanceof QName);
             return deserialize((QName) input);
         }
 
         @Override
-        public Class<?> deserialize(QName input) {
+        public Class<?> deserialize(final QName input) {
             Type type = qnamesToIdentityMap.get(input);
             if (type == null) {
                 return null;
@@ -942,13 +1216,49 @@ public class LazyGeneratedCodecRegistry implements //
             ReferencedTypeImpl typeref = new ReferencedTypeImpl(type.getPackageName(), type.getName());
             WeakReference<Class> softref = typeToClass.get(typeref);
             if (softref == null) {
+
+                try {
+                    Class<?> cls = classLoadingStrategy.loadClass(typeref.getFullyQualifiedName());
+                    if (cls != null) {
+                        serialize(cls);
+                        return cls;
+                    }
+                } catch (Exception e) {
+                    LOG.warn("Identity {} was not deserialized, because of missing class {}", input,
+                            typeref.getFullyQualifiedName());
+                }
                 return null;
             }
             return softref.get();
         }
 
         @Override
-        public QName serialize(Class input) {
+        public Object deserialize(final Object input, final InstanceIdentifier bindingIdentifier) {
+            Type type = qnamesToIdentityMap.get(input);
+            if (type == null) {
+                return null;
+            }
+            ReferencedTypeImpl typeref = new ReferencedTypeImpl(type.getPackageName(), type.getName());
+            WeakReference<Class> softref = typeToClass.get(typeref);
+            if (softref == null) {
+
+                try {
+                    Class<?> cls = classLoadingStrategy.loadClass(typeref.getFullyQualifiedName());
+                    if (cls != null) {
+                        serialize(cls);
+                        return cls;
+                    }
+                } catch (Exception e) {
+                    LOG.warn("Identity {} was not deserialized, because of missing class {}", input,
+                            typeref.getFullyQualifiedName());
+                }
+                return null;
+            }
+            return softref.get();
+        }
+
+        @Override
+        public QName serialize(final Class input) {
             Preconditions.checkArgument(BaseIdentity.class.isAssignableFrom(input));
             bindingClassEncountered(input);
             QName qname = identityQNames.get(input);
@@ -964,31 +1274,14 @@ public class LazyGeneratedCodecRegistry implements //
         }
 
         @Override
-        public Object serialize(Object input) {
+        public Object serialize(final Object input) {
             Preconditions.checkArgument(input instanceof Class);
             return serialize((Class) input);
         }
+
     }
 
-    public boolean isCodecAvailable(Class<? extends DataContainer> cls) {
-        if (containerCodecs.containsKey(cls)) {
-            return true;
-        }
-        if (identifierCodecs.containsKey(cls)) {
-            return true;
-        }
-        if (choiceCodecs.containsKey(cls)) {
-            return true;
-        }
-        if (caseCodecs.containsKey(cls)) {
-            return true;
-        }
-        if (augmentableCodecs.containsKey(cls)) {
-            return true;
-        }
-        if (augmentationCodecs.containsKey(cls)) {
-            return true;
-        }
-        return false;
+    private static final Type referencedType(final Class<?> augmentableType) {
+        return new ReferencedTypeImpl(augmentableType.getPackage().getName(), augmentableType.getSimpleName());
     }
 }