【JSR269实战】之编译时操作AST，修改字节码文件，以实现和lombok类似的功能

笔者日常： 兄弟姐妹们，还是尽量少熬夜啊。我感觉我记性有所下降，难受。

---

需求说明(本文以实现此需求为例进行说明)：

  现在有一个需求，就是要给枚举类生成一个内部类，这个内部类中以静态常量的形式记录外部枚举类所有枚举项的值，即：

• 编译前java文件是这样的：
  
• (编译时操作AST，)编译后的class文件是这样的：
  

---

编译时操作AST，修改字节码文件：

软硬件环境说明： JDK1.8、Mavne3.6.3、IntelliJ IDEA。

项目整体(步骤)说明：

第一步：创建一个普通的maven项目，并在pom中引入相关依赖。

<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0modelVersion>
    <groupId>com.pingangroupId>
    <artifactId>jsr269-custom-astartifactId>
    <version>1.0.1version>
    <name>jsr269-custom-astname>
    <description>基于JSR269, 面向AST,自定义实现一个编译时修改字节码的类。具体功能为: generate inner-class containing outer-class's all
        public-static-final parameters
    description>

    <properties>
        <project.build.sourceEncoding>UTF-8project.build.sourceEncoding>
    properties>

    <dependencies>
        <dependency>
            <groupId>com.sungroupId>
            <artifactId>toolsartifactId>
            <version>1.8version>
            <scope>systemscope>
            <systemPath>${java.home}/../lib/tools.jarsystemPath>
        dependency>
    dependencies>

    <build>

        

        <resources>
            <resource>
                <directory>src/main/resourcesdirectory>
                <excludes>
                    <exclude>META-INF/**/*exclude>
                excludes>
            resource>
        resources>
        <plugins>
            <plugin>
                <groupId>org.apache.maven.pluginsgroupId>
                <artifactId>maven-compiler-pluginartifactId>
                <version>3.8.1version>
                <configuration>
                    
                    <source>1.8source>
                    <target>1.8target>
                configuration>
            plugin>

            <plugin>
                <groupId>org.apache.maven.pluginsgroupId>
                <artifactId>maven-resources-pluginartifactId>
                <version>3.1.0version>
                <executions>
                    <execution>
                        <id>process-METAid>
                        <phase>prepare-packagephase>
                        <goals>
                            <goal>copy-resourcesgoal>
                        goals>
                        <configuration>
                            <outputDirectory>target/classesoutputDirectory>
                            <resources>
                                <resource>
                                    <directory>${basedir}/src/main/resources/directory>
                                    <includes>
                                        <include>**/*include>
                                    includes>
                                resource>
                            resources>
                        configuration>
                    execution>
                executions>
            plugin>
        plugins>
    build>
project>

第二步：自定义编译时注解。

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * 生成一个公开的静态内部类， 这个内部类中持有了外部类的public-static-final常量
 *
 * @author JustryDeng
 * @date 2020/5/13 20:50:51
 */
@Target(ElementType.TYPE)
@Retention(RetentionPolicy.SOURCE)
public @interface EnumInnerConstant {

    /** 默认的内部类名 */
    String innerClassName() default "JustryDeng";
}

第三步：编写处理器。

import com.pingan.jsr269ast.annotation.EnumInnerConstant;
import com.sun.tools.javac.api.JavacTrees;
import com.sun.tools.javac.code.Flags;
import com.sun.tools.javac.code.TypeTag;
import com.sun.tools.javac.processing.JavacProcessingEnvironment;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.TreeMaker;
import com.sun.tools.javac.tree.TreeTranslator;
import com.sun.tools.javac.util.Context;
import com.sun.tools.javac.util.List;
import com.sun.tools.javac.util.Name;
import com.sun.tools.javac.util.Names;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import javax.lang.model.type.TypeKind;
import javax.tools.Diagnostic;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Pattern;

/**
 * generate inner-class containing outer-class's all public-static-final parameters
 *
 * @author JustryDeng
 * @date 2020/5/13 20:53:30
 */
@SupportedSourceVersion(SourceVersion.RELEASE_8)
@SupportedAnnotationTypes("com.pingan.jsr269ast.annotation.EnumInnerConstant")
public class EnumInnerConstantProcessor extends AbstractProcessor {

    /** 消息记录器 */
    private Messager messager;

    /** 可将Element转换为JCTree的工具。(注: 简单的讲，处理AST, 就是处理一个又一个CTree) */
    private JavacTrees trees;

    /** JCTree制作器 */
    private TreeMaker treeMaker;

    /** 名字处理器*/
    private Names names;

    /** 内部类类名校验 */
    private static final String INNER_CLASS_NAME_REGEX = "[A-Z][A-Za-z0-9]+";
    private static final Pattern INNER_CLASS_NAME_PATTERN = Pattern.compile(INNER_CLASS_NAME_REGEX);

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        this.messager = processingEnv.getMessager();
        this.trees = JavacTrees.instance(processingEnv);
        Context context = ((JavacProcessingEnvironment) processingEnv).getContext();
        this.treeMaker = TreeMaker.instance(context);
        this.names = Names.instance(context);
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        messager.printMessage(Diagnostic.Kind.NOTE, "roundEnv -> " + roundEnv);
        // 获取被@EnumInnerConstant注解标记的所有元素(可能是类、变量、方法等等)
        Set<? extends Element> elementSet = roundEnv.getElementsAnnotatedWith(EnumInnerConstant.class);
        elementSet.forEach(element -> {
            /*
             * 存储参数名与参数值的map、存储参数名与参数类型的map、一个辅助计数的map
             * 

             * 注: 照理来说，这里是单线程的。但考虑到本人对AST的处理机制也不是很熟，为
             *     保证万无一失，这里直接用线程安全的类吧。
             */
            Map<String, Object> paramsNameValueMap = new ConcurrentHashMap<>(8);
            Map<String, JCTree.JCExpression> paramsNameTypeMap = new ConcurrentHashMap<>(8);
            Map<String, AtomicInteger> paramIndexHelper = new ConcurrentHashMap<>(4);

            // 获取到注解信息
            EnumInnerConstant annotation = element.getAnnotation(EnumInnerConstant.class);
            String originInnerClassName = annotation.innerClassName();
            // 内部类类名校验
            String innerClassName = checkInnerClassName(originInnerClassName);
            // 将Element转换为JCTree
            JCTree jcTree = trees.getTree(element);
            String className = (((JCTree.JCClassDecl)jcTree).sym).type.toString();
            String enumFlag = "enum";
            if (!enumFlag.equalsIgnoreCase(jcTree.getKind().name())) {
                // 为保证错误信息能在各种情况下都能被看到, 这里用多种方式记录错误信息
                String errorMessage = "@EnumInnerConstant only support enum-class, [" + className + "] is not supported";
                System.err.println(errorMessage);
                messager.printMessage(Diagnostic.Kind.ERROR, errorMessage);
                throw new RuntimeException(errorMessage);
            }
            /*
             * 通过JCTree.accept(JCTree.Visitor)访问JCTree对象的内部信息。
             *
             * JCTree.Visitor有很多方法，我们可以通过重写对应的方法,(从该方法的形参中)来获取到我们想要的信息:
             * 如: 重写visitClassDef方法， 获取到类的信息;
             *     重写visitMethodDef方法， 获取到方法的信息;
             *     重写visitVarDef方法， 获取到变量的信息;
             *     重写visitLabelled方法， 获取到常量的信息;
             *     重写visitBlock方法， 获取到方法体的信息;
             *     重写visitImport方法， 获取到导包信息;
             *     重写visitForeachLoop方法， 获取到for循环的信息;
             *     ......
             */
            jcTree.accept(new TreeTranslator() {

                @Override
                public void visitClassDef(JCTree.JCClassDecl jcClassDecl) {
                    // 不要放在 jcClassDecl.defs = jcClassDecl.defs.append(a);之后，否者会递归
                    super.visitClassDef(jcClassDecl);
                    // 生成内部类， 并将内部类写进去
                    JCTree.JCClassDecl innerClass = generateInnerClass(innerClassName, paramsNameValueMap, paramsNameTypeMap);
                    jcClassDecl.defs = jcClassDecl.defs.append(innerClass);
                }

                @Override
                public void visitVarDef(JCTree.JCVariableDecl jcVariableDecl) {
                    boolean isEnumConstant = className.equals(jcVariableDecl.vartype.type.toString());
                    if (!isEnumConstant) {
                        super.visitVarDef(jcVariableDecl);
                        return;
                    }
                    Name name = jcVariableDecl.getName();
                    String paramName = name.toString();
                    /*
                     * 枚举项本身也属于变量, 每个枚举项里面，可能还有变量。 这里继
                     * 续JCTree.accept(JCTree.Visitor)进入，访问这个枚举项的内部信息。
                     */
                    jcVariableDecl.accept(new TreeTranslator() {
                        @Override
                        public void visitLiteral(JCTree.JCLiteral jcLiteral) {
                            Object paramValue = jcLiteral.getValue();
                            if (paramValue == null) {
                                return;
                            }
                            TypeTag typetag = jcLiteral.typetag;
                            JCTree.JCExpression paramType;
                            if (isPrimitive(typetag)) {
                                // 如果是基本类型，那么可以直接生成
                                paramType = treeMaker.TypeIdent(typetag);
                            } else if (paramValue instanceof String) {
                                // 如果不是基本类型，那么需要拼接生成
                                paramType = generateJcExpression("java.lang.String");
                            } else {
                                return;
                            }
                            AtomicInteger atomicInteger = paramIndexHelper.get(paramName);
                            if (atomicInteger == null) {
                                atomicInteger = new AtomicInteger(0);
                                paramIndexHelper.put(paramName, atomicInteger);
                            }
                            int paramIndex = atomicInteger.getAndIncrement();
                            String key = paramName + "_" + paramIndex;
                            paramsNameTypeMap.put(key, paramType);
                            paramsNameValueMap.put(key, paramValue);
                            super.visitLiteral(jcLiteral);
                        }
                    });
                    super.visitVarDef(jcVariableDecl);
                }
            });
        });
        return false;
    }

    /**
     * 内部类类名 合法性校验
     *
     * @param innerClassName
     *            内部类类名
     * @return  校验后的内部类类名
     * @date 2020/5/17 13:45:27
     */
    private String checkInnerClassName (String innerClassName) {
        if (innerClassName == null || innerClassName.trim().length() == 0) {
            // 为保证错误信息能在各种情况下都能被看到, 这里用多种方式记录错误信息
            String errorMessage = "@EnumInnerConstant. inner-class-name cannot be empty";
            System.err.println(errorMessage);
            messager.printMessage(Diagnostic.Kind.ERROR, errorMessage);
            throw new RuntimeException(errorMessage);
        }
        innerClassName = innerClassName.trim();
        if (!INNER_CLASS_NAME_PATTERN.matcher(innerClassName).matches()) {
            // 为保证错误信息能在各种情况下都能被看到, 这里用多种方式记录错误信息
            String errorMessage = "@EnumInnerConstant. inner-class-name must match regex " + INNER_CLASS_NAME_REGEX;
            System.err.println(errorMessage);
            messager.printMessage(Diagnostic.Kind.ERROR, errorMessage);
            throw new RuntimeException(errorMessage);
        }
        return innerClassName;
    }

    /**
     * 判断typeTag是否属于基本类型
     *
     * @param typeTag
     *            typeTag
     * @return 是否属于基本类型
     * @date 2020/5/17 13:10:54
     */
    private boolean isPrimitive(TypeTag typeTag) {
        if (typeTag == null) {
            return false;
        }
        TypeKind typeKind;
        try {
            typeKind = typeTag.getPrimitiveTypeKind();
        } catch (Throwable e) {
            return false;
        }
        if (typeKind == null) {
            return false;
        }
        return typeKind.isPrimitive();
    }


    /**
     * 生成内部类
     *
     * @return 生成出来的内部类
     * @date 2020/5/16 15:43:56
     */
    private JCTree.JCClassDecl generateInnerClass(String innerClassName, Map<String, Object> paramsInfoMap,Map<String, JCTree.JCExpression> paramsNameTypeMap) {
        JCTree.JCClassDecl jcClassDecl1 = treeMaker.ClassDef(
                treeMaker.Modifiers(Flags.PUBLIC + Flags.STATIC),
                names.fromString(innerClassName),
                List.nil(),
                null,
                List.nil(),
                List.nil());
        List<JCTree.JCVariableDecl> collection = generateAllParameters(paramsInfoMap, paramsNameTypeMap);
        collection.forEach(x -> jcClassDecl1.defs = jcClassDecl1.defs.append(x));
        return jcClassDecl1;
    }

    /**
     * 生成参数
     *
     * @param paramNameValueMap
     *           参数名-参数值map
     * @param paramsNameTypeMap
     *           参数名-参数类型map
     *
     * @return  参数JCTree集合
     * @date 2020/5/16 15:44:47
     */
    private List<JCTree.JCVariableDecl> generateAllParameters(Map<String, Object> paramNameValueMap,
                                                              Map<String, JCTree.JCExpression> paramsNameTypeMap) {
        List<JCTree.JCVariableDecl> jcVariableDeclList = List.nil();
        JCTree.JCVariableDecl statement;
        if (paramNameValueMap != null && paramNameValueMap.size() != 0) {
            for (Map.Entry<String, Object> entry : paramNameValueMap.entrySet()) {
                // 定义变量
                statement = treeMaker.VarDef(
                        // 访问修饰符
                        treeMaker.Modifiers(Flags.PUBLIC + Flags.STATIC + Flags.FINAL),
                        // 参数名
                        names.fromString(entry.getKey()),
                        // 参数类型
                        paramsNameTypeMap.get(entry.getKey()),
                        // 参数值
                        treeMaker.Literal(entry.getValue()));

                jcVariableDeclList = jcVariableDeclList.append(statement);
            }
        }
        return jcVariableDeclList;
    }

    /**
     * 根据全类名获取JCTree.JCExpression
     *
     * 如: 类变量 public static final String ABC = "abc";中, String就需要
     *     调用此方法generateJCExpression("java.lang.String")进行获取。
     *      追注: 其余的复杂类型，也可以通过这种方式进行获取。
     *      追注: 对于基本数据类型，可以直接通过类TreeMaker.TypeIdent获得，
     *           如: treeMaker.TypeIdent(TypeTag.INT)可获得int的JCTree.JCExpression
     *
     * @param fullNameOfTheClass
     *            全类名
     * @return  全类名对应的JCTree.JCExpression
     * @date 2020/5/16 15:47:32
     */
    private JCTree.JCExpression generateJcExpression(String fullNameOfTheClass) {
        String[] fullNameOfTheClassArray = fullNameOfTheClass.split("\\.");
        JCTree.JCExpression expr = treeMaker.Ident(names.fromString(fullNameOfTheClassArray[0]));
        for (int i = 1; i < fullNameOfTheClassArray.length; i++) {
            expr = treeMaker.Select(expr, names.fromString(fullNameOfTheClassArray[i]));
        }
        return expr;
    }

}

第四步：配置META-INF/services/javax.annotation.processing.Processor文件，使编译项目时触发我们自定义的处理器。

提示： 容器启动时，会扫描jar包下的META-INF/services/里的文件并作相应解析。

测试一下：

1. 把上面编写的项目install到本地仓库。上面的项目的maven信息是这样的：
   

2. 打开一个新的项目，在新项目的pom.xml中引入依赖。
   

3. 在新项目中，创建一个模型,并使用我们自定义的注解。
   

4. 编译新项目，并查看class文件。
   

由此可见，编译时操作AST，修改字节码文件成功 ！

---

调试说明：

方式一：debug调试

提示： 此方式需要一个提供Processor的项目，以及一个使用该Processor的项目。

• 第一步： 对使用Processor的项目进行mvnDebug clean compile。
  
  注: 我这里是直接在开发工具里进行的mvn，你也可以在其它命令行窗口里面进行mvn。
• 第二步： 在我们的Processor项目中,配置Remote调试。
  
  
  

方式二：输出调试

提示： 此方式只需要一个Processor项目即可。

• 第一步： 先利用开发工具,对我们自定义的注解及处理器进行编译。
  
• 第二步： 使用第一步编译后的处理器协助编译，编译目标类。
  
• 第三步： 查看第二步的Messager输出日志，进行调试。

---

  到此为止，我们实现的功能就与lombok非常相近了，唯一差的一点是：lombok针对不同IDE有提供不同的插件，使得IDE能够识别到lombok编译后的内容 ！针对IDE插件的开发，个人兴趣不大，如果以后时间非常充足的话，我可能会花一点时间去搞一搞。

^_^ 如有不当之处，欢迎指正

^_^ 参考连接
         https://blog.mythsman.com/p…bf/
         https://www.jianshu.com…35a
         https://blog.csdn.net/a_zhenzhen…063
         https://blog.csdn.net/sunqu…274

^_^ 测试代码托管链接
         https://github.com/JustryDeng/CommonRep…

^_^ 本文已经被收录进《程序员成长笔记(七)》，笔者JustryDeng