如何在注解处理器中为不可变类自动生成可变方法的异常重写？

针对你的需求——在编译时修改AST，为不可变实现类自动生成父类/接口中标记@Behaviour(MUTABLE)方法的重写（逻辑为抛出异常），可以通过Javac插件实现（普通注解处理器无法修改已有类的AST，只能生成新文件），以下是具体步骤和代码实现：

一、核心思路 #

1. 识别目标类：筛选出标记@Implementation(type=IMPLEMENTATION)且@EngineNature(behaviour=IMMUTABLE)的不可变实现类。
2. 收集MUTABLE方法：递归遍历目标类的父类、实现的接口，收集所有被@Behaviour(MUTABLE)标记的方法。
3. 修改AST：为每个未被重写的MUTABLE方法生成带@Override的重写方法，方法体抛出UnsupportedOperationException。
4. 集成编译流程：将插件注册到Javac编译过程中，在AST解析完成后执行修改逻辑。

二、代码实现 #

1. 编写Javac插件类 #

实现com.sun.source.util.Plugin接口，注册AST处理器：

import com.sun.source.util.JavacTask;
import com.sun.source.util.Plugin;
import com.sun.source.util.TaskEvent;
import com.sun.source.util.TaskListener;

public class ImmutableMethodGeneratorPlugin implements Plugin {
    @Override
    public String getName() {
        // 插件名称，编译时通过-Xplugin参数指定
        return "immutable-method-generator";
    }

    @Override
    public void init(JavacTask task, String... args) {
        // 注册任务监听器，在PARSE阶段完成后处理AST
        task.addTaskListener(new TaskListener() {
            @Override
            public void started(TaskEvent e) {}

            @Override
            public void finished(TaskEvent e) {
                if (e.getKind() == TaskEvent.Kind.PARSE) {
                    ImmutableMethodVisitor visitor = new ImmutableMethodVisitor(task);
                    visitor.scan(e.getCompilationUnit(), null);
                }
            }
        });
    }
}

2. 编写AST遍历与修改的Visitor #

继承TreePathScanner实现类的遍历、方法收集与AST修改：

import com.sun.tools.javac.api.JavacTaskImpl;
import com.sun.tools.javac.code.Flags;
import com.sun.tools.javac.code.Types;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.JCTreeMaker;
import com.sun.tools.javac.tree.TreePathScanner;
import com.sun.tools.javac.util.List;
import com.sun.tools.javac.util.Names;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeMirror;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Set;

class ImmutableMethodVisitor extends TreePathScanner<Void, Void> {
    private final JavacTask task;
    private final Types typeUtils;
    private final Elements elementUtils;
    private final com.sun.source.util.Trees trees;

    public ImmutableMethodVisitor(JavacTask task) {
        this.task = task;
        try {
            this.typeUtils = task.getTypes();
            this.elementUtils = task.getElements();
            this.trees = com.sun.source.util.Trees.instance(task);
        } catch (Exception e) {
            throw new RuntimeException("Failed to initialize visitor", e);
        }
    }

    @Override
    public Void visitClass(JCTree.JCClassDecl classTree, Void unused) {
        Element classElement = trees.getElement(getCurrentPath());
        if (!(classElement instanceof TypeElement typeElement)) {
            return super.visitClass(classTree, unused);
        }

        // 筛选不可变实现类
        Implementation implAnnotation = typeElement.getAnnotation(Implementation.class);
        EngineNature natureAnnotation = typeElement.getAnnotation(EngineNature.class);
        if (implAnnotation == null || implAnnotation.type() != Implementation.Type.IMPLEMENTATION
                || natureAnnotation == null || natureAnnotation.behaviour() != Behaviour.Type.IMMUTABLE) {
            return super.visitClass(classTree, unused);
        }

        // 收集所有MUTABLE方法
        Set<ExecutableElement> mutableMethods = collectMutableMethods(typeElement);

        // 生成并添加重写方法
        ArrayList<JCTree> newMembers = new ArrayList<>(classTree.defs);
        for (ExecutableElement method : mutableMethods) {
            if (isMethodAlreadyOverridden(typeElement, method)) {
                continue;
            }
            newMembers.add(generateOverrideMethod(method));
        }
        classTree.defs = List.from(newMembers);

        return super.visitClass(classTree, unused);
    }

    // 检查当前类是否已重写该方法
    private boolean isMethodAlreadyOverridden(TypeElement typeElement, ExecutableElement method) {
        for (Element enclosed : typeElement.getEnclosedElements()) {
            if (enclosed instanceof ExecutableElement currentMethod) {
                if (typeUtils.isOverrideEquivalent(currentMethod, method)) {
                    return true;
                }
            }
        }
        return false;
    }

    // 递归收集父类/接口中的MUTABLE方法
    private Set<ExecutableElement> collectMutableMethods(TypeElement typeElement) {
        Set<ExecutableElement> mutableMethods = new HashSet<>();
        // 遍历超类
        TypeMirror superType = typeElement.getSuperclass();
        while (!superType.toString().equals(Object.class.getName())) {
            TypeElement superClass = (TypeElement) typeUtils.asElement(superType);
            addMutableMethodsFromType(superClass, mutableMethods);
            superType = superClass.getSuperclass();
        }
        // 遍历接口及父接口
        for (TypeMirror interfaceType : typeElement.getInterfaces()) {
            TypeElement interfaceElement = (TypeElement) typeUtils.asElement(interfaceType);
            addMutableMethodsFromType(interfaceElement, mutableMethods);
            collectInterfaceMethods(interfaceElement, mutableMethods);
        }
        return mutableMethods;
    }

    private void collectInterfaceMethods(TypeElement interfaceElement, Set<ExecutableElement> mutableMethods) {
        for (TypeMirror parentInterface : interfaceElement.getInterfaces()) {
            TypeElement parentElement = (TypeElement) typeUtils.asElement(parentInterface);
            addMutableMethodsFromType(parentElement, mutableMethods);
            collectInterfaceMethods(parentElement, mutableMethods);
        }
    }

    private void addMutableMethodsFromType(TypeElement typeElement, Set<ExecutableElement> mutableMethods) {
        for (Element enclosed : typeElement.getEnclosedElements()) {
            if (enclosed instanceof ExecutableElement method) {
                Behaviour behaviour = method.getAnnotation(Behaviour.class);
                if (behaviour != null && behaviour.value() == Behaviour.Type.MUTABLE) {
                    mutableMethods.add(method);
                }
            }
        }
    }

    // 生成重写方法的AST节点
    private JCTree.JCMethodDecl generateOverrideMethod(ExecutableElement method) {
        JavacTaskImpl javacTaskImpl = (JavacTaskImpl) task;
        JCTreeMaker treeMaker = javacTaskImpl.getContext().get(JCTreeMaker.class);
        Names names = javacTaskImpl.getContext().get(Names.class);

        // 生成@Override注解
        JCTree.JCAnnotation overrideAnn = treeMaker.Annotation(
                treeMaker.Ident(names.fromString("Override")),
                List.nil()
        );

        // 方法修饰符：public + @Override
        JCTree.JCModifiers modifiers = treeMaker.Modifiers(Flags.PUBLIC, List.of(overrideAnn));

        // 方法名称、返回类型
        JCTree.JCIdent methodName = treeMaker.Ident(names.fromString(method.getSimpleName().toString()));
        JCTree.JCExpression returnType = treeMaker.Type((JCTree.JCType) method.getReturnType());

        // 方法参数
        List<JCTree.JCVariableDecl> params = List.nil();
        for (VariableElement param : method.getParameters()) {
            JCTree.JCVariableDecl paramDecl = treeMaker.VarDef(
                    treeMaker.Modifiers(Flags.PARAMETER),
                    names.fromString(param.getSimpleName().toString()),
                    treeMaker.Type((JCTree.JCType) param.asType()),
                    null
            );
            params = params.append(paramDecl);
        }

        // 方法异常声明
        List<JCTree.JCExpression> exceptions = List.nil();
        for (TypeMirror exception : method.getThrownTypes()) {
            exceptions = exceptions.append(treeMaker.Type((JCTree.JCType) exception));
        }

        // 方法体：抛出UnsupportedOperationException
        JCTree.JCExpression exceptionClass = treeMaker.Ident(names.fromString(UnsupportedOperationException.class.getName()));
        JCTree.JCLiteral message = treeMaker.Literal(
                String.format("Mutable operation '%s' is not supported in immutable implementation", method.getSimpleName())
        );
        JCTree.JCThrow throwStmt = treeMaker.Throw(
                treeMaker.NewClass(null, List.nil(), exceptionClass, List.of(message), null)
        );
        JCTree.JCBlock body = treeMaker.Block(0, List.of(throwStmt));

        // 构建方法节点
        return treeMaker.MethodDef(
                modifiers, methodName, returnType, List.nil(), params, exceptions, body, null
        );
    }
}

3. 注册插件 #

在项目的src/main/resources/META-INF/services目录下创建文件com.sun.source.util.Plugin，内容为插件的全类名：

com.yourpackage.ImmutableMethodGeneratorPlugin

4. Maven编译配置 #

修改pom.xml的编译插件配置，启用Javac插件并引入依赖：

<build>
    <plugins>
        <plugin>
            <groupId>org.apache.maven.plugins</groupId>
            <artifactId>maven-compiler-plugin</artifactId>
            <version>3.11.0</version>
            <configuration>
                <source>17</source>
                <target>17</target>
                <compilerArgs>
                    <!-- 启用自定义Javac插件 -->
                    <arg>-Xplugin:immutable-method-generator</arg>
                </compilerArgs>
                <annotationProcessorPaths>
                    <!-- 包含你的注解处理器和插件所在模块 -->
                    <path>
                        <groupId>your.group.id</groupId>
                        <artifactId>data-engine-core</artifactId>
                        <version>1.0-SNAPSHOT</version>
                    </path>
                </annotationProcessorPaths>
            </configuration>
        </plugin>
    </plugins>
</build>

三、注意事项 #

1. 依赖JDK内部API：Javac插件依赖com.sun.tools.javac包，该包位于JDK的tools.jar中，Maven编译时会自动引入（无需额外配置），但需确保编译使用的是JDK而非JRE。
2. 兼容性：不同JDK版本的Javac内部API可能存在差异，该方案在OpenJDK17下测试可用。
3. 避免重复生成：代码中已加入“检查方法是否已重写”的逻辑，防止重复生成相同方法。

内容的提问来源于stack exchange，提问作者DevSw