Rename SchemaRootCodecContext
[mdsal.git] / binding / mdsal-binding-dom-codec / src / main / java / org / opendaylight / mdsal / binding / dom / codec / impl / BindingCodecContext.java
index 30527f50e6e938174f46b46f2ef25f070baa84ab..5a36251901cce18f614b6969c504fc94f431a55a 100644 (file)
@@ -12,53 +12,54 @@ import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
 import static java.util.Objects.requireNonNull;
 
-import com.google.common.annotations.Beta;
+import com.google.common.base.Strings;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
+import java.io.File;
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.lang.reflect.ParameterizedType;
 import java.lang.reflect.Type;
+import java.lang.reflect.WildcardType;
 import java.time.Instant;
-import java.util.AbstractMap.SimpleEntry;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
-import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.ServiceLoader;
-import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
-import javax.inject.Inject;
-import javax.inject.Singleton;
 import org.eclipse.jdt.annotation.NonNull;
 import org.eclipse.jdt.annotation.Nullable;
 import org.kohsuke.MetaInfServices;
+import org.opendaylight.mdsal.binding.dom.codec.api.BindingAugmentationCodecTreeNode;
 import org.opendaylight.mdsal.binding.dom.codec.api.BindingCodecTreeNode;
 import org.opendaylight.mdsal.binding.dom.codec.api.BindingDataObjectCodecTreeNode;
 import org.opendaylight.mdsal.binding.dom.codec.api.BindingInstanceIdentifierCodec;
 import org.opendaylight.mdsal.binding.dom.codec.api.BindingNormalizedNodeWriterFactory;
 import org.opendaylight.mdsal.binding.dom.codec.api.BindingStreamEventWriter;
-import org.opendaylight.mdsal.binding.dom.codec.impl.NodeCodecContext.CodecContextFactory;
-import org.opendaylight.mdsal.binding.dom.codec.loader.CodecClassLoader;
+import org.opendaylight.mdsal.binding.dom.codec.api.CommonDataObjectCodecTreeNode;
 import org.opendaylight.mdsal.binding.dom.codec.spi.AbstractBindingNormalizedNodeSerializer;
 import org.opendaylight.mdsal.binding.dom.codec.spi.BindingDOMCodecServices;
 import org.opendaylight.mdsal.binding.dom.codec.spi.BindingSchemaMapping;
-import org.opendaylight.mdsal.binding.model.api.GeneratedType;
+import org.opendaylight.mdsal.binding.loader.BindingClassLoader;
 import org.opendaylight.mdsal.binding.runtime.api.BindingRuntimeContext;
+import org.opendaylight.mdsal.binding.runtime.api.ListRuntimeType;
 import org.opendaylight.mdsal.binding.spec.reflect.BindingReflections;
 import org.opendaylight.yangtools.concepts.Delegator;
-import org.opendaylight.yangtools.concepts.IllegalArgumentCodec;
 import org.opendaylight.yangtools.concepts.Immutable;
 import org.opendaylight.yangtools.util.ClassLoaderUtils;
 import org.opendaylight.yangtools.yang.binding.Action;
+import org.opendaylight.yangtools.yang.binding.Augmentation;
+import org.opendaylight.yangtools.yang.binding.BaseIdentity;
+import org.opendaylight.yangtools.yang.binding.BaseNotification;
 import org.opendaylight.yangtools.yang.binding.DataContainer;
 import org.opendaylight.yangtools.yang.binding.DataObject;
 import org.opendaylight.yangtools.yang.binding.Identifiable;
@@ -71,6 +72,7 @@ import org.opendaylight.yangtools.yang.binding.RpcInput;
 import org.opendaylight.yangtools.yang.binding.RpcOutput;
 import org.opendaylight.yangtools.yang.common.QName;
 import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier;
+import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier.NodeIdentifier;
 import org.opendaylight.yangtools.yang.data.api.schema.ChoiceNode;
 import org.opendaylight.yangtools.yang.data.api.schema.ContainerNode;
 import org.opendaylight.yangtools.yang.data.api.schema.LeafSetNode;
@@ -80,18 +82,18 @@ import org.opendaylight.yangtools.yang.data.api.schema.UnkeyedListNode;
 import org.opendaylight.yangtools.yang.data.api.schema.ValueNode;
 import org.opendaylight.yangtools.yang.data.api.schema.stream.NormalizedNodeStreamWriter;
 import org.opendaylight.yangtools.yang.data.impl.schema.ImmutableNormalizedNodeStreamWriter;
-import org.opendaylight.yangtools.yang.data.impl.schema.NormalizedNodeResult;
+import org.opendaylight.yangtools.yang.data.impl.schema.NormalizationResultHolder;
 import org.opendaylight.yangtools.yang.model.api.AnydataSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.AnyxmlSchemaNode;
-import org.opendaylight.yangtools.yang.model.api.DataNodeContainer;
 import org.opendaylight.yangtools.yang.model.api.DataSchemaNode;
-import org.opendaylight.yangtools.yang.model.api.DocumentedNode.WithStatus;
 import org.opendaylight.yangtools.yang.model.api.LeafListSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.LeafSchemaNode;
-import org.opendaylight.yangtools.yang.model.api.ListSchemaNode;
+import org.opendaylight.yangtools.yang.model.api.TypeAware;
 import org.opendaylight.yangtools.yang.model.api.TypeDefinition;
 import org.opendaylight.yangtools.yang.model.api.TypedDataSchemaNode;
+import org.opendaylight.yangtools.yang.model.api.meta.EffectiveStatement;
 import org.opendaylight.yangtools.yang.model.api.stmt.SchemaNodeIdentifier.Absolute;
+import org.opendaylight.yangtools.yang.model.api.stmt.TypeDefinitionAware;
 import org.opendaylight.yangtools.yang.model.api.type.IdentityrefTypeDefinition;
 import org.opendaylight.yangtools.yang.model.api.type.InstanceIdentifierTypeDefinition;
 import org.opendaylight.yangtools.yang.model.api.type.LeafrefTypeDefinition;
@@ -99,9 +101,7 @@ import org.opendaylight.yangtools.yang.model.api.type.UnionTypeDefinition;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-@Beta
 @MetaInfServices(value = BindingDOMCodecServices.class)
-@Singleton
 public final class BindingCodecContext extends AbstractBindingNormalizedNodeSerializer
         implements BindingDOMCodecServices, Immutable, CodecContextFactory, DataObjectSerializerRegistry {
     private final class DataObjectSerializerProxy implements DataObjectSerializer, Delegator<DataObjectStreamer<?>> {
@@ -123,6 +123,13 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
     }
 
     private static final Logger LOG = LoggerFactory.getLogger(BindingCodecContext.class);
+    private static final @NonNull NodeIdentifier FAKE_NODEID = new NodeIdentifier(QName.create("fake", "fake"));
+    private static final File BYTECODE_DIRECTORY;
+
+    static {
+        final String dir = System.getProperty("org.opendaylight.mdsal.binding.dom.codec.loader.bytecodeDumpDirectory");
+        BYTECODE_DIRECTORY = Strings.isNullOrEmpty(dir) ? null : new File(dir);
+    }
 
     private final LoadingCache<Class<?>, DataObjectStreamer<?>> streamers = CacheBuilder.newBuilder().build(
         new CacheLoader<Class<?>, DataObjectStreamer<?>>() {
@@ -142,23 +149,23 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
             }
         });
 
-    private final @NonNull CodecClassLoader loader = CodecClassLoader.create();
+    private final @NonNull BindingClassLoader loader =
+        BindingClassLoader.create(BindingCodecContext.class, BYTECODE_DIRECTORY);
     private final @NonNull InstanceIdentifierCodec instanceIdentifierCodec;
     private final @NonNull IdentityCodec identityCodec;
     private final @NonNull BindingRuntimeContext context;
-    private final SchemaRootCodecContext<?> root;
+    private final @NonNull RootCodecContext<?> root;
 
     public BindingCodecContext() {
         this(ServiceLoader.load(BindingRuntimeContext.class).findFirst()
             .orElseThrow(() -> new IllegalStateException("Failed to load BindingRuntimeContext")));
     }
 
-    @Inject
     public BindingCodecContext(final BindingRuntimeContext context) {
         this.context = requireNonNull(context, "Binding Runtime Context is required.");
-        this.root = SchemaRootCodecContext.create(this);
-        this.identityCodec = new IdentityCodec(context);
-        this.instanceIdentifierCodec = new InstanceIdentifierCodec(this);
+        root = new RootCodecContext<>(this);
+        identityCodec = new IdentityCodec(context);
+        instanceIdentifierCodec = new InstanceIdentifierCodec(this);
     }
 
     @Override
@@ -167,7 +174,7 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
     }
 
     @Override
-    public CodecClassLoader getLoader() {
+    public BindingClassLoader getLoader() {
         return loader;
     }
 
@@ -199,49 +206,54 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
     @Override
     public Entry<YangInstanceIdentifier, BindingStreamEventWriter> newWriterAndIdentifier(
             final InstanceIdentifier<?> path, final NormalizedNodeStreamWriter domWriter) {
-        final List<YangInstanceIdentifier.PathArgument> yangArgs = new LinkedList<>();
-        final DataContainerCodecContext<?,?> codecContext = getCodecContextNode(path, yangArgs);
-        return new SimpleEntry<>(YangInstanceIdentifier.create(yangArgs), codecContext.createWriter(domWriter));
+        final var yangArgs = new ArrayList<YangInstanceIdentifier.PathArgument>();
+        final var codecContext = getCodecContextNode(path, yangArgs);
+        return Map.entry(YangInstanceIdentifier.of(yangArgs),
+            new BindingToNormalizedStreamWriter(codecContext, domWriter));
     }
 
     @Override
     public BindingStreamEventWriter newWriter(final InstanceIdentifier<?> path,
             final NormalizedNodeStreamWriter domWriter) {
-        return getCodecContextNode(path, null).createWriter(domWriter);
+        return new BindingToNormalizedStreamWriter(getCodecContextNode(path, null), domWriter);
     }
 
     @Override
     public BindingStreamEventWriter newRpcWriter(final Class<? extends DataContainer> rpcInputOrOutput,
             final NormalizedNodeStreamWriter domWriter) {
-        return root.getRpc(rpcInputOrOutput).createWriter(domWriter);
+        return new BindingToNormalizedStreamWriter(root.getRpc(rpcInputOrOutput), domWriter);
     }
 
     @Override
-    public BindingStreamEventWriter newNotificationWriter(final Class<? extends Notification> notification,
+    public BindingStreamEventWriter newNotificationWriter(final Class<? extends Notification<?>> notification,
             final NormalizedNodeStreamWriter domWriter) {
-        return root.getNotification(notification).createWriter(domWriter);
+        return new BindingToNormalizedStreamWriter(root.getNotification(notification), domWriter);
     }
 
     @Override
     public BindingStreamEventWriter newActionInputWriter(final Class<? extends Action<?, ?, ?>> action,
             final NormalizedNodeStreamWriter domWriter) {
-        return getActionCodec(action).input().createWriter(domWriter);
+        return new BindingToNormalizedStreamWriter(getActionCodec(action).input(), domWriter);
     }
 
     @Override
     public BindingStreamEventWriter newActionOutputWriter(final Class<? extends Action<?, ?, ?>> action,
             final NormalizedNodeStreamWriter domWriter) {
-        return getActionCodec(action).output().createWriter(domWriter);
+        return new BindingToNormalizedStreamWriter(getActionCodec(action).output(), domWriter);
     }
 
-    DataContainerCodecContext<?,?> getCodecContextNode(final InstanceIdentifier<?> binding,
+    @NonNull DataContainerCodecContext<?,?> getCodecContextNode(final InstanceIdentifier<?> binding,
             final List<YangInstanceIdentifier.PathArgument> builder) {
-        DataContainerCodecContext<?,?> currentNode = root;
-        for (final InstanceIdentifier.PathArgument bindingArg : binding.getPathArguments()) {
-            currentNode = currentNode.bindingPathArgumentChild(bindingArg, builder);
-            checkArgument(currentNode != null, "Supplied Instance Identifier %s is not valid.", binding);
+        DataContainerCodecContext<?, ?> current = root;
+        for (var bindingArg : binding.getPathArguments()) {
+            final var next = current.bindingPathArgumentChild(bindingArg, builder);
+            if (next == null) {
+                throw new IllegalArgumentException("%s is not valid: parent %s does not have a child %s".formatted(
+                    binding, current.bindingArg(), bindingArg));
+            }
+            current = next;
         }
-        return currentNode;
+        return current;
     }
 
     /**
@@ -258,14 +270,26 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
      */
     @Nullable BindingDataObjectCodecTreeNode<?> getCodecContextNode(final @NonNull YangInstanceIdentifier dom,
             final @Nullable Collection<InstanceIdentifier.PathArgument> bindingArguments) {
-        NodeCodecContext currentNode = root;
+        CodecContext currentNode = root;
         ListNodeCodecContext<?> currentList = null;
 
-        for (final YangInstanceIdentifier.PathArgument domArg : dom.getPathArguments()) {
+        for (var domArg : dom.getPathArguments()) {
             checkArgument(currentNode instanceof DataContainerCodecContext,
                 "Unexpected child of non-container node %s", currentNode);
-            final DataContainerCodecContext<?,?> previous = (DataContainerCodecContext<?, ?>) currentNode;
-            final NodeCodecContext nextNode = previous.yangPathArgumentChild(domArg);
+            final var previous = (DataContainerCodecContext<?, ?>) currentNode;
+            var nextNode = previous.yangPathArgumentChild(domArg);
+
+            /**
+             * Compatibility case: if it's determined the node belongs to augmentation
+             * then insert augmentation path argument in between.
+             */
+            if (nextNode instanceof AugmentationNodeContext<?> augmContext) {
+                if (bindingArguments != null) {
+                    bindingArguments.add(augmContext.bindingArg());
+                }
+                currentNode = nextNode;
+                nextNode = augmContext.yangPathArgumentChild(domArg);
+            }
 
             /*
              * List representation in YANG Instance Identifier consists of two
@@ -286,17 +310,17 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
                 }
                 currentList = null;
                 currentNode = nextNode;
-            } else if (nextNode instanceof ListNodeCodecContext) {
+            } else if (nextNode instanceof ListNodeCodecContext<?> listNode) {
                 // We enter list, we do not update current Node yet,
                 // since we need to verify
-                currentList = (ListNodeCodecContext<?>) nextNode;
+                currentList = listNode;
             } else if (nextNode instanceof ChoiceNodeCodecContext) {
                 // We do not add path argument for choice, since
                 // it is not supported by binding instance identifier.
                 currentNode = nextNode;
-            } else if (nextNode instanceof DataContainerCodecContext) {
+            } else if (nextNode instanceof DataContainerCodecContext<?, ?> containerNode) {
                 if (bindingArguments != null) {
-                    bindingArguments.add(((DataContainerCodecContext<?, ?>) nextNode).getBindingPathArgument(domArg));
+                    bindingArguments.add(containerNode.getBindingPathArgument(domArg));
                 }
                 currentNode = nextNode;
             } else if (nextNode instanceof ValueNodeCodecContext) {
@@ -343,57 +367,66 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
     }
 
     @Override
-    public ImmutableMap<Method, ValueNodeCodecContext> getLeafNodes(final Class<?> parentClass,
-            final DataNodeContainer childSchema) {
-        final Map<String, DataSchemaNode> getterToLeafSchema = new HashMap<>();
-        for (final DataSchemaNode leaf : childSchema.getChildNodes()) {
-            if (leaf instanceof TypedDataSchemaNode || leaf instanceof AnyxmlSchemaNode
-                    || leaf instanceof AnydataSchemaNode) {
-                getterToLeafSchema.put(BindingSchemaMapping.getGetterMethodName(leaf), leaf);
+    public ImmutableMap<Method, ValueNodeCodecContext> getLeafNodes(final Class<?> type,
+            final EffectiveStatement<?, ?> schema) {
+        final var getterToLeafSchema = new HashMap<String, DataSchemaNode>();
+        for (var stmt : schema.effectiveSubstatements()) {
+            if (stmt instanceof TypedDataSchemaNode typedSchema) {
+                putLeaf(getterToLeafSchema, typedSchema);
+            } else if (stmt instanceof AnydataSchemaNode anydataSchema) {
+                putLeaf(getterToLeafSchema, anydataSchema);
+            } else if (stmt instanceof AnyxmlSchemaNode anyxmlSchema) {
+                putLeaf(getterToLeafSchema, anyxmlSchema);
             }
         }
-        return getLeafNodesUsingReflection(parentClass, getterToLeafSchema);
+        return getLeafNodesUsingReflection(type, getterToLeafSchema);
+    }
+
+    private static void putLeaf(final Map<String, DataSchemaNode> map, final DataSchemaNode leaf) {
+        map.put(BindingSchemaMapping.getGetterMethodName(leaf), leaf);
     }
 
     private ImmutableMap<Method, ValueNodeCodecContext> getLeafNodesUsingReflection(
             final Class<?> parentClass, final Map<String, DataSchemaNode> getterToLeafSchema) {
-        final Map<Method, ValueNodeCodecContext> leaves = new HashMap<>();
-        for (final Method method : parentClass.getMethods()) {
+        final var leaves = new HashMap<Method, ValueNodeCodecContext>();
+        for (var method : parentClass.getMethods()) {
             // Only consider non-bridge methods with no arguments
             if (method.getParameterCount() == 0 && !method.isBridge()) {
                 final DataSchemaNode schema = getterToLeafSchema.get(method.getName());
 
                 final ValueNodeCodecContext valueNode;
-                if (schema instanceof LeafSchemaNode) {
-                    final LeafSchemaNode leafSchema = (LeafSchemaNode) schema;
-
+                if (schema instanceof LeafSchemaNode leafSchema) {
+                    // FIXME: MDSAL-670: this is not right as we need to find a concrete type, but this may return
+                    //                   Object.class
                     final Class<?> valueType = method.getReturnType();
-                    final IllegalArgumentCodec<Object, Object> codec = getCodec(valueType, leafSchema.getType());
+                    final ValueCodec<Object, Object> codec = getCodec(valueType, leafSchema.getType());
                     valueNode = LeafNodeCodecContext.of(leafSchema, codec, method.getName(), valueType,
                         context.getEffectiveModelContext());
-                } else if (schema instanceof LeafListSchemaNode) {
+                } else if (schema instanceof LeafListSchemaNode leafListSchema) {
                     final Optional<Type> optType = ClassLoaderUtils.getFirstGenericParameter(
                         method.getGenericReturnType());
                     checkState(optType.isPresent(), "Failed to find return type for %s", method);
 
                     final Class<?> valueType;
-                    final Type genericType = optType.get();
-                    if (genericType instanceof Class<?>) {
-                        valueType = (Class<?>) genericType;
-                    } else if (genericType instanceof ParameterizedType) {
-                        valueType = (Class<?>) ((ParameterizedType) genericType).getRawType();
+                    final Type genericType = optType.orElseThrow();
+                    if (genericType instanceof Class<?> clazz) {
+                        valueType = clazz;
+                    } else if (genericType instanceof ParameterizedType parameterized) {
+                        valueType = (Class<?>) parameterized.getRawType();
+                    } else if (genericType instanceof WildcardType) {
+                        // FIXME: MDSAL-670: this is not right as we need to find a concrete type
+                        valueType = Object.class;
                     } else {
                         throw new IllegalStateException("Unexpected return type " + genericType);
                     }
 
-                    final LeafListSchemaNode leafListSchema = (LeafListSchemaNode) schema;
-                    final IllegalArgumentCodec<Object, Object> codec = getCodec(valueType, leafListSchema.getType());
+                    final ValueCodec<Object, Object> codec = getCodec(valueType, leafListSchema.getType());
                     valueNode = new LeafSetNodeCodecContext(leafListSchema, codec, method.getName());
-                } else if (schema instanceof AnyxmlSchemaNode) {
-                    valueNode = new OpaqueNodeCodecContext.Anyxml<>((AnyxmlSchemaNode) schema, method.getName(),
+                } else if (schema instanceof AnyxmlSchemaNode anyxmlSchema) {
+                    valueNode = new OpaqueNodeCodecContext.Anyxml<>(anyxmlSchema, method.getName(),
                             opaqueReturnType(method), loader);
-                } else if (schema instanceof AnydataSchemaNode) {
-                    valueNode = new OpaqueNodeCodecContext.Anydata<>((AnydataSchemaNode) schema, method.getName(),
+                } else if (schema instanceof AnydataSchemaNode anydataSchema) {
+                    valueNode = new OpaqueNodeCodecContext.Anydata<>(anydataSchema, method.getName(),
                             opaqueReturnType(method), loader);
                 } else {
                     verify(schema == null, "Unhandled schema %s for method %s", schema, method);
@@ -408,66 +441,111 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
     }
 
     // FIXME: this is probably not right w.r.t. nulls
-    IllegalArgumentCodec<Object, Object> getCodec(final Class<?> valueType, final TypeDefinition<?> instantiatedType) {
-        if (Class.class.equals(valueType)) {
+    ValueCodec<Object, Object> getCodec(final Class<?> valueType, final TypeDefinition<?> instantiatedType) {
+        if (BaseIdentity.class.isAssignableFrom(valueType)) {
             @SuppressWarnings({ "unchecked", "rawtypes" })
-            final IllegalArgumentCodec<Object, Object> casted = (IllegalArgumentCodec) identityCodec;
+            final ValueCodec<Object, Object> casted = (ValueCodec) identityCodec;
             return casted;
         } else if (InstanceIdentifier.class.equals(valueType)) {
             @SuppressWarnings({ "unchecked", "rawtypes" })
-            final IllegalArgumentCodec<Object, Object> casted = (IllegalArgumentCodec) instanceIdentifierCodec;
+            final ValueCodec<Object, Object> casted = (ValueCodec) instanceIdentifierCodec;
             return casted;
         } else if (BindingReflections.isBindingClass(valueType)) {
             return getCodecForBindingClass(valueType, instantiatedType);
         }
-        return ValueTypeCodec.NOOP_CODEC;
+        // FIXME: MDSAL-670: this is right for most situations, but we must never return NOOP_CODEC for
+        //                   valueType=Object.class
+        return SchemaUnawareCodec.NOOP_CODEC;
     }
 
     @SuppressWarnings("checkstyle:illegalCatch")
     // FIXME: this is probably not right w.r.t. nulls
-    private IllegalArgumentCodec<Object, Object> getCodecForBindingClass(final Class<?> valueType,
+    private ValueCodec<Object, Object> getCodecForBindingClass(final Class<?> valueType,
             final TypeDefinition<?> typeDef) {
         if (typeDef instanceof IdentityrefTypeDefinition) {
-            return ValueTypeCodec.encapsulatedValueCodecFor(valueType, typeDef, identityCodec);
+            return new CompositeValueCodec.OfIdentity(valueType, identityCodec);
         } else if (typeDef instanceof InstanceIdentifierTypeDefinition) {
-            return ValueTypeCodec.encapsulatedValueCodecFor(valueType, typeDef, instanceIdentifierCodec);
-        } else if (typeDef instanceof UnionTypeDefinition) {
-            final Callable<UnionTypeCodec> unionLoader = UnionTypeCodec.loader(valueType, (UnionTypeDefinition) typeDef,
-                this);
+            return new CompositeValueCodec.OfInstanceIdentifier(valueType, instanceIdentifierCodec);
+        } else if (typeDef instanceof UnionTypeDefinition unionType) {
             try {
-                return unionLoader.call();
-            } catch (final Exception e) {
+                return UnionTypeCodec.of(valueType, unionType, this);
+            } catch (Exception e) {
                 throw new IllegalStateException("Unable to load codec for " + valueType, e);
             }
         } else if (typeDef instanceof LeafrefTypeDefinition) {
-            final Entry<GeneratedType, WithStatus> typeWithSchema = context.getTypeWithSchema(valueType);
-            final WithStatus schema = typeWithSchema.getValue();
-            checkState(schema instanceof TypeDefinition, "Unexpected schema %s", schema);
-            return getCodec(valueType, (TypeDefinition<?>) schema);
+            final var typeWithSchema = context.getTypeWithSchema(valueType);
+            final var schema = typeWithSchema.statement();
+            final TypeDefinition<?> def;
+            if (schema instanceof TypeDefinitionAware typeDefAware) {
+                def = typeDefAware.getTypeDefinition();
+            } else if (schema instanceof TypeAware typeAware) {
+                def = typeAware.getType();
+            } else {
+                throw new IllegalStateException("Unexpected schema " + schema);
+            }
+            return getCodec(valueType, def);
         }
-        return ValueTypeCodec.getCodecFor(valueType, typeDef);
+        return SchemaUnawareCodec.of(valueType, typeDef);
     }
 
     @Override
-    public IdentifiableItemCodec getPathArgumentCodec(final Class<?> listClz, final ListSchemaNode schema) {
+    public IdentifiableItemCodec getPathArgumentCodec(final Class<?> listClz, final ListRuntimeType type) {
         final Optional<Class<Identifier<?>>> optIdentifier = ClassLoaderUtils.findFirstGenericArgument(listClz,
                 Identifiable.class);
         checkState(optIdentifier.isPresent(), "Failed to find identifier for %s", listClz);
 
-        final Class<Identifier<?>> identifier = optIdentifier.get();
+        final Class<Identifier<?>> identifier = optIdentifier.orElseThrow();
         final Map<QName, ValueContext> valueCtx = new HashMap<>();
-        for (final ValueNodeCodecContext leaf : getLeafNodes(identifier, schema).values()) {
+        for (final ValueNodeCodecContext leaf : getLeafNodes(identifier, type.statement()).values()) {
             final QName name = leaf.getDomPathArgument().getNodeType();
             valueCtx.put(name, new ValueContext(identifier, leaf));
         }
-        return IdentifiableItemCodec.of(schema, identifier, listClz, valueCtx);
+        return IdentifiableItemCodec.of(type.statement(), identifier, listClz, valueCtx);
     }
 
+    @Override
+    public <E extends DataObject> CommonDataObjectCodecTreeNode<E> streamChild(final Class<E> childClass) {
+        return root.streamChild(childClass);
+    }
+
+    @Override
     @SuppressWarnings("unchecked")
+    public <A extends Augmentation<?>> BindingAugmentationCodecTreeNode<A> getAugmentationCodec(
+            final InstanceIdentifier<A> path) {
+        final var codecContext = getCodecContextNode(path, null);
+        if (codecContext instanceof BindingAugmentationCodecTreeNode) {
+            return (BindingAugmentationCodecTreeNode<A>) codecContext;
+        }
+        throw new IllegalArgumentException(path + " does not refer to an Augmentation");
+    }
+
     @Override
-    public <T extends DataObject> BindingDataObjectCodecTreeNode<T> getSubtreeCodec(final InstanceIdentifier<T> path) {
+    @SuppressWarnings("unchecked")
+    public <T extends DataObject> BindingDataObjectCodecTreeNode<T> getDataObjectCodec(
+            final InstanceIdentifier<T> path) {
+        final var codecContext = getCodecContextNode(path, null);
+        if (codecContext instanceof BindingDataObjectCodecTreeNode) {
+            return (BindingDataObjectCodecTreeNode<T>) codecContext;
+        }
+        throw new IllegalArgumentException(path + " does not refer to a plain DataObject");
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public <T extends DataObject> CodecWithPath<T> getSubtreeCodecWithPath(final InstanceIdentifier<T> path) {
+        final var yangArgs = new ArrayList<YangInstanceIdentifier.PathArgument>();
+        final var codecContext = getCodecContextNode(path, yangArgs);
+
         // TODO Do we need defensive check here?
-        return (BindingDataObjectCodecTreeNode<T>) getCodecContextNode(path, null);
+        return new CodecWithPath<>((CommonDataObjectCodecTreeNode<T>) codecContext,
+            YangInstanceIdentifier.of(yangArgs));
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public <T extends DataObject> CommonDataObjectCodecTreeNode<T> getSubtreeCodec(final InstanceIdentifier<T> path) {
+        // TODO Do we need defensive check here?
+        return (CommonDataObjectCodecTreeNode<T>) getCodecContextNode(path, null);
     }
 
     @Override
@@ -491,29 +569,67 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
     }
 
     @Override
-    public <T extends DataObject> Entry<YangInstanceIdentifier, NormalizedNode<?,?>> toNormalizedNode(
-            final InstanceIdentifier<T> path, final T data) {
-        final NormalizedNodeResult result = new NormalizedNodeResult();
-        // We create DOM stream writer which produces normalized nodes
-        final NormalizedNodeStreamWriter domWriter = ImmutableNormalizedNodeStreamWriter.from(result);
+    public <A extends Augmentation<?>> AugmentationResult toNormalizedAugmentation(final InstanceIdentifier<A> path,
+            final A data) {
+        final var result = toNormalizedNode(path, data);
+        if (result instanceof AugmentationResult augment) {
+            return augment;
+        }
+        throw new IllegalArgumentException(path + " does not identify an Augmentation");
+    }
+
+    @Override
+    public <T extends DataObject> NodeResult toNormalizedDataObject(final InstanceIdentifier<T> path, final T data) {
+        final var result = toNormalizedNode(path, data);
+        if (result instanceof NodeResult node) {
+            return node;
+        }
+        throw new IllegalArgumentException(path + " does not identify a plain DataObject");
+    }
 
+    @Override
+    public <T extends DataObject> NormalizedResult toNormalizedNode(final InstanceIdentifier<T> path, final T data) {
         // We create Binding Stream Writer which translates from Binding to Normalized Nodes
-        final Entry<YangInstanceIdentifier, BindingStreamEventWriter> writeCtx = newWriterAndIdentifier(path,
-            domWriter);
+        final var yangArgs = new ArrayList<YangInstanceIdentifier.PathArgument>();
+        final var codecContext = getCodecContextNode(path, yangArgs);
+        final var yangPath = YangInstanceIdentifier.of(yangArgs);
+
+        // We create DOM stream writer which produces normalized nodes
+        final var result = new NormalizationResultHolder();
+        final var domWriter = ImmutableNormalizedNodeStreamWriter.from(result);
+        final var bindingWriter = new BindingToNormalizedStreamWriter(codecContext, domWriter);
+        final var augment = codecContext instanceof BindingAugmentationCodecTreeNode<?> augmentNode ? augmentNode
+            : null;
 
-        // We get serializer which reads binding data and uses Binding To Normalized Node writer to write result
         try {
-            getSerializer(path.getTargetType()).serialize(data, writeCtx.getValue());
+            // Augmentations do not have a representation, so we are faking a ContainerNode as the parent and we will be
+            // extracting the resulting children.
+            if (augment != null) {
+                domWriter.startContainerNode(FAKE_NODEID, NormalizedNodeStreamWriter.UNKNOWN_SIZE);
+            }
+
+            // We get serializer which reads binding data and uses Binding To Normalized Node writer to write result
+            getSerializer(path.getTargetType()).serialize(data, bindingWriter);
+
+            if (augment != null) {
+                domWriter.endNode();
+            }
         } catch (final IOException e) {
             LOG.error("Unexpected failure while serializing path {} data {}", path, data, e);
             throw new IllegalStateException("Failed to create normalized node", e);
         }
-        return new SimpleEntry<>(writeCtx.getKey(),result.getResult());
+
+        // Terminate the fake container and extract it to the result
+        if (augment != null) {
+            return new AugmentationResult(yangPath, augment.childPathArguments(),
+                ImmutableList.copyOf(((ContainerNode) result.getResult().data()).body()));
+        }
+        return new NodeResult(yangPath, result.getResult().data());
     }
 
     @Override
     public Entry<InstanceIdentifier<?>, DataObject> fromNormalizedNode(final YangInstanceIdentifier path,
-            final NormalizedNode<?, ?> data) {
+            final NormalizedNode data) {
         if (notBindingRepresentable(data)) {
             return null;
         }
@@ -529,21 +645,20 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
         }
 
         final DataObject lazyObj = codec.deserialize(data);
-        final InstanceIdentifier<?> bindingPath = InstanceIdentifier.create(builder);
-        return new SimpleEntry<>(bindingPath, lazyObj);
+        final InstanceIdentifier<?> bindingPath = InstanceIdentifier.unsafeOf(builder);
+        return Map.entry(bindingPath, lazyObj);
     }
 
     @Override
-    public Notification fromNormalizedNodeNotification(final Absolute path, final ContainerNode data) {
+    public BaseNotification fromNormalizedNodeNotification(final Absolute path, final ContainerNode data) {
         return getNotificationContext(path).deserialize(data);
     }
 
     @Override
-    public Notification fromNormalizedNodeNotification(final Absolute path, final ContainerNode data,
+    public BaseNotification fromNormalizedNodeNotification(final Absolute path, final ContainerNode data,
             final Instant eventInstant) {
         return eventInstant == null ? fromNormalizedNodeNotification(path, data)
                 : getNotificationContext(path).deserialize(data, eventInstant);
-
     }
 
     @Override
@@ -565,10 +680,22 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
 
     @Override
     @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
-    public ContainerNode toNormalizedNodeNotification(@NonNull final Notification data) {
+    public ContainerNode toNormalizedNodeNotification(final Notification<?> data) {
         // FIXME: Should the cast to DataObject be necessary?
         return serializeDataObject((DataObject) data,
-            (ctx, iface, domWriter) -> ctx.newNotificationWriter(iface.asSubclass(Notification.class), domWriter));
+            (ctx, iface, domWriter) -> ctx.newNotificationWriter(
+                (Class<? extends Notification<?>>) iface.asSubclass(Notification.class), domWriter));
+    }
+
+    @Override
+    public ContainerNode toNormalizedNodeNotification(final Absolute path, final BaseNotification data) {
+        checkArgument(data instanceof DataObject, "Unexpected data %s", data);
+        @SuppressWarnings("rawtypes")
+        final NotificationCodecContext notifContext = getNotificationContext(path);
+        @SuppressWarnings("unchecked")
+        final var result = notifContext.serialize((DataObject) data);
+        verify(result instanceof ContainerNode, "Unexpected result %s from %s", result, data);
+        return (ContainerNode) result;
     }
 
     @Override
@@ -592,9 +719,9 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
 
     private <T extends DataContainer> @NonNull ContainerNode serializeDataObject(final DataObject data,
             final WriterFactoryMethod<T> newWriter) {
-        final NormalizedNodeResult result = new NormalizedNodeResult();
+        final var result = new NormalizationResultHolder();
         // We create DOM stream writer which produces normalized nodes
-        final NormalizedNodeStreamWriter domWriter = ImmutableNormalizedNodeStreamWriter.from(result);
+        final var domWriter = ImmutableNormalizedNodeStreamWriter.from(result);
         final Class<? extends DataObject> type = data.implementedInterface();
         @SuppressWarnings("unchecked")
         final BindingStreamEventWriter writer = newWriter.createWriter(this, (Class<T>) type, domWriter);
@@ -604,16 +731,16 @@ public final class BindingCodecContext extends AbstractBindingNormalizedNodeSeri
             LOG.error("Unexpected failure while serializing data {}", data, e);
             throw new IllegalStateException("Failed to create normalized node", e);
         }
-        return (ContainerNode) result.getResult();
+        return (ContainerNode) result.getResult().data();
     }
 
 
-    private static boolean notBindingRepresentable(final NormalizedNode<?, ?> data) {
+    private static boolean notBindingRepresentable(final NormalizedNode data) {
         // ValueNode covers LeafNode and LeafSetEntryNode
         return data instanceof ValueNode
-                || data instanceof MapNode || data instanceof UnkeyedListNode
-                || data instanceof ChoiceNode
-                || data instanceof LeafSetNode;
+            || data instanceof MapNode || data instanceof UnkeyedListNode
+            || data instanceof ChoiceNode
+            || data instanceof LeafSetNode;
     }
 
     @SuppressWarnings("rawtypes")