聊聊如何通过APT+AST来实现AOP功能

前言

如果有使用过spring aop功能的小伙伴，应该都会知道spring aop主要是通过动态代理在运行时，对业务进行切面拦截操作。今天我们就来实现一下如何通过APT+AST在编译期时实现AOP功能。不过在此之前先科普一下APT和AST相关内容

AST(抽象语法树)

什么是AST

抽象语法树（Abstract Syntax Tree，AST），是源代码语法结构的一种抽象表示。它以树状的形式表现编程语言的语法结构，树上的每个节点都表示源代码中的一种结构。比如包、类型、修饰符、运算符、接口、返回值都可以是一个语法结构。

示例：

package com.example.adams.astdemo;
public class TestClass {
    int x = 0;
    int y = 1;
    public int testMethod(){
        int z = x + y;
        return z;
    }
}

对应的抽象语法树如下：

java的编译过程

重点关注步骤一和步骤二生成AST的过程

步骤一：词法分析，将源代码的字符流转变为 Token 列表。

通过词法分析器分析源文件中的所有字符，将所有的单词或字符都转化成符合规范的Token

规范化的token可以分成一下三种类型：

java关键字：public, static, final, String, int等等；
自定义的名称：包名，类名，方法名和变量名；
运算符或者逻辑运算符等符号：+、-、*、/、&&，|| 等等。

步骤二： 语法分析，根据 Token 流来构造树形表达式也就是 AST。

语法树的每一个节点都代表着程序代码中的一个语法结构，如类型、修饰符、运算符等。经过这个步骤后，编译器就基本不会再对源码文件进行操作了，后续的操作都建立在抽象语法树之上。

AST的应用场景

AST 定义了代码的结构，通过操作 AST，我们可以精准地定位到声明语句、赋值语句、运算语句等，实现对源代码的分析、优化、变更等操作。

注： AST操作属于编译器级别，对程序运行完全没有影响，效率相对其他AOP更高

java抽象语法树常用API类介绍

JCTree

JCTree 是语法树元素的基类，包含一个重要的字段 pos，该字段用于指明当前语法树节点（JCTree）在语法树中的位置，因此我们不能直接用 new 关键字来创建语法树节点，即使创建了也没有意义。

重点介绍几个JCTree的子类：

1、JCStatement：声明语法树节点，常见的子类如下

• JCBlock：语句块语法树节点
• JCReturn：return 语句语法树节点
• JCClassDecl：类定义语法树节点
• JCVariableDecl：字段 / 变量定义语法树节点

2、JCMethodDecl：方法定义语法树节点
3、JCModifiers：访问标志语法树节点
4、JCExpression：表达式语法树节点，常见的子类如下

• JCAssign：赋值语句
• JCAssignOp：+=
• JCIdent：标识符，可以是变量，类型，关键字等等
• JCLiteral: 字面量表达式，如123， “string”等
• JCBinary：二元操作符

JCTrees更多API的介绍可以查看如下链接

https://blog.csdn.net/u013998373/article/details/90050810

TreeMaker

TreeMaker 用于创建一系列的语法树节点，我们上面说了创建 JCTree 不能直接使用 new 关键字来创建，所以 Java 为我们提供了一个工具，就是 TreeMaker，它会在创建时为我们创建的 JCTree 对象设置 pos 字段，所以必须使用上下文相关的 TreeMaker 对象来创建语法树节点。

着重介绍一下常用的几个方法

TreeMaker.Modifiers

TreeMaker.Modifiers 方法用于创建访问标志语法树节点（JCModifiers），源码如下

public JCModifiers Modifiers(long flags) {
    return Modifiers(flags, List.< JCAnnotation >nil());
}

public JCModifiers Modifiers(long flags,
    List<JCAnnotation> annotations) {
        JCModifiers tree = new JCModifiers(flags, annotations);
        boolean noFlags = (flags & (Flags.ModifierFlags | Flags.ANNOTATION)) == 0;
        tree.pos = (noFlags && annotations.isEmpty()) ? Position.NOPOS : pos;
        return tree;
}

1. flags：访问标志
2. annotations：注解列表

其中 flags 可以使用枚举类 com.sun.tools.javac.code.Flags 来表示，例如我们可以这样用，就生成了下面的访问标志了。

示例：
创建访问修饰符 public

treeMaker.Modifiers(Flags.PUBLIC);

TreeMaker.ClassDef

TreeMaker.ClassDef 用于创建类定义语法树节点（JCClassDecl）, 源码如下：

public JCClassDecl ClassDef(JCModifiers mods,
    Name name,
    List<JCTypeParameter> typarams,
    JCExpression extending,
    List<JCExpression> implementing,
    List<JCTree> defs) {
        JCClassDecl tree = new JCClassDecl(mods,
                                     name,
                                     typarams,
                                     extending,
                                     implementing,
                                     defs,
                                     null);
        tree.pos = pos;
        return tree;
}

1. mods：访问标志，可以通过 TreeMaker.Modifiers 来创建
2. name：类名
3. typarams：泛型参数列表
4. extending：父类
5. implementing：实现的接口
6. defs：类定义的详细语句，包括字段、方法的定义等等

TreeMaker.MethodDef

TreeMaker.MethodDef 用于创建方法定义语法树节点（JCMethodDecl），源码如下

public JCMethodDecl MethodDef(JCModifiers mods,
    Name name,
    JCExpression restype,
    List<JCTypeParameter> typarams,
    List<JCVariableDecl> params,
    List<JCExpression> thrown,
    JCBlock body,
    JCExpression defaultValue) {
        JCMethodDecl tree = new JCMethodDecl(mods,
                                       name,
                                       restype,
                                       typarams,
                                       params,
                                       thrown,
                                       body,
                                       defaultValue,
                                       null);
        tree.pos = pos;
        return tree;
}

public JCMethodDecl MethodDef(MethodSymbol m,
    Type mtype,
    JCBlock body) {
        return (JCMethodDecl)
            new JCMethodDecl(
                Modifiers(m.flags(), Annotations(m.getAnnotationMirrors())),
                m.name,
                Type(mtype.getReturnType()),
                TypeParams(mtype.getTypeArguments()),
                Params(mtype.getParameterTypes(), m),
                Types(mtype.getThrownTypes()),
                body,
                null,
                m).setPos(pos).setType(mtype);
}

1. mods：访问标志
2. name：方法名
3. restype：返回类型
4. typarams：泛型参数列表
5. params：参数列表
6. thrown：异常声明列表
7. body：方法体
8. defaultValue：默认方法（可能是 interface 中的哪个 default）
9. m：方法符号
10. mtype：方法类型。包含多种类型，泛型参数类型、方法参数类型、异常参数类型、返回参数类型。

注： 返回类型 restype 填写 null 或者 treeMaker.TypeIdent(TypeTag.VOID) 都代表返回 void 类型

示例

创建方法

    public String getUserName(String userName){
       return userName;
    }

ListBuffer<JCTree.JCStatement> usernameStatement = new ListBuffer<>();
usernameStatement.append(treeMaker.Return(treeMaker.Ident(names.fromString("userName"))));
JCTree.JCBlock usernameBody = treeMaker.Block(0, usernameStatement .toList());

// 生成入参
JCTree.JCVariableDecl param = treeMaker.VarDef(treeMaker.Modifiers(Flags.PARAMETER), names.fromString("userName"),treeMaker.Ident(names.fromString("String")), null);
com.sun.tools.javac.util.List<JCTree.JCVariableDecl> parameters = com.sun.tools.javac.util.List.of(param);

JCTree.JCMethodDecl username = treeMaker.MethodDef(
        treeMaker.Modifiers(Flags.PUBLIC), 
        names.fromString("getUserName"), // 方法名
        treeMaker.Ident(names.fromString("String")), // 返回类型
        com.sun.tools.javac.util.List.nil(),
        parameters, // 入参
        com.sun.tools.javac.util.List.nil(),
        usernameBody ,
        null
);

TreeMaker.VarDef

TreeMaker.VarDef 用于创建字段 / 变量定义语法树节点（JCVariableDecl），源码如下

public JCVariableDecl VarDef(JCModifiers mods,
    Name name,
    JCExpression vartype,
    JCExpression init) {
        JCVariableDecl tree = new JCVariableDecl(mods, name, vartype, init, null);
        tree.pos = pos;
        return tree;
}

public JCVariableDecl VarDef(VarSymbol v,
    JCExpression init) {
        return (JCVariableDecl)
            new JCVariableDecl(
                Modifiers(v.flags(), Annotations(v.getAnnotationMirrors())),
                v.name,
                Type(v.type),
                init,
                v).setPos(pos).setType(v.type);
}

1. mods：访问标志
2. name：参数名称
3. vartype：类型
4. init：初始化语句
5. v：变量符号

示例：

创建变量： private String username = “lyb-geek”;

treeMaker.VarDef(treeMaker.Modifiers(Flags.PRIVATE), names.fromString("useranme"), treeMaker.Ident(names.fromString("String"),  treeMaker.Literal("lyb-geek");

TreeMaker.Ident

TreeMaker.Ident 用于创建标识符语法树节点（JCIdent）可以表示类、变量引用或者方法。源码如下

public JCIdent Ident(Name name) {
        JCIdent tree = new JCIdent(name, null);
        tree.pos = pos;
        return tree;
}

public JCIdent Ident(Symbol sym) {
        return (JCIdent)new JCIdent((sym.name != names.empty)
                                ? sym.name
                                : sym.flatName(), sym)
            .setPos(pos)
            .setType(sym.type);
}

public JCExpression Ident(JCVariableDecl param) {
        return Ident(param.sym);
}

示例：

创建username的引用

treeMaker.Ident(names.fromString("username"))))

TreeMaker.Return

TreeMaker.Return 用于创建 return 语句（JCReturn），源码如下

public JCReturn Return(JCExpression expr) {
        JCReturn tree = new JCReturn(expr);
        tree.pos = pos;
        return tree;
}

示例
return this.username

treeMaker.Return(treeMaker.Select(treeMaker.Ident(names.fromString("this")),names.fromString("useranme")));

TreeMaker.Select

TreeMaker.Select 用于创建域访问 / 方法访问（这里的方法访问只是取到名字，方法的调用需要用 TreeMaker.Apply）语法树节点（JCFieldAccess），源码如下

public JCFieldAccess Select(JCExpression selected,
    Name selector) 
{
        JCFieldAccess tree = new JCFieldAccess(selected, selector, null);
        tree.pos = pos;
        return tree;
}

public JCExpression Select(JCExpression base,
    Symbol sym) {
        return new JCFieldAccess(base, sym.name, sym).setPos(pos).setType(sym.type);
}

1. selected：. 运算符左边的表达式
2. selector：. 运算符右边的表达式

示例

获取方法logDTO.setArgs（）

treeMaker.Select(treeMaker.Ident(getNameFromString("logDTO")),
                        getNameFromString("setArgs")

TreeMaker.NewClass

TreeMaker.NewClass 用于创建 new 语句语法树节点（JCNewClass）, 源码如下：

public JCNewClass NewClass(JCExpression encl,
    List<JCExpression> typeargs,
    JCExpression clazz,
    List<JCExpression> args,
    JCClassDecl def) {
        JCNewClass tree = new JCNewClass(encl, typeargs, clazz, args, def);
        tree.pos = pos;
        return tree;
}

1. encl：不太明白此参数的含义，我看很多例子中此参数都设置为 null
2. typeargs：参数类型列表
3. clazz：待创建对象的类型
4. args：参数列表
5. def：类定义

示例：

创建 List args = new ArrayList();

 JCTree.JCNewClass argsListclass = treeMaker.NewClass(null, null, memberAccess("java.util.ArrayList"), List.nil(), null);

        
        JCTree.JCVariableDecl args = makeVarDef(treeMaker.Modifiers(0),
                memberAccess("java.util.List"),
                "args",
                argsListclass
        );

TreeMaker.Apply

TreeMaker.Apply 用于创建方法调用语法树节点（JCMethodInvocation），源码如下：

public JCMethodInvocation Apply(List<JCExpression> typeargs,
    JCExpression fn,
    List<JCExpression> args) {
        JCMethodInvocation tree = new JCMethodInvocation(typeargs, fn, args);
        tree.pos = pos;
        return tree;
}

1. typeargs：参数类型列表
2. fn：调用语句
3. args：参数列表

TreeMaker.Assign
TreeMaker.Assign 用户创建赋值语句语法树节点（JCAssign），源码如下：

public JCAssign Assign(JCExpression lhs,
    JCExpression rhs) {
        JCAssign tree = new JCAssign(lhs, rhs);
        tree.pos = pos;
        return tree;
}

1. lhs：赋值语句左边表达式
2. rhs：赋值语句右边表达式

示例
创建 username = “lyb-geek”

treeMaker.Assign(treeMaker.Ident(names.fromString("username")))), treeMaker.Literal("lyb-geek"))

TreeMaker.Exec

TreeMaker.Exec 用于创建可执行语句语法树节点（JCExpressionStatement），源码如下：

public JCExpressionStatement Exec(JCExpression expr) {
        JCExpressionStatement tree = new JCExpressionStatement(expr);
        tree.pos = pos;
        return tree;
}

注： TreeMaker.Apply 以及 TreeMaker.Assign 就需要外面包一层 TreeMaker.Exec 来获得一个 JCExpressionStatement

示例：

username = “lyb-geek”

treeMaker.Exec(treeMaker.Assign(treeMaker.Ident(names.fromString("username")),treeMaker.Binary(JCTree.Tag.PLUS,treeMaker.Literal("lyb"),treeMaker.Literal("-geek"))))

TreeMaker.Block

TreeMaker.Block 用于创建组合语句的语法树节点（JCBlock），源码如下：

public JCBlock Block(long flags,
    List<JCStatement> stats) {
        JCBlock tree = new JCBlock(flags, stats);
        tree.pos = pos;
        return tree;
}

1. flags：访问标志
2. stats：语句列表

示例

创建代码块

List<JCTree.JCStatement> jcStatementList = List.nil();
treeMaker.Block(0, jcStatementList);

TreeMaker更多详细API可以查看如下链接
http://www.docjar.com/docs/api/com/sun/tools/javac/tree/TreeMaker.html

Names

Names封装了操作标识符的方法,类、方法、参数的名称都可以通过names来获取

大家如果对AST感兴趣，可以通过https://astexplorer.net/在线体验一下

实战

示例主要通过APT+AST实现一个统计方法调用耗时以及记录日志的功能

注： 大家可以通过JavaParserJavaParser来简化对AST的操作。

本示例通过jdk自带的tools.jar工具类进行操作

1、在pom引入tools.jar gav

  <dependency>
            <groupId>com.sun</groupId>
            <artifactId>tools</artifactId>
            <version>1.8</version>
            <scope>system</scope>
            <systemPath>${java.home}/../lib/tools.jar</systemPath>
        </dependency>

2、自定义注解CostTimeRecoder

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.SOURCE)
@Documented
public @interface CostTimeRecoder {
}

3、编写注解处理器

@AutoService(Processor.class)
@SupportedOptions("debug")
public class CostTimeRecordProcessor extends AbstractComponentProcessor {


    /**
     * 元素辅助类
     */
    private Elements elementUtils;

    /**
     * 日志输出工具类
     */
    private Messager meessager;

    /**
     * 抽象语法树
     */
    private JavacTrees trees;

    /**
     * 封装了创建或者修改AST节点的一些方法
     */
    private TreeMaker treeMaker;

    /**
     * 封装了操作标识符的方法
     */
    private Names names;



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


    @Override
    public Set<String> getSupportedAnnotationTypes() {
        return Collections.singleton(CostTimeRecoder.class.getName());
    }


    @Override
    protected boolean processImpl(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {

        if (annotations == null || annotations.isEmpty()) {
            return false;
        }


        Set<? extends Element> elements = roundEnv.getElementsAnnotatedWith(CostTimeRecoder.class);
        if (elements == null || elements.isEmpty()){
            return false;
        }

        if (!roundEnv.processingOver()) {
            elements.stream() .filter(element -> element instanceof ExecutableElement)
                    .map(element -> (ExecutableElement) element)
                    .forEach(method -> {
                        TypeElement typeElement = (TypeElement)method.getEnclosingElement();
                        JCTree.JCClassDecl tree = trees.getTree(typeElement);
                        JCTree.JCMethodDecl methodDecl = trees.getTree(method);
                        CostTimeRecordAstTranslator costTimeRecordAstTranslator = new CostTimeRecordAstTranslator(treeMaker,names,meessager,tree,methodDecl);
                        costTimeRecordAstTranslator.setTrees(trees);
                        // 导入引用类,如果不配置import，则方法调用，需配置全类路径，
                        // 比如LogFactory.getLogger(),如果没导入LogFactory，则方法需写成com.github.lybgeek.log.factory.LogFactory.getLogger
                        // 配置后，仅需写成LogFactory.getLogger即可
                        costTimeRecordAstTranslator.addImportInfo(typeElement, LogFactory.class.getPackage().getName(),LogFactory.class.getSimpleName());
                        costTimeRecordAstTranslator.addImportInfo(typeElement,LogDTO.class.getPackage().getName(),LogDTO.class.getSimpleName());
//                        costTimeRecordAstTranslator.addImportInfo(typeElement, LogService.class.getPackage().getName(),LogService.class.getSimpleName());

                        tree.accept(costTimeRecordAstTranslator);

                    });
        }



        return false;
    }


    private String getPackageName(TypeElement typeElement) {
        return elementUtils.getPackageOf(typeElement).getQualifiedName()
                .toString();
    }




}

3、编写AST TreeTranslator

注：省略业务的TreeTranslator，就列出基类，可能对大家比较有用，需要业务的实现方法，直接见下方demo链接

public abstract class AbstractTreeTranslator extends TreeTranslator {

    /**
     * 封装了创建或者修改AST节点的一些方法
     */
    protected TreeMaker treeMaker;

    /**
     * 封装了操作标识符的方法
     */
    protected Names names;

    /**
     * 日志输出工具类
     */
    protected Messager meessager;

    /**
     * 抽象语法树
     */
    private JavacTrees trees;

    public AbstractTreeTranslator(TreeMaker treeMaker, Names names, Messager meessager) {
        this.treeMaker = treeMaker;
        this.names = names;
        this.meessager = meessager;
    }


    /**
     * 根据字符串获取Name
     * @param s
     * @return
     */
    public Name getNameFromString(String s) { return names.fromString(s); }


    /**
     * 创建变量语句
     * @param modifiers 访问修饰符
     * @param name 参数名称
     * @param varType 参数类型
     * @param init 初始化赋值语句
     * 示例
     *   JCTree.JCVariableDecl var = makeVarDef(treeMaker.Modifiers(0), "xiao", memberAccess("java.lang.String"), treeMaker.Literal("methodName"));
     *   生成语句为：String xiao = "methodName";
     * @return
     */
    public JCTree.JCVariableDecl makeVarDef(JCTree.JCModifiers modifiers, JCTree.JCExpression varType,String name, JCTree.JCExpression init) {
        return treeMaker.VarDef(
                modifiers,
                getNameFromString(name),
                varType,
                init
        );
    }

    /**
     * 创建 域/方法 的多级访问, 方法的标识只能是最后一个
     * @param components 比如java.lang.System.out.println
     * @return
     */
    public JCTree.JCExpression memberAccess(String components) {
        String[] componentArray = components.split("\\.");
        JCTree.JCExpression expr = treeMaker.Ident(getNameFromString(componentArray[0]));
        for (int i = 1; i < componentArray.length; i++) {
            expr = treeMaker.Select(expr, getNameFromString(componentArray[i]));
        }
        return expr;
    }

    /**
     * 给变量赋值
     * @param lhs
     * @param rhs
     * @return
     * 示例：makeAssignment(treeMaker.Ident(getNameFromString("xiao")), treeMaker.Literal("assignment test"));
     * 生成的赋值语句为：xiao = "assignment test";
     */
    public JCTree.JCExpressionStatement makeAssignment(JCTree.JCExpression lhs, JCTree.JCExpression rhs) {
        return treeMaker.Exec(
                treeMaker.Assign(
                        lhs,
                        rhs
                )
        );
    }

    /**
     * 导入方法依赖的package包
     * @param packageName
     * @param className
     * @return
     */
    public JCTree.JCImport buildImport(String packageName, String className) {
        JCTree.JCIdent ident = treeMaker.Ident(names.fromString(packageName));
        JCTree.JCImport jcImport = treeMaker.Import(treeMaker.Select(
                ident, names.fromString(className)), false);
         meessager.printMessage(Diagnostic.Kind.NOTE,jcImport.toString());
        return jcImport;
    }

    /**
     * 导入方法依赖的package包
     * @param element  class
     * @param packageName
     * @param className
     * @return
     */
    public void addImportInfo(TypeElement element, String packageName, String className) {
        TreePath treePath = getTrees().getPath(element);
        Tree leaf = treePath.getLeaf();
        if (treePath.getCompilationUnit() instanceof JCTree.JCCompilationUnit && leaf instanceof JCTree) {
            JCTree.JCCompilationUnit jccu = (JCTree.JCCompilationUnit) treePath.getCompilationUnit();

            for (JCTree jcTree : jccu.getImports()) {
                if (jcTree != null && jcTree instanceof JCTree.JCImport) {
                    JCTree.JCImport jcImport = (JCTree.JCImport) jcTree;
                    if (jcImport.qualid != null && jcImport.qualid instanceof JCTree.JCFieldAccess) {
                        JCTree.JCFieldAccess jcFieldAccess = (JCTree.JCFieldAccess) jcImport.qualid;
                        try {
                            if (packageName.equals(jcFieldAccess.selected.toString()) && className.equals(jcFieldAccess.name.toString())) {
                                return;
                            }
                        } catch (NullPointerException e) {
                            e.printStackTrace();
                        }
                    }
                }
            }
            java.util.List<JCTree> trees = new ArrayList<>();
            trees.addAll(jccu.defs);
            JCTree.JCImport jcImport = buildImport(packageName,className);
            if (!trees.contains(jcImport)) {
                trees.add(0, jcImport);
            }
            jccu.defs = List.from(trees);
        }
    }


    public JavacTrees getTrees() {
        return trees;
    }

    public void setTrees(JavacTrees trees) {
        this.trees = trees;
    }
}

4、测试

编写测试类

public class HelloService {

    @CostTimeRecoder
    public String sayHello(String username){
        try {
            TimeUnit.SECONDS.sleep(3);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        return "hello : " + username;
    }
}

测试主类

public class AptAstMainTest {

    public static void main(String[] args) {


        System.out.println(new HelloService().sayHello("zhangsan"));


    }
}

运行查看控制台

会发现多了耗时，以及日志打印。我们查看HelloService .class文件，会发现多了如下内容

public class HelloService {
    public HelloService() {
    }

    public String sayHello(String username) {
        Long startTime = System.currentTimeMillis();

        try {
            TimeUnit.SECONDS.sleep(3L);
        } catch (InterruptedException var7) {
            var7.printStackTrace();
        }

        Long endTime = System.currentTimeMillis();
        Long costTime = endTime - startTime;
        String msg = String.format("costTime = %s(ms)", costTime);
        System.out.println(msg);
        List args = new ArrayList();
        args.add(username);
        this.saveLog(costTime, args);
        return "hello : " + username;
    }

    private void saveLog(Long costTime, List args) {
        LogDTO logDTO = new LogDTO();
        logDTO.setMethodName("sayHello");
        logDTO.setClassName("com.github.lybgeek.test.service.HelloService");
        logDTO.setCostTime(costTime);
        logDTO.setArgs(args);
        LogService logService = LogFactory.getLogger();
        logService.save(logDTO);
    }
}

总结

本文主要重点介绍AST的用法，对AOP的实现基本上是一笔带过。原因主要是平时除非是对性能有特别要求，我们实现AOP通常会在运行期实现，而非在编译期实现。其次AST比较偏底层，如果出问题，排查难度会比较高。当然如果团队有对AST很熟悉的话，能兼顾性能是最好的。