BUG-1485: switch BuilderTemplate to new range enforcement
[yangtools.git] / code-generator / binding-java-api-generator / src / main / java / org / opendaylight / yangtools / sal / java / api / generator / BuilderTemplate.xtend
index 7dd355230387a27ca245954b12a8ab9dd826e705..038a2a8c18a1a7252e3d52dc3050dbea48842e5c 100644 (file)
@@ -7,7 +7,14 @@
  */
 package org.opendaylight.yangtools.sal.java.api.generator
 
-import java.util.Arrays;
+import com.google.common.collect.ImmutableSortedSet
+import com.google.common.collect.Range
+import java.util.ArrayList
+import java.util.Arrays
+import java.util.Collection
+import java.util.Collections
+import java.util.HashMap
+import java.util.HashSet
 import java.util.LinkedHashSet
 import java.util.List
 import java.util.Map
@@ -15,23 +22,24 @@ import java.util.Set
 import org.opendaylight.yangtools.binding.generator.util.ReferencedTypeImpl
 import org.opendaylight.yangtools.binding.generator.util.Types
 import org.opendaylight.yangtools.binding.generator.util.generated.type.builder.GeneratedTOBuilderImpl
+import org.opendaylight.yangtools.sal.binding.model.api.ConcreteType
 import org.opendaylight.yangtools.sal.binding.model.api.GeneratedProperty
 import org.opendaylight.yangtools.sal.binding.model.api.GeneratedTransferObject
 import org.opendaylight.yangtools.sal.binding.model.api.GeneratedType
 import org.opendaylight.yangtools.sal.binding.model.api.MethodSignature
 import org.opendaylight.yangtools.sal.binding.model.api.Type
 import org.opendaylight.yangtools.yang.binding.Augmentable
-import static org.opendaylight.yangtools.binding.generator.util.Types.*
-import java.util.HashMap
-import java.util.Collections
 import org.opendaylight.yangtools.yang.binding.DataObject
-import java.util.ArrayList
-import java.util.HashSet
-import java.util.Collection
 import org.opendaylight.yangtools.yang.binding.Identifiable
+import org.opendaylight.yangtools.concepts.Builder
+import org.opendaylight.yangtools.yang.binding.AugmentationHolder
+import org.opendaylight.yangtools.sal.binding.model.api.Restrictions
+import java.math.BigDecimal
+import java.math.BigInteger
+import com.google.common.collect.ImmutableList
 
 /**
- * Template for generating JAVA builder classes. 
+ * Template for generating JAVA builder classes.
  */
 
 class BuilderTemplate extends BaseTemplate {
@@ -46,6 +54,11 @@ class BuilderTemplate extends BaseTemplate {
      */
     val static BUILDER = 'Builder'
 
+    /**
+     * Constant with the name of the BuilderFor interface
+     */
+     val static BUILDERFOR = Builder.simpleName;
+
     /**
      * Constant with suffix for the classes which are generated from the builder classes.
      */
@@ -61,6 +74,8 @@ class BuilderTemplate extends BaseTemplate {
      */
     val Set<GeneratedProperty> properties
 
+    private static val METHOD_COMPARATOR = new AlphabeticallyTypeMemberComparator<MethodSignature>();
+
     /**
      * Constructs new instance of this class.
      * @throws IllegalArgumentException if <code>genType</code> equals <code>null</code>
@@ -68,25 +83,28 @@ class BuilderTemplate extends BaseTemplate {
     new(GeneratedType genType) {
         super(genType)
         this.properties = propertiesFromMethods(createMethods)
+        importMap.put(Builder.simpleName, Builder.package.name)
     }
 
     /**
      * Returns set of method signature instances which contains all the methods of the <code>genType</code>
      * and all the methods of the implemented interfaces.
-     * 
+     *
      * @returns set of method signature instances
      */
     def private Set<MethodSignature> createMethods() {
-        val Set<MethodSignature> methods = new LinkedHashSet
+        val Set<MethodSignature> methods = new LinkedHashSet();
         methods.addAll(type.methodDefinitions)
         collectImplementedMethods(methods, type.implements)
-        return methods
+        val Set<MethodSignature> sortedMethods = ImmutableSortedSet.orderedBy(METHOD_COMPARATOR).addAll(methods).build()
+
+        return sortedMethods
     }
 
     /**
-     * Adds to the <code>methods</code> set all the methods of the <code>implementedIfcs</code> 
-     * and recursivelly their implemented interfaces.
-     * 
+     * Adds to the <code>methods</code> set all the methods of the <code>implementedIfcs</code>
+     * and recursively their implemented interfaces.
+     *
      * @param methods set of method signatures
      * @param implementedIfcs list of implemented interfaces
      */
@@ -119,7 +137,7 @@ class BuilderTemplate extends BaseTemplate {
 
     /**
      * Returns the first element of the list <code>elements</code>.
-     * 
+     *
      * @param list of elements
      */
     def private <E> first(List<E> elements) {
@@ -128,7 +146,7 @@ class BuilderTemplate extends BaseTemplate {
 
     /**
      * Returns the name of the package from <code>fullyQualifiedName</code>.
-     * 
+     *
      * @param fullyQualifiedName string with fully qualified type name (package + type)
      * @return string with the package name
      */
@@ -137,12 +155,12 @@ class BuilderTemplate extends BaseTemplate {
         return if (lastDotIndex == -1) "" else fullyQualifiedName.substring(0, lastDotIndex)
     }
 
-       /**
-        * Returns the name of tye type from <code>fullyQualifiedName</code>
-        * 
-        * @param fullyQualifiedName string with fully qualified type name (package + type)
-        * @return string with the name of the type
-        */
+    /**
+     * Returns the name of tye type from <code>fullyQualifiedName</code>
+     *
+     * @param fullyQualifiedName string with fully qualified type name (package + type)
+     * @return string with the name of the type
+     */
     def private String getName(String fullyQualifiedName) {
         val lastDotIndex = fullyQualifiedName.lastIndexOf(Constants.DOT)
         return if (lastDotIndex == -1) fullyQualifiedName else fullyQualifiedName.substring(lastDotIndex + 1)
@@ -150,11 +168,11 @@ class BuilderTemplate extends BaseTemplate {
 
     /**
      * Creates set of generated property instances from getter <code>methods</code>.
-     * 
-     * @param set of method signature instances which should be transformed to list of properties 
+     *
+     * @param set of method signature instances which should be transformed to list of properties
      * @return set of generated property instances which represents the getter <code>methods</code>
      */
-    def private propertiesFromMethods(Set<MethodSignature> methods) {
+    def private propertiesFromMethods(Collection<MethodSignature> methods) {
         if (methods == null || methods.isEmpty()) {
             return Collections.emptySet
         }
@@ -170,14 +188,14 @@ class BuilderTemplate extends BaseTemplate {
 
     /**
      * Creates generated property instance from the getter <code>method</code> name and return type.
-     * 
+     *
      * @param method method signature from which is the method name and return type obtained
      * @return generated property instance for the getter <code>method</code>
      * @throws IllegalArgumentException<ul>
-     *         <li>if the <code>method</code> equals <code>null</code></li>
-     *         <li>if the name of the <code>method</code> equals <code>null</code></li>
-     *         <li>if the name of the <code>method</code> is empty</li>
-     *         <li>if the return type of the <code>method</code> equals <code>null</code></li>
+     *  <li>if the <code>method</code> equals <code>null</code></li>
+     *  <li>if the name of the <code>method</code> equals <code>null</code></li>
+     *  <li>if the name of the <code>method</code> is empty</li>
+     *  <li>if the return type of the <code>method</code> equals <code>null</code></li>
      * </ul>
      */
     def private GeneratedProperty propertyFromGetter(MethodSignature method) {
@@ -185,9 +203,9 @@ class BuilderTemplate extends BaseTemplate {
             throw new IllegalArgumentException("Method, method name, method return type reference cannot be NULL or empty!")
         }
         var prefix = "get";
-        if(BOOLEAN.equals(method.returnType)) {
+        if(Types.BOOLEAN.equals(method.returnType)) {
             prefix = "is";
-        } 
+        }
         if (method.name.startsWith(prefix)) {
             val fieldName = method.getName().substring(prefix.length()).toFirstLower
             val tmpGenTO = new GeneratedTOBuilderImpl("foo", "foo")
@@ -197,20 +215,22 @@ class BuilderTemplate extends BaseTemplate {
     }
 
     /**
-     * Template method which generates JAVA class body for builder class and for IMPL class. 
-     * 
+     * Template method which generates JAVA class body for builder class and for IMPL class.
+     *
      * @return string with JAVA source code
      */
     override body() '''
-
-        public class «type.name»«BUILDER» {
+        «wrapToDocumentation(formatDataForJavaDoc(type))»
+        public class «type.name»«BUILDER» implements «BUILDERFOR» <«type.importedName»> {
 
             «generateFields(false)»
 
-            «generateAugmentField(true)»
+            «generateAugmentField(false)»
 
             «generateConstructorsFromIfcs(type)»
 
+            «generateCopyConstructor(false)»
+
             «generateMethodFieldsFrom(type)»
 
             «generateGetters(false)»
@@ -227,16 +247,16 @@ class BuilderTemplate extends BaseTemplate {
 
                 «generateFields(true)»
 
-                «generateAugmentField(false)»
+                «generateAugmentField(true)»
 
-                «generateConstructor»
+                «generateCopyConstructor(true)»
 
                 «generateGetters(true)»
 
                 «generateHashCode()»
 
                 «generateEquals()»
-                
+
                 «generateToString(properties)»
             }
 
@@ -248,7 +268,7 @@ class BuilderTemplate extends BaseTemplate {
      */
     def private generateConstructorsFromIfcs(Type type) '''
         public «type.name»«BUILDER»() {
-        } 
+        }
         «IF (type instanceof GeneratedType && !(type instanceof GeneratedTransferObject))»
             «val ifc = type as GeneratedType»
             «FOR impl : ifc.implements»
@@ -262,14 +282,12 @@ class BuilderTemplate extends BaseTemplate {
      */
     def private Object generateConstructorFromIfc(Type impl) '''
         «IF (impl instanceof GeneratedType)»
-            «val implType = impl as GeneratedType»
-
-            «IF !(implType.methodDefinitions.empty)»
-                public «type.name»«BUILDER»(«implType.fullyQualifiedName» arg) {
-                    «printConstructorPropertySetter(implType)»
+            «IF !(impl.methodDefinitions.empty)»
+                public «type.name»«BUILDER»(«impl.fullyQualifiedName» arg) {
+                    «printConstructorPropertySetter(impl)»
                 }
             «ENDIF»
-            «FOR implTypeImplement : implType.implements»
+            «FOR implTypeImplement : impl.implements»
                 «generateConstructorFromIfc(implTypeImplement)»
             «ENDFOR»
         «ENDIF»
@@ -314,7 +332,7 @@ class BuilderTemplate extends BaseTemplate {
 
     def private generateMethodFieldsFromComment(GeneratedType type) '''
         /**
-         Set fields from given grouping argument. Valid argument is instance of one of following types:
+         *Set fields from given grouping argument. Valid argument is instance of one of following types:
          * <ul>
          «FOR impl : type.getAllIfcs»
          * <li>«impl.fullyQualifiedName»</li>
@@ -365,7 +383,7 @@ class BuilderTemplate extends BaseTemplate {
                 baseIfcs.add(ifc)
             }
         }
-        return baseIfcs 
+        return baseIfcs
     }
 
     private def Set<Type> getAllIfcs(Type type) {
@@ -379,7 +397,7 @@ class BuilderTemplate extends BaseTemplate {
                 baseIfcs.addAll(impl.getAllIfcs)
             }
         }
-        return baseIfcs 
+        return baseIfcs
     }
 
     private def List<String> toListOfNames(Collection<Type> types) {
@@ -390,56 +408,197 @@ class BuilderTemplate extends BaseTemplate {
         return names
     }
 
-       /**
-        * Template method which generates class attributes.
-        * 
-        * @param boolean value which specify whether field is|isn't final
-        * @return string with class attributes and their types
-        */
+    /**
+     * Template method which generates class attributes.
+     *
+     * @param boolean value which specify whether field is|isn't final
+     * @return string with class attributes and their types
+     */
     def private generateFields(boolean _final) '''
-        «IF !properties.empty»
+        «IF properties !== null»
             «FOR f : properties»
                 private«IF _final» final«ENDIF» «f.returnType.importedName» «f.fieldName»;
+                «val restrictions = f.returnType.restrictions»
+                «IF !_final && restrictions != null»
+                    «IF !(restrictions.lengthConstraints.empty)»
+                        private static «List.importedName»<«Range.importedName»<«f.returnType.importedNumber»>> «f.fieldName»_length;
+                    «ENDIF»
+                    «IF !(restrictions.rangeConstraints.empty)»
+                        private static «List.importedName»<«Range.importedName»<«f.returnType.importedNumber»>> «f.fieldName»_range;
+                    «ENDIF»
+                «ENDIF»
             «ENDFOR»
         «ENDIF»
     '''
 
-    def private generateAugmentField(boolean init) '''
+    def private generateAugmentField(boolean isPrivate) '''
         «IF augmentField != null»
-            private «Map.importedName»<«Class.importedName»<? extends «augmentField.returnType.importedName»>, «augmentField.returnType.importedName»> «augmentField.name» = new «HashMap.importedName»<>();
+            «IF isPrivate»private «ENDIF»«Map.importedName»<«Class.importedName»<? extends «augmentField.returnType.importedName»>, «augmentField.returnType.importedName»> «augmentField.name» = «Collections.importedName».emptyMap();
         «ENDIF»
     '''
 
-       /**
-        * Template method which generates setter methods
-        * 
-        * @return string with the setter methods 
-        */
+    /**
+     * Template method which generates setter methods
+     *
+     * @return string with the setter methods
+     */
     def private generateSetters() '''
         «FOR field : properties SEPARATOR '\n'»
-            public «type.name»«BUILDER» set«field.name.toFirstUpper»(«field.returnType.importedName» value) {
-                «generateRestrictions(field, "value")»
+            «val length = field.fieldName + "_length"»
+            «val restrictions = field.returnType.restrictions»
+            «IF restrictions != null»
+                «IF !restrictions.rangeConstraints.nullOrEmpty»
+                    «val rangeGenerator = AbstractRangeGenerator.forType(field.returnType)»
+                    «rangeGenerator.generateRangeChecker(field.name.toFirstUpper, restrictions.rangeConstraints)»
 
+                «ENDIF»
+            «ENDIF»
+            public «type.name»«BUILDER» set«field.name.toFirstUpper»(«field.returnType.importedName» value) {
+                «IF restrictions != null && !restrictions.rangeConstraints.nullOrEmpty»
+                if (value != null) {
+                    «val rangeGenerator = AbstractRangeGenerator.forType(field.returnType)»
+                    «IF field.returnType instanceof ConcreteType»
+                        «rangeGenerator.generateRangeCheckerCall(field.name.toFirstUpper, "value")»
+                    «ELSE»
+                        «rangeGenerator.generateRangeCheckerCall(field.name.toFirstUpper, "value.getValue()")»
+                    «ENDIF»
+                }
+                «ENDIF»
+                «generateRestrictions(field, "value", length)»
                 this.«field.fieldName» = value;
                 return this;
             }
+            «generateLengthMethod(length, field.returnType, type.name+BUILDER, length)»
+            «val range = field.fieldName + "_range"»
+            «generateRangeMethod(range, restrictions, field.returnType, type.name+BUILDER, range)»
         «ENDFOR»
         «IF augmentField != null»
 
             public «type.name»«BUILDER» add«augmentField.name.toFirstUpper»(«Class.importedName»<? extends «augmentField.returnType.importedName»> augmentationType, «augmentField.returnType.importedName» augmentation) {
+                if (augmentation == null) {
+                    return remove«augmentField.name.toFirstUpper»(augmentationType);
+                }
+
+                if (!(this.«augmentField.name» instanceof «HashMap.importedName»)) {
+                    this.«augmentField.name» = new «HashMap.importedName»<>();
+                }
+
                 this.«augmentField.name».put(augmentationType, augmentation);
                 return this;
             }
+
+            public «type.name»«BUILDER» remove«augmentField.name.toFirstUpper»(«Class.importedName»<? extends «augmentField.returnType.importedName»> augmentationType) {
+                if (this.«augmentField.name» instanceof «HashMap.importedName») {
+                    this.«augmentField.name».remove(augmentationType);
+                }
+                return this;
+            }
         «ENDIF»
     '''
 
-    /**
-     * Template method which generate constructor for IMPL class.
-     * 
-     * @return string with IMPL class constructor
-     */
-    def private generateConstructor() '''
-        private «type.name»«IMPL»(«type.name»«BUILDER» builder) {
+    def private generateRestrictions(GeneratedProperty field, String paramName, String lengthGetter) '''
+        «val Type type = field.returnType»
+        «IF type instanceof ConcreteType»
+            «createRestrictions(type, paramName, type.name.contains("["), lengthGetter)»
+        «ELSEIF type instanceof GeneratedTransferObject»
+            «createRestrictions(type, paramName, isArrayType(type as GeneratedTransferObject), lengthGetter)»
+        «ENDIF»
+    '''
+
+    def private createRestrictions(Type type, String paramName, boolean isArray, String lengthGetter) '''
+        «val restrictions = type.getRestrictions»
+        «IF restrictions !== null»
+            «val boolean isNestedType = !(type instanceof ConcreteType)»
+            «IF !restrictions.lengthConstraints.empty»
+                «generateLengthRestriction(type, paramName, lengthGetter, isNestedType, isArray)»
+            «ENDIF»
+        «ENDIF»
+    '''
+
+    def private generateLengthRestriction(Type type, String paramName, String getterName, boolean isNestedType, boolean isArray) '''
+        «val restrictions = type.getRestrictions»
+        if («paramName» != null) {
+            «val clazz = restrictions.lengthConstraints.iterator.next.min.class»
+            «printLengthConstraint(type, clazz, paramName, isNestedType, isArray)»
+            boolean isValidLength = false;
+            for («Range.importedName»<«clazz.importedNumber»> r : «getterName»()) {
+                if (r.contains(_constraint)) {
+                    isValidLength = true;
+                }
+            }
+            if (!isValidLength) {
+                throw new IllegalArgumentException(String.format("Invalid length: %s, expected: %s.", «paramName», «getterName»));
+            }
+        }
+    '''
+
+    def private generateLengthMethod(String methodName, Type type, String className, String varName) '''
+        «val Restrictions restrictions = type.restrictions»
+        «IF restrictions != null && !(restrictions.lengthConstraints.empty)»
+            «val numberClass = restrictions.lengthConstraints.iterator.next.min.class»
+            /**
+             * @deprecated This method is slated for removal in a future release. See BUG-1485 for details.
+             */
+            @Deprecated
+            public static «List.importedName»<«Range.importedName»<«numberClass.importedNumber»>> «methodName»() {
+                «IF numberClass.equals(typeof(BigDecimal))»
+                    «lengthBody(restrictions, numberClass, className, varName)»
+                «ELSE»
+                    «lengthBody(restrictions, typeof(BigInteger), className, varName)»
+                «ENDIF»
+            }
+        «ENDIF»
+    '''
+
+    def private lengthBody(Restrictions restrictions, Class<? extends Number> numberClass, String className, String varName) '''
+        if («varName» == null) {
+            synchronized («className».class) {
+                if («varName» == null) {
+                    «ImmutableList.importedName».Builder<«Range.importedName»<«numberClass.importedName»>> builder = «ImmutableList.importedName».builder();
+                    «FOR r : restrictions.lengthConstraints»
+                        builder.add(«Range.importedName».closed(«numericValue(numberClass, r.min)», «numericValue(numberClass, r.max)»));
+                    «ENDFOR»
+                    «varName» = builder.build();
+                }
+            }
+        }
+        return «varName»;
+    '''
+
+    def private generateRangeMethod(String methodName, Restrictions restrictions, Type returnType, String className, String varName) '''
+        «IF restrictions != null && !(restrictions.rangeConstraints.empty)»
+            «val number = returnType.importedNumber»
+            /**
+             * @deprecated This method is slated for removal in a future release. See BUG-1485 for details.
+             */
+            @Deprecated
+            public static «List.importedName»<«Range.importedName»<«number»>> «methodName»() {
+                «IF returnType.fullyQualifiedName.equals(BigDecimal.canonicalName)»
+                    «rangeBody(restrictions, BigDecimal, className, varName)»
+                «ELSE»
+                    «rangeBody(restrictions, BigInteger, className, varName)»
+                «ENDIF»
+            }
+        «ENDIF»
+    '''
+
+    def private rangeBody(Restrictions restrictions, Class<? extends Number> numberClass, String className, String varName) '''
+        if («varName» == null) {
+            synchronized («className».class) {
+                if («varName» == null) {
+                    «ImmutableList.importedName».Builder<«Range.importedName»<«numberClass.importedName»>> builder = «ImmutableList.importedName».builder();
+                    «FOR r : restrictions.rangeConstraints»
+                        builder.add(«Range.importedName».closed(«numericValue(numberClass, r.min)», «numericValue(numberClass, r.max)»));
+                    «ENDFOR»
+                    «varName» = builder.build();
+                }
+            }
+        }
+        return «varName»;
+    '''
+
+    def private CharSequence generateCopyConstructor(boolean impl) '''
+        «IF impl»private«ELSE»public«ENDIF» «type.name»«IF impl»«IMPL»«ELSE»«BUILDER»«ENDIF»(«type.name»«IF impl»«BUILDER»«ENDIF» base) {
             «val allProps = new ArrayList(properties)»
             «val isList = implementsIfc(type, Types.parameterizedTypeFor(Types.typeForClass(Identifiable), type))»
             «val keyType = type.getKey»
@@ -454,37 +613,52 @@ class BuilderTemplate extends BaseTemplate {
                     «removeProperty(allProps, field.name)»
                 «ENDFOR»
                 «removeProperty(allProps, "key")»
-                if (builder.getKey() == null) {
+                if (base.getKey() == null) {
                     this._key = new «keyType.importedName»(
                         «FOR keyProp : keyProps SEPARATOR ", "»
-                            builder.«keyProp.getterMethodName»()
+                            base.«keyProp.getterMethodName»()
                         «ENDFOR»
                     );
                     «FOR field : keyProps»
-                        this.«field.fieldName» = builder.«field.getterMethodName»();
+                        this.«field.fieldName» = base.«field.getterMethodName»();
                     «ENDFOR»
                 } else {
-                    this._key = builder.getKey();
+                    this._key = base.getKey();
                     «FOR field : keyProps»
                            this.«field.fieldName» = _key.«field.getterMethodName»();
                     «ENDFOR»
                 }
             «ENDIF»
             «FOR field : allProps»
-                this.«field.fieldName» = builder.«field.getterMethodName»();
+                this.«field.fieldName» = base.«field.getterMethodName»();
             «ENDFOR»
             «IF augmentField != null»
-               switch (builder.«augmentField.name».size()) {
-                case 0:
-                    this.«augmentField.name» = «Collections.importedName».emptyMap();
-                    break;
-                case 1:
-                    final «Map.importedName».Entry<«Class.importedName»<? extends «augmentField.returnType.importedName»>, «augmentField.returnType.importedName»> e = builder.«augmentField.name».entrySet().iterator().next();
-                    this.«augmentField.name» = «Collections.importedName».<«Class.importedName»<? extends «augmentField.returnType.importedName»>, «augmentField.returnType.importedName»>singletonMap(e.getKey(), e.getValue());
-                    break;
-                default :
-                    this.«augmentField.name» = new «HashMap.importedName»<>(builder.«augmentField.name»);
-                }
+                «IF impl»
+                    switch (base.«augmentField.name».size()) {
+                    case 0:
+                        this.«augmentField.name» = «Collections.importedName».emptyMap();
+                        break;
+                    case 1:
+                        final «Map.importedName».Entry<«Class.importedName»<? extends «augmentField.returnType.importedName»>, «augmentField.returnType.importedName»> e = base.«augmentField.name».entrySet().iterator().next();
+                        this.«augmentField.name» = «Collections.importedName».<«Class.importedName»<? extends «augmentField.returnType.importedName»>, «augmentField.returnType.importedName»>singletonMap(e.getKey(), e.getValue());
+                        break;
+                    default :
+                        this.«augmentField.name» = new «HashMap.importedName»<>(base.«augmentField.name»);
+                    }
+                «ELSE»
+                    if (base instanceof «type.name»«IMPL») {
+                        «type.name»«IMPL» impl = («type.name»«IMPL») base;
+                        if (!impl.«augmentField.name».isEmpty()) {
+                            this.«augmentField.name» = new «HashMap.importedName»<>(impl.«augmentField.name»);
+                        }
+                    } else if (base instanceof «AugmentationHolder.importedName») {
+                        @SuppressWarnings("unchecked")
+                        «AugmentationHolder.importedName»<«type.importedName»> casted =(«AugmentationHolder.importedName»<«type.importedName»>) base;
+                        if (!casted.augmentations().isEmpty()) {
+                            this.«augmentField.name» = new «HashMap.importedName»<>(casted.augmentations());
+                        }
+                    }
+                «ENDIF»
             «ENDIF»
         }
     '''
@@ -521,7 +695,7 @@ class BuilderTemplate extends BaseTemplate {
 
     /**
      * Template method which generate getter methods for IMPL class.
-     * 
+     *
      * @return string with getter methods
      */
     def private generateGetters(boolean addOverride) '''
@@ -546,7 +720,7 @@ class BuilderTemplate extends BaseTemplate {
 
     /**
      * Template method which generates the method <code>hashCode()</code>.
-     * 
+     *
      * @return string with the <code>hashCode()</code> method definition in JAVA format
      */
     def protected generateHashCode() '''
@@ -572,8 +746,8 @@ class BuilderTemplate extends BaseTemplate {
 
     /**
      * Template method which generates the method <code>equals()</code>.
-     * 
-     * @return string with the <code>equals()</code> method definition in JAVA format     
+     *
+     * @return string with the <code>equals()</code> method definition in JAVA format
      */
     def protected generateEquals() '''
         «IF !properties.empty || augmentField != null»
@@ -582,35 +756,50 @@ class BuilderTemplate extends BaseTemplate {
                 if (this == obj) {
                     return true;
                 }
-                if (obj == null) {
+                if (!(obj instanceof «DataObject.importedName»)) {
                     return false;
                 }
-                if (getClass() != obj.getClass()) {
+                if (!«type.importedName».class.equals(((«DataObject.importedName»)obj).getImplementedInterface())) {
                     return false;
                 }
-                «type.name»«IMPL» other = («type.name»«IMPL») obj;
+                «type.importedName» other = («type.importedName»)obj;
                 «FOR property : properties»
                     «val fieldName = property.fieldName»
                     if («fieldName» == null) {
-                        if (other.«fieldName» != null) {
+                        if (other.«property.getterMethodName»() != null) {
                             return false;
                         }
                     «IF property.returnType.name.contains("[")»
-                    } else if(!«Arrays.importedName».equals(«fieldName», other.«fieldName»)) {
+                    } else if(!«Arrays.importedName».equals(«fieldName», other.«property.getterMethodName»())) {
                     «ELSE»
-                    } else if(!«fieldName».equals(other.«fieldName»)) {
+                    } else if(!«fieldName».equals(other.«property.getterMethodName»())) {
                     «ENDIF»
                         return false;
                     }
                 «ENDFOR»
                 «IF augmentField != null»
-                    «val fieldName = augmentField.name»
-                    if («fieldName» == null) {
-                        if (other.«fieldName» != null) {
+                    if (getClass() == obj.getClass()) {
+                        // Simple case: we are comparing against self
+                        «type.name»«IMPL» otherImpl = («type.name»«IMPL») obj;
+                        «val fieldName = augmentField.name»
+                        if («fieldName» == null) {
+                            if (otherImpl.«fieldName» != null) {
+                                return false;
+                            }
+                        } else if(!«fieldName».equals(otherImpl.«fieldName»)) {
+                            return false;
+                        }
+                    } else {
+                        // Hard case: compare our augments with presence there...
+                        for («Map.importedName».Entry<«Class.importedName»<? extends «augmentField.returnType.importedName»>, «augmentField.returnType.importedName»> e : «augmentField.name».entrySet()) {
+                            if (!e.getValue().equals(other.getAugmentation(e.getKey()))) {
+                                return false;
+                            }
+                        }
+                        // .. and give the other one the chance to do the same
+                        if (!obj.equals(this)) {
                             return false;
                         }
-                    } else if(!«fieldName».equals(other.«fieldName»)) {
-                        return false;
                     }
                 «ENDIF»
                 return true;
@@ -664,5 +853,22 @@ class BuilderTemplate extends BaseTemplate {
     }
     '''
 
+    private def createDescription(GeneratedType type) {
+        return '''
+        Class that builds {@link «type.importedName»} instances.
+
+        @see «type.importedName»
+    '''
+    }
+
+    override def protected String formatDataForJavaDoc(GeneratedType type) {
+        val typeDescription = createDescription(type)
+
+        return '''
+            «IF !typeDescription.nullOrEmpty»
+            «typeDescription»
+            «ENDIF»
+        '''.toString
+    }
 }