GenerateEntityImplementationsProcessor.java

/*
 * Copyright 2021 Red Hat, Inc. and/or its affiliates
 * and other contributors as indicated by the @author tags.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.keycloak.models.map.processor;

import org.keycloak.models.map.annotations.GenerateEntityImplementations;

import java.io.IOException;
import java.io.PrintWriter;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import javax.annotation.processing.SupportedAnnotationTypes;
import javax.annotation.processing.SupportedSourceVersion;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.ElementKind;
import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.Modifier;
import javax.lang.model.element.Name;
import javax.lang.model.element.TypeElement;
import javax.lang.model.type.TypeMirror;
import javax.tools.Diagnostic;
import javax.tools.Diagnostic.Kind;
import javax.tools.JavaFileObject;
import static org.keycloak.models.map.processor.FieldAccessorType.*;
import static org.keycloak.models.map.processor.Util.isSetType;
import static org.keycloak.models.map.processor.Util.methodParameters;
import java.util.Collection;
import java.util.Comparator;
import java.util.IdentityHashMap;
import java.util.Optional;
import java.util.TreeSet;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Stream;
import javax.lang.model.element.VariableElement;
import javax.lang.model.type.TypeKind;

/**
 *
 * @author hmlnarik
 */
@SupportedAnnotationTypes("org.keycloak.models.map.annotations.GenerateEntityImplementations")
@SupportedSourceVersion(SourceVersion.RELEASE_11)
public class GenerateEntityImplementationsProcessor extends AbstractGenerateEntityImplementationsProcessor {

    private static final Collection<String> autogenerated = new TreeSet<>();
    private static final String ID_FIELD_NAME = "Id";

    private final Generator[] generators = new Generator[] {
        new ClonerGenerator(),
        new DelegateGenerator(),
        new FieldsGenerator(),
        new FieldDelegateGenerator(),
        new ImplGenerator(),
    };

    @Override
    protected void afterAnnotationProcessing() {
        if (! autogenerated.isEmpty()) {
            try {
                JavaFileObject file = processingEnv.getFiler().createSourceFile("org.keycloak.models.map.common.AutogeneratedClasses");
                try (PrintWriter pw = new PrintWriterNoJavaLang(file.openWriter())) {
                    pw.println("package org.keycloak.models.map.common;");

                    pw.println("import " + FQN_DEEP_CLONER + ".Cloner;");
                    pw.println("import " + FQN_DEEP_CLONER + ".DelegateCreator;");
                    pw.println("import java.util.function.Function;");
                    pw.println("import " + FQN_DEEP_CLONER + ".EntityFieldDelegateCreator;");
                    pw.println("// DO NOT CHANGE THIS CLASS, IT IS GENERATED AUTOMATICALLY BY " + GenerateEntityImplementationsProcessor.class.getSimpleName());
                    generatedAnnotation(pw);
                    pw.println("public final class AutogeneratedClasses {");
                    pw.println("    public static final java.util.Map<Class<?>, Cloner<?>> CLONERS_WITH_ID = new java.util.HashMap<>();");
                    pw.println("    public static final java.util.Map<Class<?>, Cloner<?>> CLONERS_WITHOUT_ID = new java.util.HashMap<>();");
                    pw.println("    public static final java.util.Map<Class<?>, DelegateCreator<?>> DELEGATE_CREATORS = new java.util.HashMap<>();");
                    pw.println("    public static final java.util.Map<Class<?>, EntityFieldDelegateCreator<?>> ENTITY_FIELD_DELEGATE_CREATORS = new java.util.HashMap<>();");
                    pw.println("    public static final java.util.Map<Class<?>, Object> EMPTY_INSTANCES = new java.util.HashMap<>();");
                    pw.println("    public static final java.util.Map<Class<?>, EntityField<?>[]> ENTITY_FIELDS = new java.util.HashMap<>();");
                    pw.println("    public static final java.util.Map<Class<?>, Function<DeepCloner, ?>> CONSTRUCTORS_DC = new java.util.HashMap<>();");
                    pw.println("    static {");
                    autogenerated.forEach(pw::println);
                    pw.println("    }");
                    pw.println("}");
                }
            } catch (IOException ex) {
                Logger.getLogger(GenerateEntityImplementationsProcessor.class.getName()).log(Level.SEVERE, null, ex);
            }
        }
    }

    @Override
    protected Generator[] getGenerators() {
        return this.generators;
    }

    @Override
    protected boolean testAnnotationElement(TypeElement e) {
        if (e.getKind() != ElementKind.INTERFACE) {
            processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, "Annotation @GenerateEntityImplementations is only applicable to an interface", e);
            return false;
        }

        return true;
    }

    protected static String toEnumConstant(String key) {
        return key.replaceAll("([a-z])([A-Z])", "$1_$2").toUpperCase();
    }

    private class FieldsGenerator implements Generator {

        @Override
        public void generate(TypeElement e) throws IOException {
            Map<String, HashSet<ExecutableElement>> methodsPerAttribute = methodsPerAttributeMapping(e);
            String className = e.getQualifiedName().toString();
            String packageName = null;
            int lastDot = className.lastIndexOf('.');
            if (lastDot > 0) {
                packageName = className.substring(0, lastDot);
            }

            String simpleClassName = className.substring(lastDot + 1);
            String mapFieldsClassName = className + "Fields";
            String mapSimpleFieldsClassName = simpleClassName + "Fields";

            JavaFileObject file = processingEnv.getFiler().createSourceFile(mapFieldsClassName);
            try (PrintWriter pw = new PrintWriterNoJavaLang(file.openWriter())) {
                if (packageName != null) {
                    pw.println("package " + packageName + ";");
                }

                generatedAnnotation(pw);
                pw.println("public enum " + mapSimpleFieldsClassName + " implements " + FQN_ENTITY_FIELD + "<" + className + "> {");
                methodsPerAttribute.keySet().stream()
                  .sorted(NameFirstComparator.ID_INSTANCE)
                  .forEach(key -> {
                      pw.println("    " + toEnumConstant(key) + " {");
                      printEntityFieldMethods(pw, className, key, methodsPerAttribute.get(key));
                      pw.println("    },");
                  });
                pw.println("}");

                autogenerated.add("        ENTITY_FIELDS.put(" + className + ".class, " + mapFieldsClassName + ".values());");
            }
        }

        private void printEntityFieldMethods(PrintWriter pw, String className, String fieldName, HashSet<ExecutableElement> methods) {
            TypeMirror fieldType = determineFieldType(fieldName, methods);
            pw.println("        public static final String FIELD_NAME = \"" + fieldName + "\";");
            pw.println("        public static final String FIELD_NAME_DASHED = \"" + fieldName.replaceAll("([^_A-Z])([A-Z])", "$1-$2").toLowerCase() + "\";");
            pw.println("        public static final String FIELD_NAME_CAMEL_CASE = \"" + fieldName.substring(0, 1).toLowerCase() + fieldName.substring(1) + "\";");
            pw.println("        @SuppressWarnings(\"unchecked\") @Override public Class<?> getFieldClass() {");
            pw.println("            return " + types.erasure(fieldType) + ".class;");
            pw.println("        }");
            pw.println("        @Override public String getName() {");
            pw.println("            return FIELD_NAME;");
            pw.println("        }");
            pw.println("        @Override public String getNameDashed() {");
            pw.println("            return FIELD_NAME_DASHED;");
            pw.println("        }");
            pw.println("        @Override public String getNameCamelCase() {");
            pw.println("            return FIELD_NAME_CAMEL_CASE;");
            pw.println("        }");

            FieldAccessorType.getMethod(FieldAccessorType.GETTER, methods, fieldName, types, fieldType).ifPresent(method -> {
                if (Util.isCollectionType((TypeElement) types.asElement(types.erasure(fieldType)))) {
                    TypeMirror firstParameterType = Util.getGenericsDeclaration(method.getReturnType()).get(0);
                    pw.println("        @SuppressWarnings(\"unchecked\") @Override public Class<?> getCollectionElementClass() {");
                    pw.println("            return " + types.erasure(firstParameterType) + ".class;");
                    pw.println("        }");

                } else if (Util.isMapType((TypeElement) types.asElement(types.erasure(fieldType)))) {
                    TypeMirror firstParameterType = Util.getGenericsDeclaration(method.getReturnType()).get(0);
                    TypeMirror secondParameterType = Util.getGenericsDeclaration(method.getReturnType()).get(1);

                    pw.println("        @SuppressWarnings(\"unchecked\") @Override public Class<?> getMapKeyClass() {");
                    pw.println("            return " + types.erasure(firstParameterType) + ".class;");
                    pw.println("        }");
                    pw.println("        @SuppressWarnings(\"unchecked\") @Override public Class<?> getMapValueClass() {");
                    pw.println("            return " + types.erasure(secondParameterType) + ".class;");
                    pw.println("        }");
                }
            });

            for (ExecutableElement ee : methods) {
                FieldAccessorType fat = FieldAccessorType.determineType(ee, fieldName, types, fieldType);
                printMethodBody(pw, fat, ee, className, fieldType);
            }
        }

        private void printMethodBody(PrintWriter pw, FieldAccessorType accessorType, ExecutableElement method, String className, TypeMirror fieldType) {
            TypeMirror firstParameterType = method.getParameters().isEmpty()
              ? types.getNullType()
              : method.getParameters().get(0).asType();

            switch (accessorType) {
                case GETTER:
                    pw.println("        @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " get(" + className + " e) {");
                    pw.println("            return (" + fieldType + ") e." + method.getSimpleName() + "();");
                    pw.println("        }");
                    return;
                case SETTER:
                    pw.println("        @SuppressWarnings(\"unchecked\") @Override public <T> void set(" + className + " e, T value) {");
                    pw.println("            e." + method.getSimpleName() + "((" + firstParameterType + ") value);");
                    pw.println("        }");
                    return;
                case COLLECTION_ADD:
                    pw.println("        @SuppressWarnings(\"unchecked\") @Override public <T> void collectionAdd(" + className + " e, T value) {");
                    pw.println("            e." + method.getSimpleName() + "((" + firstParameterType + ") value);");
                    pw.println("        }");
                    return;
                case COLLECTION_DELETE:
                {
                    String returnType = method.getReturnType().getKind() == TypeKind.VOID ? "Void" : method.getReturnType().toString();
                    TypeElement fieldTypeElement = elements.getTypeElement(types.erasure(fieldType).toString());
                    if (Util.isMapType(fieldTypeElement)) {
                        pw.println("        @SuppressWarnings(\"unchecked\") @Override public <K> " + returnType + " mapRemove(" + className + " e, K p0) {");
                    } else {
                        pw.println("        @SuppressWarnings(\"unchecked\") @Override public <T> " + returnType + " collectionRemove(" + className + " e, T p0) {");
                    }
                    if (method.getReturnType().getKind() == TypeKind.VOID) {
                        pw.println("            e." + method.getSimpleName() + "((" + firstParameterType + ") p0); return null;");
                    } else {
                        pw.println("            return e." + method.getSimpleName() + "((" + firstParameterType + ") p0);");
                    }
                    pw.println("        }");
                    return;
                }
                case COLLECTION_DELETE_BY_ID:
                {
                    String returnType = method.getReturnType().getKind() == TypeKind.VOID ? "Void" : method.getReturnType().toString();
                    pw.println("        @SuppressWarnings(\"unchecked\") @Override public <K> " + returnType + " mapRemove(" + className + " e, K p0) {");
                    if (method.getReturnType().getKind() == TypeKind.VOID) {
                        pw.println("            e." + method.getSimpleName() + "((String) p0); return null;");
                    } else {
                        pw.println("            return e." + method.getSimpleName() + "((String) p0);");
                    }
                    pw.println("        }");
                    return;
                }
                case COLLECTION_GET_BY_ID:
                    pw.println("        @SuppressWarnings(\"unchecked\") @Override public <K> " + method.getReturnType() + " mapGet(" + className + " e, K key) {");
                    pw.println("            return e." + method.getSimpleName() + "((" + firstParameterType + ") key);");
                    pw.println("        }");
                    return;
                case MAP_ADD:
                    TypeMirror secondParameterType = method.getParameters().get(1).asType();
                    pw.println("        @SuppressWarnings(\"unchecked\") @Override public <K, T> void mapPut(" + className + " e, K key, T value) {");
                    pw.println("            e." + method.getSimpleName() + "((" + firstParameterType + ") key, (" + secondParameterType + ") value);");
                    pw.println("        }");
                    return;
                case MAP_GET:
                    pw.println("        @SuppressWarnings(\"unchecked\") @Override public <K> " + method.getReturnType() + " mapGet(" + className + " e, K key) {");
                    pw.println("            return (" + method.getReturnType() + ") e." + method.getSimpleName() + "((" + firstParameterType + ") key);");
                    pw.println("        }");
            }
        }
    }

    private class ImplGenerator implements Generator {

        @Override
        public void generate(TypeElement e) throws IOException {
            Map<String, HashSet<ExecutableElement>> methodsPerAttribute = methodsPerAttributeMapping(e);
            GenerateEntityImplementations an = e.getAnnotation(GenerateEntityImplementations.class);
            TypeElement parentTypeElement = elements.getTypeElement((an.inherits() == null || an.inherits().isEmpty()) ? "void" : an.inherits());
            if (parentTypeElement == null) {
                processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, "Unable to find type " + an.inherits() + " for inherits parameter for annotation " + GenerateEntityImplementations.class.getTypeName(), e);
            }
            final List<? extends Element> allParentMembers = elements.getAllMembers(parentTypeElement);
            String className = e.getQualifiedName().toString();
            String packageName = null;
            int lastDot = className.lastIndexOf('.');
            if (lastDot > 0) {
                packageName = className.substring(0, lastDot);
            }

            String simpleClassName = className.substring(lastDot + 1);
            String mapImplClassName = className + "Impl";
            String mapSimpleClassName = simpleClassName + "Impl";
            boolean hasId = methodsPerAttribute.containsKey(ID_FIELD_NAME) || allParentMembers.stream().anyMatch(el -> "getId".equals(el.getSimpleName().toString()));
            boolean hasDeepClone = allParentMembers.stream().filter(el -> el.getKind() == ElementKind.METHOD).anyMatch(el -> "deepClone".equals(el.getSimpleName().toString()));
            boolean needsDeepClone = fieldGetters(methodsPerAttribute)
              .map(ExecutableElement::getReturnType)
              .anyMatch(fieldType -> ! isKnownCollectionOfImmutableFinalTypes(fieldType) && ! isImmutableFinalType(fieldType));
            boolean usingGeneratedCloner = ! hasDeepClone && needsDeepClone;

            JavaFileObject file = processingEnv.getFiler().createSourceFile(mapImplClassName);
            try (PrintWriter pw = new PrintWriterNoJavaLang(file.openWriter())) {
                if (packageName != null) {
                    pw.println("package " + packageName + ";");
                }

                pw.println("import java.util.Objects;");
                pw.println("import java.util.Optional;");
                pw.println("import " + FQN_DEEP_CLONER + ";");
                pw.println("// DO NOT CHANGE THIS CLASS, IT IS GENERATED AUTOMATICALLY BY " + GenerateEntityImplementationsProcessor.class.getSimpleName());
                generatedAnnotation(pw);
                pw.println("public class " + mapSimpleClassName + (an.inherits().isEmpty() ? "" : " extends " + an.inherits()) + " implements " + className + " {");

                pw.println("    "
                        + "private "
                        + mapSimpleClassName + "() { this(DeepCloner.DUMB_CLONER); } // Nullary constructor only for Jackson deserialization"
                );
                pw.println("    "
                        + "public "
                        + mapSimpleClassName + "(DeepCloner cloner) { super(); " + (!usingGeneratedCloner ? "" : "this.cloner = cloner;") + "}"
                );

                // equals, hashCode, toString
                pw.println("    @Override public boolean equals(Object o) {");
                pw.println("        if (o == this) return true; ");
                pw.println("        if (! (o instanceof " + mapSimpleClassName + ")) return false; ");
                pw.println("        " + mapSimpleClassName + " other = (" + mapSimpleClassName + ") o; ");
                pw.println("        return "
                  + fieldGetters(methodsPerAttribute)
                    .map(ExecutableElement::getSimpleName)
                    .map(Name::toString)
                    .sorted(NameFirstComparator.GET_ID_INSTANCE)
                    .map(v -> "Objects.equals(" + v + "(), other." + v + "())")
                    .collect(Collectors.joining("\n          && "))
                  + ";");
                pw.println("    }");
                pw.println("    @Override public int hashCode() {");
                pw.println("        return "
                  + (hasId
                    ? "(getId() == null ? super.hashCode() : getId().hashCode())"
                    : "Objects.hash("
                      + fieldGetters(methodsPerAttribute)
                        .filter(ee -> isImmutableFinalType(ee.getReturnType()))
                        .map(ExecutableElement::getSimpleName)
                        .map(Name::toString)
                        .sorted(NameFirstComparator.GET_ID_INSTANCE)
                        .map(v -> v + "()")
                        .collect(Collectors.joining(",\n          "))
                      + ")")
                  + ";");
                pw.println("    }");
                pw.println("    @Override public String toString() {");
                pw.println("        return String.format(\"%s@%08x\", " + (hasId ? "getId()" : "\"" + mapSimpleClassName + "\"" ) + ", System.identityHashCode(this));");
                pw.println("    }");

                // deepClone
                if (usingGeneratedCloner) {
                    pw.println("    private final DeepCloner cloner;");
                    pw.println("    public <V> V deepClone(V obj) {");
                    pw.println("        return cloner.from(obj);");
                    pw.println("    }");
                }

                // fields, getters, setters
                methodsPerAttribute.entrySet().stream().sorted(Comparator.comparing(Map.Entry::getKey, NameFirstComparator.ID_INSTANCE)).forEach(me -> {
                    HashSet<ExecutableElement> methods = me.getValue();
                    TypeMirror fieldType = determineFieldType(me.getKey(), methods);
                    if (fieldType == null) {
                        return;
                    }

                    pw.println("");
                    pw.println("    private " + fieldType + " f" + me.getKey() + ";");

                    for (ExecutableElement method : methods) {
                        FieldAccessorType fat = FieldAccessorType.determineType(method, me.getKey(), types, fieldType);
                        Optional<ExecutableElement> parentMethod = Util.findParentMethodImplementation(allParentMembers, method);

                        if (parentMethod.isPresent()) {
                            processingEnv.getMessager().printMessage(Kind.OTHER, "Method " + method + " is declared in a parent class.", method);
                        } else if (fat == FieldAccessorType.UNKNOWN || ! printMethodBody(pw, fat, method, "f" + me.getKey(), fieldType)) {
                            processingEnv.getMessager().printMessage(Kind.WARNING, "Could not determine desired semantics of method from its signature", method);
                        }
                    }
                });

                // Read-only class overrides setters to be no-op
                pw.println("    public static class Empty " + (an.inherits().isEmpty() ? "" : " extends " + an.inherits()) + " implements " + className + " {");
                pw.println("        public static final Empty INSTANCE = new Empty();");
                methodsPerAttribute.entrySet().stream().sorted(Comparator.comparing(Map.Entry::getKey, NameFirstComparator.ID_INSTANCE))
                  .map(Map.Entry::getValue)
                  .flatMap(Collection::stream)
                  .forEach(ee -> {
                    pw.println("        @Override "
                      + ee.getModifiers().stream().filter(m -> m != Modifier.ABSTRACT).map(Object::toString).collect(Collectors.joining(" "))
                      + " " + ee.getReturnType()
                      + " " + ee.getSimpleName()
                      + "(" + methodParameters(ee.getParameters()) + ") {");
                    if (ee.getReturnType().getKind() == TypeKind.VOID) {
                        pw.println("        }");
                    } else {
                        pw.println("            return null;");
                        pw.println("        }");
                    }
                });
                elements.getAllMembers(e).stream()
                  .filter(ee -> ee.getSimpleName().contentEquals("isUpdated"))
                  .filter(ExecutableElement.class::isInstance)
                  .map(ExecutableElement.class::cast)
                  .filter(ee -> ee.getReturnType().getKind() == TypeKind.BOOLEAN)
                  .forEach(ee -> {
                    pw.println("        @Override "
                      + ee.getModifiers().stream().filter(m -> m != Modifier.ABSTRACT).map(Object::toString).collect(Collectors.joining(" "))
                      + " " + ee.getReturnType()
                      + " " + ee.getSimpleName()
                      + "(" + methodParameters(ee.getParameters()) + ") {");
                    pw.println("            return false;");
                    pw.println("        }");
                });
                pw.println("    }");

                autogenerated.add("        EMPTY_INSTANCES.put(" + className + ".class, " + mapImplClassName + ".Empty.INSTANCE);");
                autogenerated.add("        CONSTRUCTORS_DC.put(" + className + ".class, " + mapImplClassName + "::new);");

                pw.println("}");
            }
        }

        private boolean printMethodBody(PrintWriter pw, FieldAccessorType accessorType, ExecutableElement method, String fieldName, TypeMirror fieldType) {
            TypeMirror firstParameterType = method.getParameters().isEmpty()
              ? types.getNullType()
              : method.getParameters().get(0).asType();
            TypeElement typeElement = elements.getTypeElement(types.erasure(fieldType).toString());

            switch (accessorType) {
                case GETTER:
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method + " {");
                    pw.println("        return " + fieldName + ";");
                    pw.println("    }");
                    return true;
                case SETTER:
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(" + firstParameterType + " p0) {");
                    if (! isImmutableFinalType(fieldType)) {
                        pw.println("        p0 = " + deepClone(firstParameterType, "p0") + ";");
                    }
                    if (isCollection(firstParameterType)) {
                        pw.println("        if (p0 != null) {");
                        pw.println("            " + removeUndefined(firstParameterType, "p0") + ";");
                        pw.println("            if (" + isUndefined("p0") + ") p0 = null;");
                        pw.println("        }");
                    }
                    pw.println("        updated |= ! Objects.equals(" + fieldName + ", p0);");
                    pw.println("        " + fieldName + " = p0;");
                    pw.println("    }");
                    return true;
                case COLLECTION_ADD:
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(" + firstParameterType + " p0) {");
                    if (! isImmutableFinalType(firstParameterType)) {
                        pw.println("        p0 = " + deepClone(firstParameterType, "p0") + ";");
                    }
                    if (isCollection(firstParameterType)) {
                        pw.println("        if (p0 != null) " + removeUndefined(firstParameterType, "p0") + ";");
                    }
                    pw.println("        if (" + isUndefined("p0") + ") return;");
                    pw.println("        if (" + fieldName + " == null) { " + fieldName + " = " + interfaceToImplementation(typeElement, "") + "; }");
                    if (isSetType(typeElement)) {
                        pw.println("        updated |= " + fieldName + ".add(p0);");
                    } else {
                        pw.println("        " + fieldName + ".add(p0);");
                        pw.println("        updated = true;");
                    }
                    pw.println("    }");
                    return true;
                case COLLECTION_DELETE:
                {
                    boolean needsReturn = method.getReturnType().getKind() != TypeKind.VOID;
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(" + firstParameterType + " p0) {");
                    pw.println("        if (" + fieldName + " == null) { return" + (needsReturn ? " false" : "") + "; }");
                    pw.println("        boolean removed = " + fieldName + ".remove(p0)" + ("java.util.Map".equals(typeElement.getQualifiedName().toString()) ? " != null" : "") + ";");
                    pw.println("        updated |= removed;");
                    if (needsReturn) pw.println("        return removed;");
                    pw.println("    }");
                    return true;
                }
                case COLLECTION_DELETE_BY_ID:
                {
                    boolean needsReturn = method.getReturnType().getKind() != TypeKind.VOID;
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(String p0) {");
                    pw.println("        boolean removed = " + fieldName + " != null && " + fieldName + ".removeIf(o -> Objects.equals(o." + getCollectionKey(fieldType, method) + ", p0));");
                    pw.println("        updated |= removed;");
                    if (needsReturn) pw.println("        return removed;");
                    pw.println("    }");
                    return true;
                }
                case COLLECTION_GET_BY_ID:
                {
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(String p0) {");
                    pw.println("        if (" + fieldName + " == null || " + fieldName + ".isEmpty()) return Optional.empty();");
                    pw.println("        return " + fieldName + ".stream().filter(o -> Objects.equals(o." + getCollectionKey(fieldType, method) + ", p0)).findFirst();");
                    pw.println("    }");
                    return true;
                }
                case MAP_ADD:
                    TypeMirror secondParameterType = method.getParameters().get(1).asType();
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(" + firstParameterType + " p0, " + secondParameterType + " p1) {");
                    if (! isImmutableFinalType(secondParameterType)) {
                        pw.println("        p1 = " + deepClone(secondParameterType, "p1") + ";");
                    }
                    if (isCollection(secondParameterType)) {
                        pw.println("        if (p1 != null) " + removeUndefined(secondParameterType, "p1") + ";");
                    }
                    pw.println("        boolean valueUndefined = " + isUndefined("p1") + ";");
                    pw.println("        if (valueUndefined) { if (" + fieldName + " != null) { updated |= " + fieldName + ".remove(p0) != null; } return; }");
                    pw.println("        if (" + fieldName + " == null) { " + fieldName + " = " + interfaceToImplementation(typeElement, "") + "; }");

                    pw.println("        Object v = " + fieldName + ".put(p0, p1);");
                    pw.println("        updated |= ! Objects.equals(v, p1);");
                    pw.println("    }");
                    return true;
                case MAP_GET:
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(" + firstParameterType + " p0) {");
                    pw.println("        return " + fieldName + " == null ? null : " + fieldName + ".get(p0);");
                    pw.println("    }");
                    return true;
            }

            return false;
        }
    }

    private class FieldDelegateGenerator implements Generator {

        @Override
        public void generate(TypeElement e) throws IOException {
            Map<String, HashSet<ExecutableElement>> methodsPerAttribute = methodsPerAttributeMapping(e);
            String className = e.getQualifiedName().toString();
            String packageName = null;
            int lastDot = className.lastIndexOf('.');
            if (lastDot > 0) {
                packageName = className.substring(0, lastDot);
            }

            String simpleClassName = className.substring(lastDot + 1);
            String mapClassName = className + "FieldDelegate";
            String mapSimpleClassName = simpleClassName + "FieldDelegate";
            String fieldsClassName = className + "Fields";

            GenerateEntityImplementations an = e.getAnnotation(GenerateEntityImplementations.class);
            TypeElement parentTypeElement = elements.getTypeElement((an.inherits() == null || an.inherits().isEmpty()) ? "void" : an.inherits());
            if (parentTypeElement == null) {
                return;
            }

            JavaFileObject file = processingEnv.getFiler().createSourceFile(mapClassName);
            IdentityHashMap<ExecutableElement, String> m2field = new IdentityHashMap<>();
            methodsPerAttribute.forEach((f, s) -> s.forEach(m -> m2field.put(m, f)));   // Create reverse map
            try (PrintWriter pw = new PrintWriterNoJavaLang(file.openWriter())) {
                if (packageName != null) {
                    pw.println("package " + packageName + ";");
                }

                generatedAnnotation(pw);
                pw.println("public class " + mapSimpleClassName + (an.inherits().isEmpty() ? "" : " extends " + an.inherits()) + " implements " + className + ", " + FQN_HAS_ENTITY_FIELD_DELEGATE + "<" + className + ">" + " {");
                pw.println("    private final " + FQN_ENTITY_FIELD_DELEGATE + "<" + className + "> entityFieldDelegate;");
                pw.println("    public " + mapSimpleClassName + "(" + FQN_ENTITY_FIELD_DELEGATE + "<" + className + "> entityFieldDelegate) {");
                pw.println("        this.entityFieldDelegate = entityFieldDelegate;");
                pw.println("    }");
                pw.println("    public " + FQN_ENTITY_FIELD_DELEGATE + "<" + className + "> getEntityFieldDelegate() {");
                pw.println("        return this.entityFieldDelegate;");
                pw.println("    }");

                pw.println("    @Override public boolean isUpdated() {");
                pw.println("        return entityFieldDelegate.isUpdated();");
                pw.println("    }");

                pw.println("    @Override public void markUpdatedFlag() {");
                pw.println("        entityFieldDelegate.markUpdatedFlag();");
                pw.println("    }");

                pw.println("    @Override public void clearUpdatedFlag() {");
                pw.println("        entityFieldDelegate.clearUpdatedFlag();");
                pw.println("    }");

                pw.println("    @Override public String toString() {");
                pw.println("        return \"%\" + String.valueOf(entityFieldDelegate);");
                pw.println("    }");

                getAllAbstractMethods(e)
                  .forEach(ee -> {
                      String originalField = m2field.get(ee);
                      if (originalField == null) {
                          return;
                      }
                      TypeMirror fieldType = determineFieldType(originalField, methodsPerAttribute.get(originalField));
                      String field = fieldsClassName + "." + toEnumConstant(originalField);

                      FieldAccessorType fat = FieldAccessorType.determineType(ee, originalField, types, fieldType);
                      printMethodBody(pw, fat, ee, field, fieldType);
                  });

                autogenerated.add("        ENTITY_FIELD_DELEGATE_CREATORS.put(" + className + ".class, (EntityFieldDelegateCreator<" + className + ">) " + mapClassName + "::new);");

                pw.println("}");
            }
        }

        private boolean printMethodBody(PrintWriter pw, FieldAccessorType accessorType, ExecutableElement method, String fieldName, TypeMirror fieldType) {
            TypeMirror firstParameterType = method.getParameters().isEmpty()
              ? types.getNullType()
              : method.getParameters().get(0).asType();

            switch (accessorType) {
                case GETTER:
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method + " {");
                    pw.println("        return (" + fieldType + ") entityFieldDelegate.get(" + fieldName + ");");
                    pw.println("    }");
                    return true;
                case SETTER:
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(" + firstParameterType + " p0) {");
                    pw.println("        entityFieldDelegate.set(" + fieldName + ", p0);");
                    pw.println("    }");
                    return true;
                case COLLECTION_ADD:
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(" + firstParameterType + " p0) {");
                    pw.println("        entityFieldDelegate.collectionAdd(" + fieldName + ", p0);");
                    pw.println("    }");
                    return true;
                case COLLECTION_DELETE:
                {
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(" + firstParameterType + " p0) {");
                    TypeElement fieldTypeElement = elements.getTypeElement(types.erasure(fieldType).toString());
                    String removeMethod = Util.isMapType(fieldTypeElement) ? "mapRemove" : "collectionRemove";
                    if (method.getReturnType().getKind() == TypeKind.VOID) {
                        pw.println("        entityFieldDelegate." + removeMethod + "(" + fieldName + ", p0);");
                    } else {
                        pw.println("        return (" + method.getReturnType() + ") entityFieldDelegate." + removeMethod + "(" + fieldName + ", p0);");
                    }
                    pw.println("    }");
                    return true;
                }
                case COLLECTION_DELETE_BY_ID:
                {
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(String p0) {");
                    if (method.getReturnType().getKind() == TypeKind.VOID) {
                        pw.println("        entityFieldDelegate.mapRemove(" + fieldName + ", p0);");
                    } else {
                        pw.println("        return (" + method.getReturnType() + ") entityFieldDelegate.mapRemove(" + fieldName + ", p0);");
                    }
                    pw.println("    }");
                    return true;
                }
                case MAP_ADD:
                    TypeMirror secondParameterType = method.getParameters().get(1).asType();
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(" + firstParameterType + " p0, " + secondParameterType + " p1) {");
                    pw.println("        entityFieldDelegate.mapPut(" + fieldName + ", p0, p1);");
                    pw.println("    }");
                    return true;
                case COLLECTION_GET_BY_ID:
                case MAP_GET:
                    pw.println("    @SuppressWarnings(\"unchecked\") @Override public " + method.getReturnType() + " " + method.getSimpleName() + "(" + firstParameterType + " p0) {");
                    pw.println("        return (" + method.getReturnType() + ") entityFieldDelegate.mapGet(" + fieldName + ", p0);");
                    pw.println("    }");
                    return true;
            }

            return false;
        }
    }

    private class DelegateGenerator implements Generator {
        @Override
        public void generate(TypeElement e) throws IOException {
            Map<String, HashSet<ExecutableElement>> methodsPerAttribute = methodsPerAttributeMapping(e);
            String className = e.getQualifiedName().toString();
            String packageName = null;
            int lastDot = className.lastIndexOf('.');
            if (lastDot > 0) {
                packageName = className.substring(0, lastDot);
            }

            String simpleClassName = className.substring(lastDot + 1);
            String mapClassName = className + "Delegate";
            String mapSimpleClassName = simpleClassName + "Delegate";
            String fieldsClassName = className + "Fields";

            GenerateEntityImplementations an = e.getAnnotation(GenerateEntityImplementations.class);
            TypeElement parentTypeElement = elements.getTypeElement((an.inherits() == null || an.inherits().isEmpty()) ? "void" : an.inherits());
            if (parentTypeElement == null) {
                return;
            }

            JavaFileObject file = processingEnv.getFiler().createSourceFile(mapClassName);
            IdentityHashMap<ExecutableElement, String> m2field = new IdentityHashMap<>();
            methodsPerAttribute.forEach((f, s) -> s.forEach(m -> m2field.put(m, f)));   // Create reverse map
            try (PrintWriter pw = new PrintWriterNoJavaLang(file.openWriter())) {
                if (packageName != null) {
                    pw.println("package " + packageName + ";");
                }

                generatedAnnotation(pw);
                pw.println("public class " + mapSimpleClassName + " implements " + className + ", org.keycloak.models.map.common.delegate.HasDelegateProvider<" + className + "> {");
                pw.println("    private final org.keycloak.models.map.common.delegate.DelegateProvider<" + className + "> delegateProvider;");
                pw.println("    public " + mapSimpleClassName + "(org.keycloak.models.map.common.delegate.DelegateProvider<" + className + "> delegateProvider) {");
                pw.println("        this.delegateProvider = delegateProvider;");
                pw.println("    }");
                pw.println("    public org.keycloak.models.map.common.delegate.DelegateProvider<" + className + "> getDelegateProvider() {");
                pw.println("        return this.delegateProvider;");
                pw.println("    }");

                pw.println("    @Override public String toString() {");
                pw.println("        return \"/\" + String.valueOf(this.delegateProvider);");
                pw.println("    }");

                getAllAbstractMethods(e)
                  .forEach(ee -> {
                      printMethodHeader(pw, ee);
                      String field = m2field.get(ee);
                      field = field == null ? "null" : fieldsClassName + "." + toEnumConstant(field);
                      if (ee.getReturnType().getKind() == TypeKind.BOOLEAN && "isUpdated".equals(ee.getSimpleName().toString())) {
                          pw.println("        return delegateProvider.isUpdated();");
                      } else if (ee.getReturnType().getKind() == TypeKind.VOID) {  // write operation
                          pw.println("        delegateProvider.getDelegate(false, "
                            + Stream.concat(Stream.of(field), ee.getParameters().stream().map(VariableElement::getSimpleName)).collect(Collectors.joining(", "))
                            + ")." + ee.getSimpleName() + "("
                            + ee.getParameters().stream().map(VariableElement::getSimpleName).collect(Collectors.joining(", "))
                            + ");");
                      } else {
                          pw.println("        return delegateProvider.getDelegate(true, "
                            + Stream.concat(Stream.of(field), ee.getParameters().stream().map(VariableElement::getSimpleName)).collect(Collectors.joining(", "))
                            + ")." + ee.getSimpleName() + "("
                            + ee.getParameters().stream().map(VariableElement::getSimpleName).collect(Collectors.joining(", "))
                            + ");");
                      }
                      pw.println("    }");
                  });

                pw.println("}");

                autogenerated.add("        DELEGATE_CREATORS.put(" + className + ".class, (DelegateCreator<" + className + ">) " + mapClassName + "::new);");
            }
        }
    }

    protected void printMethodHeader(final PrintWriter pw, ExecutableElement ee) {
        pw.println("    @Override "
          + ee.getModifiers().stream().filter(m -> m != Modifier.ABSTRACT).map(Object::toString).collect(Collectors.joining(" "))
          + " " + ee.getReturnType()
          + " " + ee.getSimpleName()
          + "(" + methodParameters(ee.getParameters()) + ") {");
    }

    private class ClonerGenerator implements Generator {

        @Override
        public void generate(TypeElement e) throws IOException {
            Map<String, HashSet<ExecutableElement>> methodsPerAttribute = methodsPerAttributeMapping(e);
            String className = e.getQualifiedName().toString();
            String packageName = null;
            int lastDot = className.lastIndexOf('.');
            if (lastDot > 0) {
                packageName = className.substring(0, lastDot);
            }

            String simpleClassName = className.substring(lastDot + 1);
            String clonerImplClassName = className + "Cloner";
            String clonerSimpleClassName = simpleClassName + "Cloner";

            JavaFileObject enumFile = processingEnv.getFiler().createSourceFile(clonerImplClassName);
            try (PrintWriter pw = new PrintWriter(enumFile.openWriter()) {
                @Override
                public void println(String x) {
                    super.println(x == null ? x : x.replaceAll("java.lang.", ""));
                }
            }) {
                if (packageName != null) {
                    pw.println("package " + packageName + ";");
                }
                pw.println("import " + FQN_DEEP_CLONER + ";");
                pw.println("// DO NOT CHANGE THIS CLASS, IT IS GENERATED AUTOMATICALLY BY " + GenerateEntityImplementationsProcessor.class.getSimpleName());
                generatedAnnotation(pw);
                pw.println("public class " + clonerSimpleClassName + " {");

                if (methodsPerAttribute.containsKey(ID_FIELD_NAME)) {
                    pw.println("    public static " + className + " deepClone(" + className + " original, " + className + " target) {");

                    // If the entity has an ID, set the ID first and then set all other attributes.
                    // This was important when working with Jpa storage as the ID is the one field needed to persist an entity.
                    HashSet<ExecutableElement> idMethods = methodsPerAttribute.get(ID_FIELD_NAME);
                    TypeMirror idFieldType = determineFieldType(ID_FIELD_NAME, idMethods);
                    cloneField(e, ID_FIELD_NAME, idMethods, idFieldType, pw);

                    pw.println("        return deepCloneNoId(original, target);");
                    pw.println("    }");

                    autogenerated.add("        CLONERS_WITH_ID.put(" + className + ".class, (Cloner<" + className + ">) " + clonerImplClassName + "::deepClone);");

                    pw.println("    public static " + className + " deepCloneNoId(" + className + " original, " + className + " target) {");

                    methodsPerAttribute.entrySet().stream().sorted(Comparator.comparing(Map.Entry::getKey)).forEach(me -> {
                        final String fieldName = me.getKey();
                        HashSet<ExecutableElement> methods = me.getValue();
                        TypeMirror fieldType = determineFieldType(fieldName, methods);
                        if (fieldType == null || ID_FIELD_NAME.equals(fieldName)) {
                            return;
                        }

                        cloneField(e, fieldName, methods, fieldType, pw);
                    });
                    pw.println("        target.clearUpdatedFlag();");
                    pw.println("        return target;");
                    pw.println("    }");

                    autogenerated.add("        CLONERS_WITHOUT_ID.put(" + className + ".class, (Cloner<" + className + ">) " + clonerImplClassName + "::deepCloneNoId);");
                } else {
                    pw.println("    public static " + className + " deepClone(" + className + " original, " + className + " target) {");

                    methodsPerAttribute.entrySet().stream().sorted(Comparator.comparing(Map.Entry::getKey)).forEach(me -> {
                        final String fieldName = me.getKey();
                        HashSet<ExecutableElement> methods = me.getValue();
                        TypeMirror fieldType = determineFieldType(fieldName, methods);
                        if (fieldType == null) {
                            return;
                        }

                        cloneField(e, fieldName, methods, fieldType, pw);
                    });
                    pw.println("        target.clearUpdatedFlag();");
                    pw.println("        return target;");
                    pw.println("    }");

                    autogenerated.add("        CLONERS_WITH_ID.put(" + className + ".class, (Cloner<" + className + ">) " + clonerImplClassName + "::deepClone);");
                }
                pw.println("}");
            }
        }

        private void cloneField(TypeElement e, final String fieldName, HashSet<ExecutableElement> methods, TypeMirror fieldType, final PrintWriter pw) {
            ExecutableElement getter = FieldAccessorType.getMethod(GETTER, methods, fieldName, types, fieldType).orElse(null);
            if (getter == null) {
                processingEnv.getMessager().printMessage(Kind.WARNING, "Could not determine getter for " + fieldName + " property");
                return;
            }

            Optional<ExecutableElement> setter = FieldAccessorType.getMethod(SETTER, methods, fieldName, types, fieldType);
            Optional<ExecutableElement> addToCollection = FieldAccessorType.getMethod(COLLECTION_ADD, methods, fieldName, types, fieldType);
            Optional<ExecutableElement> updateMap = FieldAccessorType.getMethod(MAP_ADD, methods, fieldName, types, fieldType);

            if (setter.isPresent()) {
                final Name setterName = setter.get().getSimpleName();
                // Setter always deep-clones whatever comes from the original, so we don't clone the value here.
                pw.println("        target." + setterName + "(original."  + getter.getSimpleName() + "());");
            } else if (addToCollection.isPresent()) {
                pw.println("        if (original." + getter.getSimpleName() + "() != null) {");
                pw.println("            original." + getter.getSimpleName() + "().forEach(target::" + addToCollection.get().getSimpleName() + ");");
                pw.println("        }");
            } else if (updateMap.isPresent()) {
                pw.println("        if (original." + getter.getSimpleName() + "() != null) {");
                pw.println("            original." + getter.getSimpleName() + "().forEach(target::" + updateMap.get().getSimpleName() + ");");
                pw.println("        }");
            } else {
                processingEnv.getMessager().printMessage(Kind.ERROR, "Could not determine way to clone " + fieldName + " property", e);
            }
        }
    }
}