Kotlin Compiler Plugin 探索与实践

653次阅读  |  发布于2年以前

1 . 背景

不知是否源于实现成本,网上有关编译器插件的资料少得可怜。实际上,Java 有实现无反射元编程的 Manifold,Kotlin 有备受推崇的 android-kotlin-synthetics 和 kapt 的替代品 google/ksp,这些项目无论从功能维度还是编译性能都是传统 codegen 方案(APT、Transform)不能比拟的。

结合日常开发中编码上的痛点,本人过了一遍 kotlin 编译相关源码(部分),照猫画虎手撸了个 Kotlin-Compiler-Plugin(后简称 KCP),落地后无论代码优雅程度还是编译性能都有了蛮大提升,故在此抛砖引玉一波,欢迎讨论~

2 . 实验及效果

看一个例子:

@ProtoMessage("webcast.data.GiftTip")
data class GiftTip(
    @SerializedName("display_text") @JvmField var displayText: Text? = null,
    @SerializedName("background_color") @JvmField var backgroundColor: String? = null,
    @SerializedName("prefix_image") @JvmField var perfixImage: ImageModel? = null
)

上面代码是一个带注解的 POJO,用途是承接网络请求数据,存在以下问题:

  1. 出于一些原因(保留无参构造、避免 NPE、适配反序列化等),小小的一个类、三个字段被加上共计 7 个注解,视觉干扰严重;
  2. 每个参数都是可空类型,实际用起来少不了一连串的问号或判空;
  3. 字段必须是 var,以防止 backingField 被标记为 final,进而影响(构造方法外非反射方式的)赋值;

上述问题,在原生 Kotlin 范畴内很难解决。在其语法规范中,除委托属性外,所有属性都必须被赋予默认值。一般默认值定义在声明之后,即写成 val foo = "bar"。在大多数情况下这样做无可厚非,但若类仅被用于反序列化,就显得多余了:

  1. 属性初始值应由原始数据决定,默认值则被用作字段缺省时的 fallback;
  2. 绝大部分字段的默认值是有共识的(number 为 0,struct 为 null 等),逐个指定一遍有冗余;
  3. 对于 data class,不应要求反序列化框架调用其有参构造,保留一个无参构造方法是必要的;

经过分析,这里的痛点是写法上不够优雅,可以借助 KCP 实现一个语法糖提高代码可读性。

此外,KAPT 进行注解处理也会拖慢编译速度,上述例子中 @ProtoMessage 是一个自定义注解,用来为 POJO 类生成 Protobuf 辅助类。KCP 在编译速度上显著优于 KAPT,因此可以使用 KCP 重写注解处理以提升性能

实验结果

通过 KCP 让代码更优雅的同时提高了编译速度,最终效果如下:

// 旧方案(kapt)
@ProtoMessage("webcast.xxx.TestMessage")
class TestKotlin {

    @SerializedName("room_id")
    @JvmField
    var roomId: Long = 0L

    @SerializedName("display_text")
    @JvmField
    var displayText: Text? = null

    @SerializedName("url_list")
    @JvmField
    var urlList: MutableList<String>? = null

    @SerializedName("nickname")
    @JvmField
    var nickname: String = "defaultName"
}
// 新方案(kcp)
@ProtoMessage("webcast.xxx.TestMessage")
class TestKotlin {

    val roomId by serialized<Long>()

    val displayText by serialized<Text>()

    val urlList by serialized<List<String>>()

    val nickname by serialized(defaultValue = "defaultName")
}

性能方面,工程编译耗时线下、线上分别减少 20% 和 16.4%(全模块,非增量,均值):

耗时\方案 kapt(旧方案) kcp(新方案)
本地测试 07:03 05:38
线上测试 10:35 08:51

接下来介绍一下具体的实现,优于具体的处理逻辑比较复杂,进一步展开前先卖个关子,介绍一下 Kotlin(JVM,非 IR)的编译流程。

3 . Kotlin(JVM) 编译简述

以下表述仅适用于面向 JVM 且使用非 IR 后端的情况(Kotlin 版本 < 1.5)。IR 后端在 Kotlin 版本来到 1.5 后成为默认选项。 To Learn More:The New JVM IR Backend Is Stable、Exploring Kotlin IR

Kotlin 编译流程示意如下:

左侧是 kotlinc 内部流程,分为 prepare、compile 两个步骤,其中 compile 又分为前端、后端。右侧是不同阶段的产物类型。

准备

背景介绍:Kotlin 在形成语法树之前的步骤均基于 intellij-core,与 IDE 逻辑是高度复用的,有兴趣可以看看 KotlinCoreEnvironment。

语法树与其他 IntelliJ 所支持的语言一样,基于 PSI 体系,有关介绍移步 What Is the PSI?

图中 Source Files 转换到 KtElements 的过程实际上也属于编译器前端的范畴,单独拆出来是因为其逻辑是与 IDE Plugin 共用的,源码工程中也是分开的。诸如 ASTNode 的生成、PsiElement 的转换等发生在这一步。

如此做法优势明显:如果有涉及语法、语义的改造(例如 android-kotlin-synthetics),开发者可以方便地让同一份代码复用在编写、编译两个环境中。

Kotlin 不愧是 Jetbrains 亲儿子,点个赞~

然而理想总是比现实丰满:想要在这里做扩展,我们需要付出更多成本。问题主要出在 IDE Plugin 的覆盖上,官方仓库的插件一般会直接集成到 IDE 中,而我们就没这么霸道了。一旦自定义的 IDE Plugin 未安装,相应语法就会失去代码联想等功能,并可能带来大量“爆红”。再考虑到前端扩展的难度会略大于后端,不建议做此类扩展(除非你的插件真的很 diao)。

这部分不过多展开,感兴趣移步 IntelliJ Platform SDK。

编译器前端

抽离了词法、语法分析,编译器前端只剩下很薄一层了。这里的主要工作就是遍历语法树,过程中发现、report 错误,并将 Elements 转换成 Descriptor 记录到 BindingContext 中。

所谓 Descriptor,即某个 Element 的“描述符”(对于 Kotlin 属性,包含名称、返回类型、是否 inlinegettersetter 等信息)。有了它,元素的信息检索就不再需要在语法树上游走了。而所谓 BindingContext,就是这些描述符的索引,本质上是 Map 的 Map,写成伪代码类似于 Map<Type, Map<Key, Descriptor>>

上述过程可以用下图表示:

不同颜色的实心圆圈是实际位于语法树中的 PsiElementKtElement),实线线段表示了语法树形状。虚线圆圈表示 CallReference 等元素间的非直接关系,用虚线连接。这些 PsiElementCallReference 等最后都会被转换成 Descriptor,最终汇聚在 BindingContext 中(每个“堆叠形状”是一类 Map,比如 Property 集中在一个 Map,Class 集中在另一个 Map)。

编译器后端

backend 的职责很简单:结合 BindingContext 将 Module 的内容以字节码形式输出到文件。流程是从 KtFile 出发,沿语法树边走边生成。具体代码就不展开分析了,有兴趣可以从 DefaultCodegenFactory#generateModule 看起。

下面列几个重要角色,理解它们有助于快速熟悉流程,上手 KCP 编写。

MemberCodegen<T : KtPureElement>

这是生成器的入口,泛型参数是面向的 Element 类型,调用其 generate() 方法即可生成对应 Element 的字节码。绝大多是时候我们用到的是 ImplementationBodyCodegen,面向 KtPureClassOrObject(即 class 或 object)。

MemberCodegen 的核心逻辑大多位于 generateBody 方法,构造方法、成员属性、成员方法等一系列元素的字节码生成均由这里触发。

同时,这个类包含了一系列常用的工具,下面捡常用的说说。

ClassBuilder

顾名思义,这是实际产生字节码的地方。之所以要有 ClassBuilder 这层抽象,除了衔接两套代码,还提供了更高的代码扩展性。除输出字节码的 AbstractClassBuilder.Concrete 外,大量 Kotlin 内置的字节码优化、字节码兼容以及部分 KCP 的逻辑都依靠代理 ClassBuilder 实现。

工具类

注意:这里不会把 Kotlin 内置类型(kotlin.Any、kotlin.String 等)转换为实际的 JVM 类型,那是后面需要做的事情。

4 . 实操

回顾一下预期:使用简短的语法,标记当前字段的类型、序列化后名称、缺省值,最好有一定扩展性。

下面介绍方案的细节由来和实现。

语法设计

原先的方案中,我们依赖 apt/kapt 做代码生成,解码的过程与 JSON 解析类似,先 new 一个对象出来,再给各字段赋值。这样带来一些问题:

  1. 涉及字段不能为 final,因为赋值在构造方法后进行;
  2. 涉及字段必须为可空类型或带有默认值;
  3. 涉及字段不能为 private,因为赋值操作发生在外部类;

不考虑实现的前提下,对于来自反序列化的数据类,字段最优雅的形态一定是 public final:开箱即用,不用关心数据的具体解析方式。那么问题来了,如何做到这一点呢?

最初的方案比较立项,是用自定义 DSL 全面替代字段声明,达到语法上的最简,类似这样:

class TestKotlin : SerializedModel("webcast.xxx.TestKotlin", {
    long("userId")
    int("id")
    string("nickName")
})

这个方案有两个缺陷:

  1. 与原方案差异过大,有一定学习成本;
  2. 强依赖 IDE 插件辅助(为当前类添加 member),缺失会引起大量爆红;

尤其是第二点非常痛,为避免日后影响新人开发体验,只能换一个方案,想到了使用委托。

委托

在 JVM 中,注解是官方提供的最简单有效的标记元素的方式,也是现阶段用的最多的姿势。然而在我们的场景中,注解太多且重复,会使代码显得十分冗长。因此不如另辟蹊径,关注 Kotlin 的语言特性:委托。

无论简洁程度还是表达的语义,委托都非常适合我们的场景:定义属性时,将 getter、setter 交给被委托者实现,背后的逻辑对使用方隐藏。至于原先的注解值,可以放在委托方法的参数中声明。

如果不引入编译器插件,委托到 serialized 方法的字段反编译成 Java 类似这样:

public final class Test {
    @NotNull private final DecodeDelegate test$delegate = KtxKt.serialized("test");

    @NotNull public final String getTest() { /* invoke getValue function and return */ }
}

可见,相较于普通属性,委托属性在字节码层面区别不大。其同样有一个 backingField,区别体现在 getter/setter 上。我们要做的是另建一个 backingField(类型与属性类型一致)并干掉原有的,最后重写 getter/setter。相当于把一个委托属性退化成普通的成员属性。

一波设计后,用于委托的方法如下:

fun <R> serialized(key: String = "") = DecodeDelegate<R>() // 无默认值

fun <R> serialized(key: String = "", defaultValue: R) = DecodeDelegate<R>() // 有默认值

/**
 * Delegate 实现(Stub),正常情况下会被编译器插件替换,方法不被调用
 */
class DecodeDelegate<R> {
    inline operator fun getValue(thisRef: Any?, property: KProperty<*>): R = error("stub")
    inline operator fun setValue(thisRef: Any?, property: KProperty<*>, value: R): R = error("stub")
}

实际效果:

// 源码
class Test {
    val test by serialized<String>()
}
// 反编译后
public final class Test {
    @SerializedName("test)
    protected String _$$test;

    public final @NotNull String getTest() { return _$$test; }
}

为支持类继承,所有 backingField 可见性都是 protected。同时为避免打包到 aar 后误调用,字段名均经过混淆。

至此,我们“偷天换日”,得到了一堆“黑盒”的委托属性。

数据解析

大功告成了吗?并不。我们还没有解决委托属性背后数据的来源问题:所有属性都需要支持从 JSON 和 Protobuf 反序列化。

JSON 好办,给每个隐藏的 backingField 加上 @SerializedName,保留一个无参构造方法,其它交给 Gson 即可。

Protobuf 处理起来会棘手一些。为了不引入反射,我们决定把字段赋值放在构造方法内,即 新增一个以 ProtoReader (用于流式解析,类似 Gson 的 JsonReader 的作用)为入参的隐藏构造方法。

最后,二进制产物反编译成 Java 类似这样:

public final class Test {
    @SerializedName("test")
    protected String _$$test;

    public final @NotNull String getTest() { return _$$test; }

    public Test(@NotNull ProtoReader) {
        for(/*...*/) {
            switch(tag) {
                case 0: _$$test = /* decode */;
            }
        }
        if (_$$test == null) {
            _$$test = "";
        }
    }
}

核心逻辑 & 实现

总结一下,按顺序,编译时插件要做以下事情:

接下来我们会用到三个 EP:

  1. AnalysisHandlerExtension
  2. ClassBuilderInterceptorExtension
  3. ExpressionCodegenExtension

后文会逐一介绍各 EP 的实现。

信息收集 & 发现错误

AnalysisHandlerExtension 的回调发生在编译器前端,主要面向 PsiElement 和收集到的 Descriptor,可以用来干预、监听源码文件分析的过程和结果。“臭名昭著”的 kapt、“备受期待”的 ksp 均在此展开。

接口大概长这样:

interface AnalysisHandlerExtension {
    /*...*/

    fun doAnalysis(/*...*/): AnalysisResult? = null

    fun analysisCompleted(/*...*/): AnalysisResult? = null
}

为尽早发现潜在的语法错误,我们将字段与 IDL 的匹配检查放在 analysisCompleted 方法进行。这里要做三件事:

  1. 按照 IDL 定义,解析所有可能用到的 proto 文件,拿到 message 集合;
  2. 根据收集的 message,检查目标类的声明是否合法;
  3. 收集各字段的声明信息,为后续代码生成做准备;

第一步代码照搬现有方案,不再赘述;

第二步的检查规则同样与现有方案相同,不再赘述。这一步看起繁琐,实则效率很高。考虑到 analysisCompleted 时已有完整的 BindingContext,这里遍历分析的速度极快(大型模块 1~2ms 的量级),耗时可几乎忽略不计。

第三步复杂一些,我们为每个字段创建一个 SerializeConfig 如下:

open class SerializeConfig constructor(
    val key: String,
    val strategy: Strategy
) {
    sealed class Strategy {

        object Optional : Strategy()

        sealed class WithDefaultValue : Strategy() {

            class Expr(val resolvedCall: ResolvedCall<out CallableDescriptor>) : WithDefaultValue()

            class Const(val constantValue: ConstantValue<*>) : WithDefaultValue()
        }

        object Annotated : Strategy()
    }

    /*...*/
}

SerializeConfig 的核心是 Strategy,决定了该字段的解码策略。我们按场景把策略分为 4 类,结合代码差异如下:

class Test {
    // 字段无缺省值。如果类型可空则默认为空,否则若未被赋值,调用 getter 时抛出异常
    val simpleField by serialized<String>()

    // 字段有缺省值,为编译期常量。若反序列化后未被赋值,会用 defaultValue 填充。
    val constField by serialized<String>(defaultValue = "hello world")

    // 字段有缺省值,需在运行时求值。若反序列化后未被赋值,会调用 defaultValue 的表达式求值并填充。
    val dynamicField by serialized<String>(defaultValue = createEmptyString())

    // 对旧方案的兼容
    @SerializedName("xxx")
    var legacyField: String = ""
}

一个有意思的点:当 defaultValue 不是编译时常量时,其表达式只有在反序列化未赋值时才会被调用,达到类似 lazy 的效果。

可见,策略由字段的特征决定,除兼容策略外,主要取决于 serialized 方法的入参。这里我们可以通过 BindingContext.DELEGATED_PROPERTY_RESOLVED_CALL 来获取某成员属性的委托方法调用,拿到的类型是 ResolvedCall<FunctionDescriptor>。从命名可以看出,该对象包含“此方法调用”的全部信息,我们需要的是其入参信息,即名为 defaultValue 的入参的表达式。

上面实际上要对得到的 ResolvedCall<FunctionDescriptor> 再找一次 ResolvedCall。第一次拿到的是对 operator fun getValue 的调用,第二次才是 serialized 的调用。

若调用方传递了此参数,则表达式一定为 ExpressionValueArgument。这时再通过 getCompileTimeConstant 工具方法判断其是否可以转换为编译时常量即可。

代理 ClassBuilder

从这一步起,我们要开始干“实事儿”了。通过实现 ClassBuilderInterceptorExtension,我们可以方便地代理内置的 ClassBuilder

几个重要方法需要实现:

public class DelegatingClassBuilder implements ClassBuilder {

    // 类定义走到这里
    void defineClass(
            @Nullable PsiElement origin,
            int version,
            int access,
            @NotNull String name,
            @Nullable String signature,
            @NotNull String superName,
            @NotNull String[] interfaces
    ) {...}

    // 字段定义走到这里(注意和 property 的差异)
    @NotNull
    FieldVisitor newField(
            @NotNull JvmDeclarationOrigin origin,
            int access,
            @NotNull String name,
            @NotNull String desc,
            @Nullable String signature,
            @Nullable Object value
    ) {...}

    // 方法定义走到这里(包括 getter、setter 等)
    @NotNull
    MethodVisitor newMethod(
            @NotNull JvmDeclarationOrigin origin,
            int access,
            @NotNull String name,
            @NotNull String desc,
            @Nullable String signature,
            @Nullable String[] exceptions
    ) {...}
}

这是一个典型的代理模式,若不打算干预,直接返回方法的 super 实现即可。几个返回值名字看着眼熟,点进去有惊喜:就是 ASM 那套,换了包名而已。

这里我们通过 BindingContext 拿到字段或方法的归属,拦截委托给 serialized 方法的字段,阻止其产生任何产物(真正的产物会在后面的步骤手动插入)。

考虑到字节码是 FieldVisitorMethodVisitor 在遍历元素时产生的,我们在 newField 和 newMethod 中选择性地返回空实现 Visitor 即可。注意:在实际过滤逻辑中,我们不仅要判断元素的归属,还要防止把后面代码生成时追加的部分误删除。

另外,由于后面还要生成额外的构造方法,为了在运行时辨别类的身份,我们还要在 defineClass 中为类签名上追加一个接口(marker interface),运行时检测到是该接口实现,就可以直接调用特定的构造方法。

代码生成

代码生成是最核心的部分,逻辑复杂,分为三大部分工作

  1. 写入委托字段实际的 backingFields 和 accessors;
  2. 改写已存在的 constructors;
  3. 写入用于解码 Protobuf 的 constructor;

第二步为什么要修改已有的构造方法呢?

对于 Kotlin class,除了 init 代码块外,primary constructor 和带有初始值的 property 都会影响构造方法的内容。而在我们的场景下,不仅要考虑字段未被反序列化时的缺省值,也要考虑该对象被其他构造方法创建时的情况。因此,所有手写的构造方法都需要被重新生成(注意不是修改,因为修改已存在的字节码需要考虑的因素太多,难度过大)。过程不复杂,可以遵循如下顺序:

  1. 调用 super;
  2. 将 primary constructor 的参数赋值到字段(处理构造方法中的 var/val);
  3. 提取代理字段的缺省值,赋值给对应 field;
  4. 将非代理字段(即类自带的普通字段)初始化;
  5. 执行 init 代码块;

上述做法在绝大多数场景下是没问题的,但严格来说还是有一些风险:当 property 和 init 代码块穿插编写时,初始化顺序会与预期有微小出入。对此似乎没有很好的解决方案,官方的 kotlin-serialization 插件也只记了个 todo。

接下来需要手撸一个包含 Protobuf 解析逻辑的 constructor。

给出一个类:

@ProtoMessage(value = "webcast.data.Image")
class ImageModel {
    val uri by serialized<String>()
    val height by serialized<Long>(defaultValue = 1L)
    val urlList by serialized<Array<String>>("url_list")
}

产物(反编译至 Java):

public final class ImageModel implements ModelXModified {
  // 省略字段声明和 accessors

  public ImageModel() {
    this._$$uri = "";
    this._$$height = 1L;
    this._$$urlList = new String[0];
  }

  public ImageModel(@NotNull ProtoReader reader) {
    ArrayList arrayList = new ArrayList();
    long l = reader.beginMessage();
    int i;
    while ((i = reader.nextTag()) != -1) {
      switch (i) {
        case 1:
          arrayList.add(ProtoScalarTypeDecoder.decodeString(reader));
          continue;
        case 2:
          this._$$uri = ProtoScalarTypeDecoder.decodeString(reader);
          continue;
        case 3:
          this._$$height = ProtoScalarTypeDecoder.decodeInt64(reader);
          continue;
      } 
      ProtoScalarTypeDecoder.skipUnknown(reader);
    } 
    reader.endMessage(l);
    this._$$urlList = (String[])arrayList.toArray(new String[arrayList.size()]);
    if (this._$$uri == null)
      this._$$uri = ""; 
    if (this._$$height == 0L)
      this._$$height = 1L; 
    if (this._$$urlList.length == 0)
      this._$$urlList = new String[0]; 
  }
}

构造方法大致分为三个部分:

  1. 字段初始化:这里仅针对集合类字段(Map / List / Array),其它字段在解码过程中赋值;
  2. 流式解码:采用类似 Visitor 的模式,通过 Reader 顺序解析二进制流或数组;
  3. 处理数组:将类型为 ArrayList 的临时变量中的元素拷贝到数字字段;
  4. 赋缺省值:针对未被初始化的字段赋缺省值;

代码中1/4的代码都是在在处理各式各样的类型转换。

对于一个 proto field:

repeated int32 id;

包括但不限于以下这些声明都是合法的:

val ids1 by serialized<List<Long>>()
val ids2 by serialized<List<Long?>>()
val ids3 by serialized<List<Long>?>()
val ids4 by serialized<Set<Long>>()
val ids5 by serialized<Set<Int>>()
val ids6 by serialized<IntArray>()
val ids7 by serialized<LongArray>()
val ids8 by serialized<Array<Int>>()

平时这些问题编译器都会替我们处理好,但这次我们作为编译器的一部分,只能自己解决了。

首先我们把问题简化,上述类型涉及到三个维度:

  1. 集合类型(Array、List、Set 等)
  2. 集合中元素的类型
  3. 内外类型是否可空

集合类型和元素类型都可以相对简单地映射到 proto message 上,唯有可空处理是个两难:在 Protobuf 编码的二进制数据中,null 值和默认值是没有区别的。

小知识:当数字类型的 field 为 0、boolean 类型的 field 为 false、字符串 field 为空字符串等情况下,field 并不会被编码到二进制数据中(为了节省空间)。因此客户端从收到的二进制反序列化时并不能判断服务端的数据究竟是 null 还是默认值。

这样一来,一个字段最后究竟 fallback 到默认值还是 null 就完全由我们说了算。参考之前基于 apt 的方案后,我们制定了以下规则:

  1. 当字段非 repeated、非 map,且声明为可空,则 fallback 到 null;
  2. 当字段非 repeated、非 map,且声明为不可空,则 fallback 到指定的缺省值或默认值;
  3. 当字段为 repeated 或 map,则 fallback 到无内容的集合对象;

上述处理方式引入一个有意思的点:写在 serialized<>() 中的 defaultValue 即缺省值只有在特定情况下起作用,因此如果是方法调用表达式,就“意外”收获到了懒加载的特性。

5 . 总结

总的来说,相比 Android 开发,KCP 的开发体验要差一些。网上基本没有文档,只能带着问题阅读源码。下面列举了一些趟坑经验,但愿能帮到有意尝试的同学。

特性设计

  1. 特性设计尽量直观、易懂,尽量做到开箱即用;
  2. 尽量不要直接修改语法层面的规则,避免对 IDE Plugin 的过度依赖;
  3. 尽量避免引入太多与 Kotlin 源码的耦合(即 EP 之外的引用),以避免过高的后期维护成本;

代码生成

  1. 由于语言规则完全由自己定义,搞 KCP 时一定要做好编译工作前的语法检查(通过 AnalysisHandlerExtension)。冗余无所谓,缺少检查导致生成了无法运行的字节码,就是火葬场;
  2. 尽量简化代码生成,将复杂、繁琐的逻辑抽象成方法并挪至 runtime 依赖中,避免直接手撸重复或复杂逻辑(以可以忽略不计的效率牺牲,换取极大的健壮性、可阅读性提升);
  3. 如果之前字节码了解不多,建议先系统学一下。ASM 中报错信息比较难懂,遇到坑更多是靠人肉 试错 分析;
  4. 做好单元测试,推荐使用 kotlin-compile-testing,可在单测中“现场编译”测试用例,能解决最基本的断点调试需求。

IR

所谓 IR,简单说就是一个“中间件”,处理多平台(JVM、Native、JS)通用的逻辑。在 Kotlin 1.5 之前,IR 默认是关闭的。考虑到我们主流 App 还停留在 1.3.x 上,现阶段还不需要过多关注。参考 The New JVM IR Backend Is Stable、Exploring Kotlin IR

源码

有学习价值的源码:

Copyright© 2013-2020

All Rights Reserved 京ICP备2023019179号-8