CustomMemberGatherer.java

/*
 *    Copyright 2006-2026 the original author or authors.
 *
 *    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
 *
 *       https://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.mybatis.generator.merge.java;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Stream;

import com.github.javaparser.ast.body.BodyDeclaration;
import com.github.javaparser.ast.body.EnumConstantDeclaration;
import com.github.javaparser.ast.body.TypeDeclaration;
import com.github.javaparser.ast.comments.Comment;
import com.github.javaparser.ast.expr.AnnotationExpr;
import com.github.javaparser.ast.expr.Expression;
import com.github.javaparser.ast.expr.MemberValuePair;
import com.github.javaparser.ast.expr.StringLiteralExpr;
import org.mybatis.generator.api.MyBatisGenerator;
import org.mybatis.generator.config.MergeConstants;

/**
 * Inspect a TypeDeclaration looking for members that should be merged. We only look for immediate children of the
 * TypeDeclaration, we do not recurse into nested types. There are three resulting lists of members:
 *
 * <ol>
 *     <li>customBodyMembers: members that should be merged with the target file because they are not generated by
 *     MyBatis Generator</li>
 *     <li>doNotDeleteBodyMembers: members that were originally generated by MyBatis Generator, but marked
 *     "do_not_delete_during_merge". This means that any matching member in the newly generated file should be removed,
 *     and these members should be merged. Currently, the only known use of this feature is the legacy Example classes.
 *     Those classes have a "Criteria" inner class that can be modified as an extension point.</li>
 *     <li>customEnumConstants: enum constants that should be merged with the target file because they are not
 *     generated by MyBatis Generator</li>
 * </ol>
 *
 * <p>If all lists are empty, there is no need to merge any members into a new file, and the new file can be used as is.
 */
public class CustomMemberGatherer {
    private final List<BodyDeclaration<?>> customBodyMembers = new ArrayList<>();
    private final List<BodyDeclaration<?>> doNotDeleteBodyMembers = new ArrayList<>();
    private final List<EnumConstantDeclaration> customEnumConstants = new ArrayList<>();

    public CustomMemberGatherer(TypeDeclaration<?> typeDeclaration) {
        typeDeclaration.getMembers().forEach(this::gatherCustomBodyMemberIfNeeded);

        if (typeDeclaration.isEnumDeclaration()) {
            typeDeclaration.asEnumDeclaration().getEntries().forEach(this::gatherCustomEnumConstantIfNeeded);
        }
    }

    public boolean hasAnyMembersToMerge() {
        return !customBodyMembers.isEmpty() || !doNotDeleteBodyMembers.isEmpty() || !customEnumConstants.isEmpty();
    }

    public Stream<BodyDeclaration<?>> allCustomBodyMembers() {
        return Stream.of(customBodyMembers.stream(), doNotDeleteBodyMembers.stream())
                .flatMap(Function.identity());
    }

    public Stream<BodyDeclaration<?>> doNotDeleteBodyMembers() {
        return doNotDeleteBodyMembers.stream();
    }

    public Stream<EnumConstantDeclaration> customEnumConstants() {
        return customEnumConstants.stream();
    }

    private void gatherCustomBodyMemberIfNeeded(BodyDeclaration<?> member) {
        // Generated annotation:
        //  - Generated Keep - add to do not delete list
        //  - Generated remove - return
        //  - Not Generated - check for old Javadoc comments

        GeneratedType generatedType = checkForGeneratedAnnotation(member);
        if (generatedType == GeneratedType.GENERATED_KEEP) {
            doNotDeleteBodyMembers.add(member);
            return;
        } else if (generatedType == GeneratedType.GENERATED_REMOVE) {
            return;
        }

        generatedType = checkForGeneratedJavadocTag(member);
        if (generatedType == GeneratedType.NOT_GENERATED) {
            customBodyMembers.add(member);
        } else if (generatedType == GeneratedType.GENERATED_KEEP) {
            doNotDeleteBodyMembers.add(member);
        }
    }

    private void gatherCustomEnumConstantIfNeeded(EnumConstantDeclaration member) {
        GeneratedType generatedType = checkForGeneratedAnnotation(member);
        if (generatedType == GeneratedType.GENERATED_KEEP) {
            customEnumConstants.add(member);
            return;
        } else if (generatedType == GeneratedType.GENERATED_REMOVE) {
            return;
        }

        generatedType = checkForGeneratedJavadocTag(member);
        if (generatedType == GeneratedType.NOT_GENERATED || generatedType == GeneratedType.GENERATED_KEEP) {
            customEnumConstants.add(member);
        }
    }

    private GeneratedType checkForGeneratedAnnotation(BodyDeclaration<?> member) {
        return member.getAnnotations().stream()
                .filter(this::isOurGeneratedAnnotation)
                .findFirst()
                .map(a -> {
                    if (hasDoNotDeleteComment(a)) {
                        return GeneratedType.GENERATED_KEEP;
                    } else {
                        return GeneratedType.GENERATED_REMOVE;
                    }
                })
                .orElse(GeneratedType.NOT_GENERATED);
    }

    private boolean isOurGeneratedAnnotation(AnnotationExpr annotationExpr) {
        if (!isGeneratedAnnotation(annotationExpr)) {
            return false;
        }

        if (annotationExpr.isSingleMemberAnnotationExpr()) {
            Expression value = annotationExpr.asSingleMemberAnnotationExpr().getMemberValue();
            if (value.isStringLiteralExpr()) {
                return annotationValueMatchesMyBatisGenerator(value.asStringLiteralExpr());
            }
        } else if (annotationExpr.isNormalAnnotationExpr()) {
            return annotationExpr.asNormalAnnotationExpr().getPairs().stream()
                    .filter(this::isValuePair)
                    .map(MemberValuePair::getValue)
                    .filter(Expression::isStringLiteralExpr)
                    .map(Expression::asStringLiteralExpr)
                    .findFirst()
                    .map(this::annotationValueMatchesMyBatisGenerator)
                    .orElse(false);
        }

        return false;
    }

    private boolean hasDoNotDeleteComment(AnnotationExpr annotationExpr) {
        // check the comments value for the do_not_delete marker string
        if (annotationExpr.isSingleMemberAnnotationExpr()) {
            // no comments in a single member annotation - only the single "value" member"
            return false;
        } else if (annotationExpr.isNormalAnnotationExpr()) {
            return annotationExpr.asNormalAnnotationExpr().getPairs().stream()
                    .filter(this::isCommentsPair)
                    .map(MemberValuePair::getValue)
                    .filter(Expression::isStringLiteralExpr)
                    .map(Expression::asStringLiteralExpr)
                    .findFirst()
                    .map(StringLiteralExpr::asString)
                    .map(s -> s.contains(MergeConstants.DO_NOT_DELETE_DURING_MERGE))
                    .orElse(false);
        }

        return false;
    }

    private boolean isGeneratedAnnotation(AnnotationExpr annotationExpr) {
        String annotationName = annotationExpr.getNameAsString();
        // Check for @Generated annotation (both javax and jakarta packages)
        return "Generated".equals(annotationName) //$NON-NLS-1$
                || "javax.annotation.Generated".equals(annotationName) //$NON-NLS-1$
                || "jakarta.annotation.Generated".equals(annotationName); //$NON-NLS-1$
    }

    private boolean isValuePair(MemberValuePair pair) {
        return pair.getName().asString().equals("value"); //$NON-NLS-1$
    }

    private boolean isCommentsPair(MemberValuePair pair) {
        return pair.getName().asString().equals("comments"); //$NON-NLS-1$
    }

    private boolean annotationValueMatchesMyBatisGenerator(StringLiteralExpr expr) {
        return expr.asString().equals(MyBatisGenerator.class.getName());
    }

    private GeneratedType checkForGeneratedJavadocTag(BodyDeclaration<?> member) {
        return member.getComment()
                .map(Comment::getContent)
                .map(this::checkJavadocTag)
                .orElse(GeneratedType.NOT_GENERATED);
    }

    // Check if the comment contains any of the javadoc tags
    private GeneratedType checkJavadocTag(String comment) {
        for (String tag : MergeConstants.getOldElementTags()) {
            if (comment.contains(tag)) {
                if (comment.contains(MergeConstants.DO_NOT_DELETE_DURING_MERGE)) {
                    return GeneratedType.GENERATED_KEEP;
                } else {
                    return GeneratedType.GENERATED_REMOVE;
                }
            }
        }
        return GeneratedType.NOT_GENERATED;
    }

    private enum GeneratedType {
        NOT_GENERATED,
        GENERATED_REMOVE,
        GENERATED_KEEP
    }
}