自定义Java注解处理器：实现将Context.save调用插入到注解变量定义后而非方法末尾

这个问题我之前帮别人排查过，核心就是没找对代码插入的精准位置！你之前的实现应该是默认把生成的代码加到了方法的最后，而没有定位到被@Remember标注的变量定义语句的后面。下面分两种常用的注解处理器实现方式给你解决方案：

一、基于Javac源码API的处理方案（推荐，源码级修改更直观） #

如果你是用JDK自带的javax.annotation.processing.AbstractProcessor做源码级处理，那可以通过Javac的Tree API精准定位变量位置，然后插入代码：

步骤1：初始化Tree工具类 #

在处理器的init方法中获取Trees实例，用于操作语法树：

private Trees trees;

@Override
public synchronized void init(ProcessingEnvironment processingEnv) {
    super.init(processingEnv);
    trees = Trees.getInstance(processingEnv);
}

步骤2：定位变量并插入代码 #

在process方法中，遍历所有被@Remember标注的局部变量，找到它们在方法体中的位置，然后紧跟插入Context.save调用：

@Override
public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
    Set<? extends Element> rememberElements = roundEnv.getElementsAnnotatedWith(Remember.class);
    for (Element element : rememberElements) {
        if (element.getKind() != ElementKind.LOCAL_VARIABLE) {
            continue; // 只处理局部变量
        }
        VariableElement varElement = (VariableElement) element;
        Remember rememberAnno = varElement.getAnnotation(Remember.class);
        String key = rememberAnno.value();
        String varName = varElement.getSimpleName().toString();

        // 获取变量对应的语法树路径
        TreePath varPath = trees.getPath(varElement);
        VariableTree varTree = (VariableTree) varPath.getLeaf();
        // 获取方法体的Block节点
        BlockTree methodBody = (BlockTree) varPath.getParentPath().getLeaf();
        List<? extends StatementTree> originalStatements = methodBody.getStatements();

        // 找到变量语句在方法体中的索引
        int varIndex = originalStatements.indexOf(varTree);
        if (varIndex == -1) {
            continue;
        }

        // 构建Context.save的调用语句
        TreeMaker treeMaker = trees.getTreeMaker();
        ExpressionTree saveCall = treeMaker.MethodInvocation(
                List.nil(), // 泛型参数
                treeMaker.Select(
                        treeMaker.Ident(trees.getName("Context")),
                        trees.getName("save")
                ),
                List.of(
                        treeMaker.Literal(key),
                        treeMaker.Ident(trees.getName(varName))
                )
        );
        ExpressionStatementTree saveStatement = treeMaker.ExpressionStatement(saveCall);

        // 插入到变量语句的下一行
        List<StatementTree> newStatements = new ArrayList<>(originalStatements);
        newStatements.add(varIndex + 1, saveStatement);

        // 替换原方法体
        MethodTree originalMethod = (MethodTree) varPath.getParentPath().getParentPath().getLeaf();
        MethodTree newMethod = treeMaker.Method(
                originalMethod.getModifiers(),
                originalMethod.getName(),
                originalMethod.getReturnType(),
                originalMethod.getTypeParameters(),
                originalMethod.getParameters(),
                originalMethod.getThrows(),
                treeMaker.Block(methodBody.getModifiers().getFlags(), newStatements),
                originalMethod.getDefaultValue()
        );

        // 更新类中的方法定义
        JCTree.JCCompilationUnit compilationUnit = (JCTree.JCCompilationUnit) varPath.getCompilationUnit();
        JCTree.JCClassDecl classDecl = (JCTree.JCClassDecl) compilationUnit.getTypeDecls().get(0);
        List<JCTree.JCStatement> classMembers = classDecl.getDefs();
        for (int i = 0; i < classMembers.size(); i++) {
            if (classMembers.get(i).equals(originalMethod)) {
                classMembers.set(i, (JCTree.JCStatement) newMethod);
                break;
            }
        }
    }
    return true;
}

这段代码的核心是通过originalStatements.indexOf(varTree)找到变量在方法体语句列表中的位置，然后在varIndex + 1处插入新语句，确保紧跟在变量定义之后。

二、基于ASM的字节码处理方案 #

如果你是用ASM直接修改字节码，那需要找到变量存储指令的位置，然后插入Context.save的调用：

核心实现代码 #

public class RememberClassVisitor extends ClassVisitor {
    public RememberClassVisitor(ClassVisitor classVisitor) {
        super(Opcodes.ASM9, classVisitor);
    }

    @Override
    public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
        MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
        return new RememberMethodVisitor(mv);
    }

    private static class RememberMethodVisitor extends MethodVisitor {
        public RememberMethodVisitor(MethodVisitor mv) {
            super(Opcodes.ASM9, mv);
        }

        @Override
        public AnnotationVisitor visitLocalVariableAnnotation(int typeRef, TypePath typePath, Label[] start, Label[] end, int[] index, String descriptor, boolean visible) {
            // 检查是否是@Remember注解
            if ("Lcom/example/Remember;".equals(descriptor)) {
                return new RememberAnnotationVisitor(mv, start, end, index);
            }
            return super.visitLocalVariableAnnotation(typeRef, typePath, start, end, index, descriptor, visible);
        }
    }

    private static class RememberAnnotationVisitor extends AnnotationVisitor {
        private final MethodVisitor mv;
        private final int[] varIndexes;
        private String key;

        public RememberAnnotationVisitor(MethodVisitor mv, Label[] start, Label[] end, int[] varIndexes) {
            super(Opcodes.ASM9);
            this.mv = mv;
            this.varIndexes = varIndexes;
        }

        @Override
        public void visit(String name, Object value) {
            if ("value".equals(name)) {
                key = (String) value; // 获取注解的key值
            }
            super.visit(name, value);
        }

        @Override
        public void visitEnd() {
            // 在变量定义的位置插入Context.save调用
            for (int index : varIndexes) {
                // 加载key字符串
                mv.visitLdcInsn(key);
                // 加载变量值（这里假设是引用类型，基本类型需要调整指令）
                mv.visitVarInsn(Opcodes.ALOAD, index);
                // 调用Context.save(String, Object)
                mv.visitMethodInsn(Opcodes.INVOKESTATIC, 
                        "com/example/Context", 
                        "save", 
                        "(Ljava/lang/String;Ljava/lang/Object;)V", 
                        false);
            }
            super.visitEnd();
        }
    }
}

注意事项 #

• 要确保ClassReader保留局部变量的注解信息，创建ClassReader时需要传入ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES以外的选项，比如ClassReader.EXPAND_FRAMES。
• 如果变量是基本类型，需要把ALOAD换成对应的ILOAD/LLOAD等指令，同时注意自动装箱（因为Context.save的第二个参数是Object）。

总结 #

不管用哪种方式，核心都是精准定位被标注变量的定义位置，然后紧跟其后插入Context.save调用。你之前的问题就是没有做这个定位，而是默认把代码加到了方法末尾，按照上面的方案修改后就能达到你想要的效果了。

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