Merge "Simplify method isMutualExclusive in Subnet. Remove redundant 'if' statements."
[controller.git] / opendaylight / md-sal / sal-binding-broker / src / main / java / org / opendaylight / controller / sal / binding / dom / serializer / impl / LazyGeneratedCodecRegistry.java
index a55d4b9276731f3ede54036272d07ddf762f3cab..cabb1bc4e2c5359582bc95e1e30514f557801b64 100644 (file)
@@ -1,7 +1,9 @@
 package org.opendaylight.controller.sal.binding.dom.serializer.impl;
 
+import java.awt.CompositeContext;
 import java.lang.ref.WeakReference;
 import java.lang.reflect.Field;
+import java.lang.reflect.ParameterizedType;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -12,11 +14,13 @@ import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Objects;
+import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.Set;
 import java.util.WeakHashMap;
 
+import org.apache.commons.lang3.text.translate.AggregateTranslator;
 import org.opendaylight.controller.sal.binding.dom.serializer.api.AugmentationCodec;
 import org.opendaylight.controller.sal.binding.dom.serializer.api.ChoiceCaseCodec;
 import org.opendaylight.controller.sal.binding.dom.serializer.api.ChoiceCodec;
@@ -24,8 +28,10 @@ import org.opendaylight.controller.sal.binding.dom.serializer.api.CodecRegistry;
 import org.opendaylight.controller.sal.binding.dom.serializer.api.DataContainerCodec;
 import org.opendaylight.controller.sal.binding.dom.serializer.api.DomCodec;
 import org.opendaylight.controller.sal.binding.dom.serializer.api.IdentifierCodec;
+import org.opendaylight.controller.sal.binding.dom.serializer.api.IdentitityCodec;
 import org.opendaylight.controller.sal.binding.dom.serializer.api.InstanceIdentifierCodec;
 import org.opendaylight.controller.sal.binding.dom.serializer.api.ValueWithQName;
+import org.opendaylight.controller.sal.binding.impl.util.ClassLoaderUtils;
 import org.opendaylight.controller.sal.core.api.model.SchemaServiceListener;
 import org.opendaylight.yangtools.binding.generator.util.ReferencedTypeImpl;
 import org.opendaylight.yangtools.binding.generator.util.Types;
@@ -33,6 +39,7 @@ import org.opendaylight.yangtools.concepts.Delegator;
 import org.opendaylight.yangtools.concepts.Identifiable;
 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;
@@ -40,6 +47,9 @@ import org.opendaylight.yangtools.yang.binding.Identifier;
 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.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;
@@ -48,6 +58,7 @@ import org.opendaylight.yangtools.yang.model.api.LeafListSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.LeafSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.SchemaContext;
 import org.opendaylight.yangtools.yang.model.api.SchemaPath;
+import org.opendaylight.yangtools.yang.model.api.UsesNode;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -57,10 +68,14 @@ import static org.opendaylight.controller.sal.binding.dom.serializer.impl.Interm
 import org.opendaylight.yangtools.sal.binding.generator.impl.ModuleContext;
 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.model.api.Module;
 import org.opendaylight.yangtools.yang.model.util.SchemaContextUtil;
 
+import com.google.common.collect.FluentIterable;
+import com.google.common.util.concurrent.CycleDetectingLockFactory.WithExplicitOrdering;
+
 public class LazyGeneratedCodecRegistry implements //
         CodecRegistry, //
         SchemaServiceListener, //
@@ -70,26 +85,31 @@ public class LazyGeneratedCodecRegistry implements //
     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;
 
     // Concrete class to codecs
-    private Map<Class<?>, DataContainerCodec<?>> containerCodecs = new WeakHashMap<>();
-    private Map<Class<?>, IdentifierCodec<?>> identifierCodecs = new WeakHashMap<>();
-    private Map<Class<?>, ChoiceCodecImpl<?>> choiceCodecs = new WeakHashMap<>();
-    private Map<Class<?>, ChoiceCaseCodecImpl<?>> caseCodecs = new WeakHashMap<>();
-    private Map<Class<?>, AugmentableCompositeCodec> augmentableCodecs = new WeakHashMap<>();
-
+    private static final Map<Class<?>, DataContainerCodec<?>> containerCodecs = new WeakHashMap<>();
+    private static final Map<Class<?>, IdentifierCodec<?>> identifierCodecs = new WeakHashMap<>();
+    private static final Map<Class<?>, ChoiceCodecImpl<?>> choiceCodecs = new WeakHashMap<>();
+    private static final Map<Class<?>, ChoiceCaseCodecImpl<?>> caseCodecs = new WeakHashMap<>();
+    private static final Map<Class<?>, AugmentableCompositeCodec> augmentableCodecs = new WeakHashMap<>();
+    private static final Map<Class<?>, AugmentationCodec<?>> augmentationCodecs = new WeakHashMap<>();
+    private static final Map<Class<?>, QName> identityQNames = new WeakHashMap<>();
+    private static final Map<QName, Type> qnamesToIdentityMap = new ConcurrentHashMap<>();
     /** Binding type to encountered classes mapping **/
     @SuppressWarnings("rawtypes")
-    Map<Type, WeakReference<Class>> typeToClass = new ConcurrentHashMap<>();
+    private static final Map<Type, WeakReference<Class>> typeToClass = new ConcurrentHashMap<>();
 
     @SuppressWarnings("rawtypes")
-    private ConcurrentMap<Type, ChoiceCaseCodecImpl> typeToCaseNodes = new ConcurrentHashMap<>();
+    private static final ConcurrentMap<Type, ChoiceCaseCodecImpl> typeToCaseCodecs = new ConcurrentHashMap<>();
 
     private CaseClassMapFacade classToCaseRawCodec = new CaseClassMapFacade();
 
-    Map<SchemaPath, GeneratedTypeBuilder> pathToType = new ConcurrentHashMap<>();
+    private static final Map<SchemaPath, GeneratedTypeBuilder> pathToType = new ConcurrentHashMap<>();
+    private static final Map<List<QName>, Type> pathToInstantiatedType = new ConcurrentHashMap<>();
+    private static final Map<Type, QName> typeToQname = new ConcurrentHashMap<>();
 
     private SchemaContext currentSchema;
 
@@ -108,24 +128,83 @@ public class LazyGeneratedCodecRegistry implements //
 
     @Override
     public <T extends Augmentation<?>> AugmentationCodec<T> getCodecForAugmentation(Class<T> object) {
-        // TODO Auto-generated method stub
-        return null;
+        AugmentationCodec<T> codec = null;
+        @SuppressWarnings("rawtypes")
+        AugmentationCodec potentialCodec = augmentationCodecs.get(object);
+        if (potentialCodec != null) {
+            codec = potentialCodec;
+        } else
+            try {
+                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);
+        }
+        return codec;
+    }
+
+    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;
+        }
     }
 
     @Override
     public Class<?> getClassForPath(List<QName> names) {
         DataSchemaNode node = getSchemaNode(names);
         SchemaPath path = node.getPath();
-        GeneratedTypeBuilder type = pathToType.get(path);
-        ReferencedTypeImpl typeref = new ReferencedTypeImpl(type.getPackageName(), type.getName());
+        Type type = pathToType.get(path);
+        if (type != null) {
+            type = new ReferencedTypeImpl(type.getPackageName(), type.getName());
+        } else {
+            type = pathToInstantiatedType.get(names);
+        }
         @SuppressWarnings("rawtypes")
-        WeakReference<Class> weakRef = typeToClass.get(typeref);
-        if(weakRef == null) {
-            LOG.error("Could not find loaded class for path: {} and type: {}",path,typeref.getFullyQualifiedName());
+        WeakReference<Class> weakRef = typeToClass.get(type);
+        if (weakRef == null) {
+            LOG.error("Could not find loaded class for path: {} and type: {}", path, type.getFullyQualifiedName());
         }
         return weakRef.get();
     }
 
+    @Override
+    public void putPathToClass(List<QName> names, Class<?> cls) {
+        Type reference = Types.typeForClass(cls);
+        pathToInstantiatedType.put(names, reference);
+        bindingClassEncountered(cls);
+    }
+
     @Override
     public IdentifierCodec<?> getKeyCodecForPath(List<QName> names) {
         @SuppressWarnings("unchecked")
@@ -147,30 +226,31 @@ public class LazyGeneratedCodecRegistry implements //
         return newWrapper;
     }
 
-    @Override
     @SuppressWarnings("rawtypes")
     public void bindingClassEncountered(Class cls) {
-        
+
         ConcreteType typeRef = Types.typeForClass(cls);
-        if(typeToClass.containsKey(typeRef)) {
+        if (typeToClass.containsKey(typeRef)) {
             return;
         }
-        LOG.info("Binding Class {} encountered.",cls);
+        LOG.trace("Binding Class {} encountered.", cls);
         WeakReference<Class> weakRef = new WeakReference<>(cls);
         typeToClass.put(typeRef, weakRef);
-        if(DataObject.class.isAssignableFrom(cls)) {
-            @SuppressWarnings({"unchecked","unused"})
+        if (Augmentation.class.isAssignableFrom(cls)) {
+
+        } else if (DataObject.class.isAssignableFrom(cls)) {
+            @SuppressWarnings({ "unchecked", "unused" })
             Object cdc = getCodecForDataObject((Class<? extends DataObject>) cls);
         }
     }
-    
+
     @Override
     public void onClassProcessed(Class<?> cls) {
         ConcreteType typeRef = Types.typeForClass(cls);
-        if(typeToClass.containsKey(typeRef)) {
+        if (typeToClass.containsKey(typeRef)) {
             return;
         }
-        LOG.info("Binding Class {} encountered.",cls);
+        LOG.trace("Binding Class {} encountered.", cls);
         WeakReference<Class> weakRef = new WeakReference<>((Class) cls);
         typeToClass.put(typeRef, weakRef);
     }
@@ -248,9 +328,21 @@ public class LazyGeneratedCodecRegistry implements //
         return newWrapper;
     }
 
+    @Override
+    public IdentitityCodec<?> getIdentityCodec() {
+        return identityRefCodec;
+    }
+
+    @Override
+    public <T extends BaseIdentity> IdentitityCodec<T> getCodecForIdentity(Class<T> codec) {
+        bindingClassEncountered(codec);
+        return identityRefCodec;
+    }
+
     @Override
     public void onCodecCreated(Class<?> cls) {
         CodecMapping.setIdentifierCodec(cls, instanceIdentifierCodec);
+        CodecMapping.setIdentityRefCodec(cls, identityRefCodec);
     }
 
     @Override
@@ -276,10 +368,12 @@ public class LazyGeneratedCodecRegistry implements //
             return potential;
         }
         ConcreteType typeref = Types.typeForClass(caseClass);
-        ChoiceCaseCodecImpl caseCodec = typeToCaseNodes.get(typeref);
+        ChoiceCaseCodecImpl caseCodec = typeToCaseCodecs.get(typeref);
 
+        checkState(caseCodec != null, "Case Codec was not created proactivelly for %s", caseClass.getName());
+        checkState(caseCodec.getSchema() != null, "Case schema is not available for %s", caseClass.getName());
         @SuppressWarnings("unchecked")
-        Class<? extends BindingCodec> newCodec = generator.caseCodecFor(caseClass, caseCodec.schema);
+        Class<? extends BindingCodec> newCodec = generator.caseCodecFor(caseClass, caseCodec.getSchema());
         BindingCodec newInstance = newInstanceOf(newCodec);
         caseCodec.setDelegate(newInstance);
         caseCodecs.put(caseClass, caseCodec);
@@ -294,7 +388,12 @@ public class LazyGeneratedCodecRegistry implements //
 
     public void onModuleContextAdded(SchemaContext schemaContext, Module module, ModuleContext context) {
         pathToType.putAll(context.getChildNodes());
-
+        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);
     }
 
@@ -302,16 +401,22 @@ public class LazyGeneratedCodecRegistry implements //
         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);
-            typeToCaseNodes.putIfAbsent(typeref, value);
+            typeToCaseCodecs.putIfAbsent(typeref, value);
         }
     }
 
@@ -323,7 +428,7 @@ public class LazyGeneratedCodecRegistry implements //
     @SuppressWarnings({ "unchecked", "rawtypes" })
     @Override
     public void onChoiceCodecCreated(Class<?> choiceClass,
-            Class<? extends BindingCodec<Map<QName, Object>, Object>> choiceCodec) {
+            Class<? extends BindingCodec<Map<QName, Object>, Object>> choiceCodec, ChoiceNode schema) {
         ChoiceCodec<?> oldCodec = choiceCodecs.get(choiceClass);
         checkState(oldCodec == null);
         BindingCodec<Map<QName, Object>, Object> delegate = newInstanceOf(choiceCodec);
@@ -332,6 +437,28 @@ public class LazyGeneratedCodecRegistry implements //
         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);
+                }
+            }
+        }
+
     }
 
     @Override
@@ -344,8 +471,7 @@ public class LazyGeneratedCodecRegistry implements //
     }
 
     @Override
-    public void onDataContainerCodecCreated(Class<?> dataClass,
-            Class<? extends BindingCodec<Map<QName, Object>, Object>> dataCodec) {
+    public void onDataContainerCodecCreated(Class<?> dataClass, Class<? extends BindingCodec<?, ?>> dataCodec) {
         if (Augmentable.class.isAssignableFrom(dataClass)) {
             AugmentableCompositeCodec augmentableCodec = getAugmentableCodec(dataClass);
             CodecMapping.setAugmentationCodec(dataCodec, augmentableCodec);
@@ -353,7 +479,7 @@ public class LazyGeneratedCodecRegistry implements //
 
     }
 
-    private AugmentableCompositeCodec getAugmentableCodec(Class<?> dataClass) {
+    public AugmentableCompositeCodec getAugmentableCodec(Class<?> dataClass) {
         AugmentableCompositeCodec ret = augmentableCodecs.get(dataClass);
         if (ret != null) {
             return ret;
@@ -434,15 +560,15 @@ public class LazyGeneratedCodecRegistry implements //
     @SuppressWarnings("rawtypes")
     private static class ChoiceCaseCodecImpl<T extends DataContainer> implements ChoiceCaseCodec<T>, //
             Delegator<BindingCodec> {
-        private final boolean augmenting;
+        private boolean augmenting;
         private BindingCodec delegate;
 
-        private final Set<String> validNames;
-        private final Set<QName> validQNames;
+        private Set<String> validNames;
+        private Set<QName> validQNames;
         private ChoiceCaseNode schema;
 
-        public ChoiceCaseCodecImpl(ChoiceCaseNode caseNode) {
-            this.delegate = NOT_READY_CODEC;
+        public void setSchema(ChoiceCaseNode caseNode) {
+            this.schema = schema;
             this.schema = caseNode;
             validNames = new HashSet<>();
             validQNames = new HashSet<>();
@@ -454,6 +580,15 @@ public class LazyGeneratedCodecRegistry implements //
             augmenting = caseNode.isAugmenting();
         }
 
+        public ChoiceCaseCodecImpl() {
+            this.delegate = NOT_READY_CODEC;
+        }
+
+        public ChoiceCaseCodecImpl(ChoiceCaseNode caseNode) {
+            this.delegate = NOT_READY_CODEC;
+            setSchema(caseNode);
+        }
+
         @Override
         public ValueWithQName<T> deserialize(Node<?> input) {
             throw new UnsupportedOperationException("Direct invocation of this codec is not allowed.");
@@ -478,7 +613,7 @@ public class LazyGeneratedCodecRegistry implements //
 
         @Override
         public boolean isAcceptable(Node<?> input) {
-            if (false == (input instanceof CompositeNode)) {
+            if (input instanceof CompositeNode) {
                 if (augmenting) {
                     return checkAugmenting((CompositeNode) input);
                 } else {
@@ -492,10 +627,8 @@ public class LazyGeneratedCodecRegistry implements //
             QName parent = input.getNodeType();
             for (Node<?> childNode : input.getChildren()) {
                 QName child = childNode.getNodeType();
-                if (false == Objects.equals(parent.getNamespace(), child.getNamespace())) {
-                    continue;
-                }
-                if (false == Objects.equals(parent.getRevision(), child.getRevision())) {
+                if (!Objects.equals(parent.getNamespace(), child.getNamespace())
+                        || !Objects.equals(parent.getRevision(), child.getRevision())) {
                     continue;
                 }
                 if (validNames.contains(child.getLocalName())) {
@@ -558,14 +691,14 @@ public class LazyGeneratedCodecRegistry implements //
 
         @Override
         public Set<java.util.Map.Entry<Class, BindingCodec<Object, Object>>> entrySet() {
-            return null;
+            return Collections.emptySet();
         }
 
         @Override
         public BindingCodec get(Object key) {
             if (key instanceof Class) {
                 Class cls = (Class) key;
-                //bindingClassEncountered(cls);
+                // bindingClassEncountered(cls);
                 ChoiceCaseCodecImpl caseCodec = getCaseCodecFor(cls);
                 return caseCodec.getDelegate();
             }
@@ -584,7 +717,7 @@ public class LazyGeneratedCodecRegistry implements //
 
         @Override
         public BindingCodec get(Object key) {
-            if (false == (key instanceof CompositeNode)) {
+            if (!(key instanceof CompositeNode)) {
                 return null;
             }
             for (java.util.Map.Entry<Class, ChoiceCaseCodecImpl<?>> entry : choiceCases.entrySet()) {
@@ -595,8 +728,7 @@ public class LazyGeneratedCodecRegistry implements //
             }
             return null;
         }
-        
-        
+
     }
 
     /**
@@ -636,7 +768,7 @@ public class LazyGeneratedCodecRegistry implements //
 
         @Override
         public Collection<BindingCodec<?, ?>> values() {
-            return null;
+            return Collections.emptySet();
         }
 
         private UnsupportedOperationException notModifiable() {
@@ -644,7 +776,7 @@ public class LazyGeneratedCodecRegistry implements //
         }
 
         @Override
-        public BindingCodec<Map<QName, Object>, Object> put(T key, BindingCodec<?,?> value) {
+        public BindingCodec<Map<QName, Object>, Object> put(T key, BindingCodec<?, ?> value) {
             throw notModifiable();
         }
 
@@ -660,18 +792,17 @@ public class LazyGeneratedCodecRegistry implements //
 
         @Override
         public boolean isEmpty() {
-            return false;
+            return true;
         }
 
         @Override
         public Set<T> keySet() {
-            return null;
+            return Collections.emptySet();
         }
 
         @Override
         public Set<java.util.Map.Entry<T, BindingCodec<?, ?>>> entrySet() {
-            // TODO Auto-generated method stub
-            return null;
+            return Collections.emptySet();
         }
 
         @Override
@@ -685,7 +816,7 @@ public class LazyGeneratedCodecRegistry implements //
 
         private final Class augmentableType;
 
-        Map<Class, BindingCodec> rawAugmentationCodecs = new WeakHashMap<>();
+        Map<Class, AugmentationCodec<?>> localAugmentationCodecs = new WeakHashMap<>();
 
         public AugmentableCompositeCodec(Class type) {
             checkArgument(Augmentable.class.isAssignableFrom(type));
@@ -709,14 +840,8 @@ public class LazyGeneratedCodecRegistry implements //
                 augmentationField.setAccessible(true);
                 Map<Class, Augmentation> augMap = (Map<Class, Augmentation>) augmentationField.get(input);
                 return augMap;
-            } catch (NoSuchFieldException e) {
-
-            } catch (SecurityException e) {
-
-            } catch (IllegalArgumentException e) {
-
-            } catch (IllegalAccessException e) {
-
+            } catch (IllegalArgumentException | IllegalAccessException | NoSuchFieldException | SecurityException e) {
+                LOG.debug("Could not read augmentations for {}", input, e);
             }
             return Collections.emptyMap();
         }
@@ -724,53 +849,33 @@ public class LazyGeneratedCodecRegistry implements //
         private List serializeImpl(Map<Class, Augmentation> input) {
             List ret = new ArrayList<>();
             for (Entry<Class, Augmentation> entry : input.entrySet()) {
-                BindingCodec codec = getRawCodecForAugmentation(entry.getKey());
-                List output = (List) codec.serialize(new ValueWithQName(null, entry.getValue()));
-                ret.addAll(output);
+                AugmentationCodec codec = getCodecForAugmentation(entry.getKey());
+                CompositeNode node = codec.serialize(new ValueWithQName(null, entry.getValue()));
+                ret.addAll(node.getChildren());
             }
             return ret;
         }
 
-        private BindingCodec getRawCodecForAugmentation(Class key) {
-            BindingCodec ret = rawAugmentationCodecs.get(key);
-            if (ret != null) {
-                return ret;
-            }
-            try {
-                Class<? extends BindingCodec> retClass = generator.augmentationTransformerFor(key);
-                ret = retClass.newInstance();
-                rawAugmentationCodecs.put(key, ret);
-                return ret;
-            } catch (InstantiationException e) {
-
-            } catch (IllegalAccessException e) {
-
-            }
-            return null;
+        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) {
             Map<Class, Augmentation> ret = new HashMap<>();
             if (input instanceof CompositeNode) {
-                for (Entry<Class, BindingCodec> codec : rawAugmentationCodecs.entrySet()) {
-                    Augmentation value = (Augmentation) codec.getValue().deserialize(input);
-                    if (value != null) {
-                        ret.put(codec.getKey(), value);
+                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());
                     }
                 }
             }
             return ret;
         }
 
-        public Map<Class, BindingCodec> getRawAugmentationCodecs() {
-            return rawAugmentationCodecs;
-        }
-
-        public void setRawAugmentationCodecs(Map<Class, BindingCodec> rawAugmentationCodecs) {
-            this.rawAugmentationCodecs = rawAugmentationCodecs;
-        }
-
         public Class getAugmentableType() {
             return augmentableType;
         }
@@ -799,4 +904,104 @@ public class LazyGeneratedCodecRegistry implements //
             return getDelegate().serialize(input);
         }
     }
+
+    private static class AugmentationCodecWrapper<T extends Augmentation<?>> implements AugmentationCodec<T>,
+            Delegator<BindingCodec> {
+
+        private BindingCodec delegate;
+
+        public AugmentationCodecWrapper(BindingCodec<Map<QName, Object>, Object> rawCodec) {
+            this.delegate = rawCodec;
+        }
+
+        @Override
+        public BindingCodec getDelegate() {
+            return delegate;
+        }
+
+        @Override
+        public CompositeNode serialize(ValueWithQName<T> input) {
+            @SuppressWarnings("unchecked")
+            List<Map<QName, Object>> rawValues = (List<Map<QName, Object>>) getDelegate().serialize(input);
+            List<Node<?>> serialized = new ArrayList<>(rawValues.size());
+            for (Map<QName, Object> val : rawValues) {
+                serialized.add(toNode(val));
+            }
+            return new CompositeNodeTOImpl(input.getQname(), null, serialized);
+        }
+
+        @Override
+        @SuppressWarnings("unchecked")
+        public ValueWithQName<T> deserialize(Node<?> input) {
+            Object rawCodecValue = getDelegate().deserialize((Map<QName, Object>) input);
+            return new ValueWithQName<T>(input.getNodeType(), (T) rawCodecValue);
+        }
+    }
+
+    private class IdentityCompositeCodec implements IdentitityCodec {
+
+        @Override
+        public Object deserialize(Object input) {
+            checkArgument(input instanceof QName);
+            return deserialize((QName) input);
+        }
+
+        @Override
+        public Class<?> deserialize(QName input) {
+            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) {
+                return null;
+            }
+            return softref.get();
+        }
+
+        @Override
+        public QName serialize(Class input) {
+            checkArgument(BaseIdentity.class.isAssignableFrom(input));
+            bindingClassEncountered(input);
+            QName qname = identityQNames.get(input);
+            if (qname != null) {
+                return qname;
+            }
+            ConcreteType typeref = Types.typeForClass(input);
+            qname = typeToQname.get(typeref);
+            if (qname != null) {
+                identityQNames.put(input, qname);
+            }
+            return qname;
+        }
+
+        @Override
+        public Object serialize(Object input) {
+            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;
+    }
 }
\ No newline at end of file