告别KAPT！使用 KSP 为 Kotlin 编译提速

本文已参与好文召集令活动，点击查看：后端、大前端双赛道投稿，2万元奖池等你挑战！

今年初 Android 发布了 Kotlin Symbol Processing（KSP）的首个 Alpha 版，几个月过去，KSP 已经更新到 Beta3 了， 目前 API 已经基本稳定，相信距离稳定版发布也不会很远了。

为什么使用 KSP ？

不少人吐槽 Kotlin 的编译速度，KAPT 便是拖慢编译的元凶之一。

很多库都会使用注解简化模板代码，例如 Room、Dagger、Retrofit 等，Kotlin 代码使用 KAPT 处理注解。 KAPT 本质上是基于 APT 工作的，APT 只能处理 Java 注解，因此需要先生成 APT 可解析的 stub (Java代码)，这拖慢了 Kotlin 的整体编译速度。

KSP 正是在这个背景下诞生的，它基于 Kotlin Compiler Plugin（简称KCP） 实现，不需要生成额外的 stub，编译速度是 KAPT 的 2 倍以上

KSP 与 KCP

Kotlin Compiler Plugin 在 kotlinc 过程中提供 hook 时机，可以再次期间解析 AST、修改字节码产物等，Kotlin 的不少语法糖都是 KCP 实现的，例如 data class、 @Parcelize、kotlin-android-extension 等, 如今火爆的 Compose 其编译期工作也是借助 KCP 完成的。

理论上 KCP 的能力是 KAPT 的超集，可以替代 KAPT 以提升编译速度。但是 KCP 的开发成本太高，涉及 Gradle Plugin、Kotlin Plugin 等的使用，API 涉及一些编译器知识的了解，一般开发者很难掌握。

一个标准 KCP 的开发涉及以下诸多内容：

• Plugin：Gradle 插件用来读取 Gradle 配置传递给 KCP（Kotlin Plugin）
• Subplugin：为 KCP 提供自定义 KP 的 maven 库地址等配置信息
• CommandLineProcessor：将参数转换为 KP 可识别参数
• ComponentRegistrar：注册 Extension 到 KCP 不同流程中
• Extension：实现自定义的 KP 功能

KSP 简化了上述流程，开发者无需了解编译器工作原理，处理注解等成本像 KAPT 一样低。

KSP 与 KAPT

KSP 顾名思义，在 Symbols 级别对 Kotlin 的 AST 进行处理，访问类、类成员、函数、相关参数等类型的元素。可以类比 PSI 中的 Kotlin AST

一个 Kotlin 源文件经 KSP 解析后的结果如下：

KSFile
  packageName: KSName
  fileName: String
  annotations: List<KSAnnotation>  (File annotations)
  declarations: List<KSDeclaration>
    KSClassDeclaration // class, interface, object
      simpleName: KSName
      qualifiedName: KSName
      containingFile: String
      typeParameters: KSTypeParameter
      parentDeclaration: KSDeclaration
      classKind: ClassKind
      primaryConstructor: KSFunctionDeclaration
      superTypes: List<KSTypeReference>
      // contains inner classes, member functions, properties, etc.
      declarations: List<KSDeclaration>
    KSFunctionDeclaration // top level function
      simpleName: KSName
      qualifiedName: KSName
      containingFile: String
      typeParameters: KSTypeParameter
      parentDeclaration: KSDeclaration
      functionKind: FunctionKind
      extensionReceiver: KSTypeReference?
      returnType: KSTypeReference
      parameters: List<KSVariableParameter>
      // contains local classes, local functions, local variables, etc.
      declarations: List<KSDeclaration>
    KSPropertyDeclaration // global variable
      simpleName: KSName
      qualifiedName: KSName
      containingFile: String
      typeParameters: KSTypeParameter
      parentDeclaration: KSDeclaration
      extensionReceiver: KSTypeReference?
      type: KSTypeReference
      getter: KSPropertyGetter
        returnType: KSTypeReference
      setter: KSPropertySetter
        parameter: KSVariableParameter
    KSEnumEntryDeclaration
      // same as KSClassDeclaration
复制代码

这是 KSP 中的 Kotlin AST 抽象。 类似的， APT/KAPT 中有对 Java 的 AST 抽象，其中能找到一些对应关系，比如 Java 使用 Element 描述包、类、方法或者变量等， KSP 中使用 Declaration

Java/APT	Kotlin/KSP	Description
PackageElement	KSFile	表示一个包程序元素。提供对有关包及其成员的信息的访问
ExecuteableElement	KSFunctionDeclaration	表示某个类或接口的方法、构造方法或初始化程序（静态或实例），包括注释类型元素
TypeElement	KSClassDeclaration	表示一个类或接口程序元素。提供对有关类型及其成员的信息的访问。注意，枚举类型是一种类，而注解类型是一种接口
VariableElement	KSVariableParameter / KSPropertyDeclaration	表示一个字段、enum 常量、方法或构造方法参数、局部变量或异常参数

Declaration 之下还有 Type 信息 ，比如函数的参数、返回值类型等，在 APT 中使用 TypeMirror 承载类型信息 ，KSP 中详细的能力由 KSType 实现。

KSP 的开发流程和 KAPT 类似：

1. 解析源码AST
2. 生成代码
3. 生成的代码与源码一起参与 Kotlin 编译

需要注意 KSP 不能用来修改原代码，只能用来生成新代码

KSP 入口：SymbolProcessorProvider

KSP 通过 SymbolProcessor 来具体执行。SymbolProcessor 需要通过一个 SymbolProcessorProvider 来创建。因此 SymbolProcessorProvider 就是 KSP 执行的入口

interface SymbolProcessorProvider {
    fun create(environment: SymbolProcessorEnvironment): SymbolProcessor
}
复制代码

SymbolProcessorEnvironment 获取一些 KSP 运行时的依赖，注入到 Processor

interface SymbolProcessor {
    fun process(resolver: Resolver): List<KSAnnotated> // Let's focus on this
    fun finish() {}
    fun onError() {}
}
复制代码

process() 提供一个 Resolver , 解析 AST 上的 symbols。 Resolver 使用访问者模式去遍历 AST。

如下，Resolver 使用 FindFunctionsVisitor 找出当前 KSFile 中 top-level 的 function 以及 Class 成员方法：

class HelloFunctionFinderProcessor : SymbolProcessor() {
    ...
    val functions = mutableListOf<String>()
    val visitor = FindFunctionsVisitor()

    override fun process(resolver: Resolver) {
        //使用 FindFunctionsVisitor 遍历访问 AST
        resolver.getAllFiles().map { it.accept(visitor, Unit) }
    }

    inner class FindFunctionsVisitor : KSVisitorVoid() {
        override fun visitClassDeclaration(classDeclaration: KSClassDeclaration, data: Unit) {
            //访问 Class 节点
            classDeclaration.getDeclaredFunctions().map { it.accept(this, Unit) }
        }

        override fun visitFunctionDeclaration(function: KSFunctionDeclaration, data: Unit) {
            // 访问 function 节点
            functions.add(function)
        }

        override fun visitFile(file: KSFile, data: Unit) {
            //访问 file
            file.declarations.map { it.accept(this, Unit) }
        }
    }
    ...
}
复制代码

KSP API 示例

举几个例子看一下 KSP 的 API 是如何工作的

访问类中的所有成员方法

fun KSClassDeclaration.getDeclaredFunctions(): List<KSFunctionDeclaration> {
    return this.declarations.filterIsInstance<KSFunctionDeclaration>()
}
复制代码

判断一个类或者方法是否是局部类或局部方法

fun KSDeclaration.isLocal(): Boolean {
    return this.parentDeclaration != null && this.parentDeclaration !is KSClassDeclaration
}
复制代码

判断一个类成员是否对其他Declaration可见

fun KSDeclaration.isVisibleFrom(other: KSDeclaration): Boolean {
    return when {
        // locals are limited to lexical scope
        this.isLocal() -> this.parentDeclaration == other
        // file visibility or member
        this.isPrivate() -> {
            this.parentDeclaration == other.parentDeclaration
                    || this.parentDeclaration == other
                    || (
                        this.parentDeclaration == null
                            && other.parentDeclaration == null
                            && this.containingFile == other.containingFile
                    )
        }
        this.isPublic() -> true
        this.isInternal() && other.containingFile != null && this.containingFile != null -> true
        else -> false
    }
}
复制代码

获取注解信息

// Find out suppressed names in a file annotation:
// @file:kotlin.Suppress("Example1", "Example2")
fun KSFile.suppressedNames(): List<String> {
    val ignoredNames = mutableListOf<String>()
    annotations.forEach {
        if (it.shortName.asString() == "Suppress" && it.annotationType.resolve()?.declaration?.qualifiedName?.asString() == "kotlin.Suppress") {
            it.arguments.forEach {
                (it.value as List<String>).forEach { ignoredNames.add(it) }
            }
        }
    }
    return ignoredNames
}
复制代码

代码生成的示例

最后看一个相对完整的例子，用来替代APT的代码生成

@IntSummable
data class Foo(
  val bar: Int = 234,
  val baz: Int = 123
)
复制代码

我们希望通过KSP处理@IntSummable，生成以下代码

public fun Foo.sumInts(): Int {
  val sum = bar + baz
  return sum
}
复制代码

Dependencies

开发 KSP 需要添加依赖：

plugins {
    kotlin("jvm") version "1.4.32"
}

repositories {
    mavenCentral()
    google()
}

dependencies {
    implementation(kotlin("stdlib"))
    implementation("com.google.devtools.ksp:symbol-processing-api:1.5.10-1.0.0-beta01")
}
复制代码

IntSummableProcessorProvider

我们需要一个入口的 Provider 来构建 Processor

import com.google.devtools.ksp.symbol.*

class IntSummableProcessorProvider : SymbolProcessorProvider {

    override fun create(environment: SymbolProcessorEnvironment): SymbolProcessor {
        return IntSummableProcessor(
            options = environment.options,
            codeGenerator = environment.codeGenerator,
            logger = environment.logger
        )
    }
}
复制代码

通过 SymbolProcessorEnvironment 可以为 Processor 注入了 options、CodeGenerator、logger 等所需依赖

IntSummableProcessor

class IntSummableProcessor() : SymbolProcessor {
    
    private lateinit var intType: KSType

    override fun process(resolver: Resolver): List<KSAnnotated> {
        intType = resolver.builtIns.intType
        val symbols = resolver.getSymbolsWithAnnotation(IntSummable::class.qualifiedName!!).filterNot{ it.validate() }

        symbols.filter { it is KSClassDeclaration && it.validate() }
            .forEach { it.accept(IntSummableVisitor(), Unit) }

        return symbols.toList()
    }
}    
复制代码

• builtIns.intType 获取到 kotlin.Int 的 KSType， 在后面需要使用。
• getSymbolsWithAnnotation 获取注解为 IntSummable 的 symbols 列表
• 当 symbol 是 Class 时，使用 Visitor 对其进行处理

IntSummableVisitor

Visitor 的接口一般如下，D 和 R 代表 Visitor 的输入和输出，

interface KSVisitor<D, R> {
    fun visitNode(node: KSNode, data: D): R

    fun visitAnnotated(annotated: KSAnnotated, data: D): R
    
    // etc.
}
复制代码

我们的需求没有输入输出，所以实现KSVisitorVoid即可，本质上是一个 KSVisitor<Unit, Unit>：

inner class Visitor : KSVisitorVoid() {
    
    override fun visitClassDeclaration(classDeclaration: KSClassDeclaration, data: Unit) {
        val qualifiedName = classDeclaration.qualifiedName?.asString()

        //1. 合法性检查
        if (!classDeclaration.isDataClass()) {
            logger.error(
                "@IntSummable cannot target non-data class $qualifiedName",
                classDeclaration
            )
            return
        }

        if (qualifiedName == null) {
            logger.error(
                "@IntSummable must target classes with qualified names",
                classDeclaration
            )
            return
        }
        
        //2. 解析Class信息
        //...
        
        //3. 代码生成
        //...
        
    }
    
    private fun KSClassDeclaration.isDataClass() = modifiers.contains(Modifier.DATA)
}
复制代码

如上，我们判断这个Class是不是data class、其类名是否合法

解析Class信息

接下来需要获取 Class 中的相关信息，用于我们的代码生成：

inner class IntSummableVisitor : KSVisitorVoid() {

    private lateinit var className: String
    private lateinit var packageName: String
    private val summables: MutableList<String> = mutableListOf()

    override fun visitClassDeclaration(classDeclaration: KSClassDeclaration, data: Unit) {
        //1. 合法性检查
        //...
        
        //2. 解析Class信息
        val qualifiedName = classDeclaration.qualifiedName?.asString()
        className = qualifiedName
        packageName = classDeclaration.packageName.asString()

        classDeclaration.getAllProperties()
            .forEach {
                it.accept(this, Unit)
            }

        if (summables.isEmpty()) {
            return
        }
        
        //3. 代码生成
        //...
    }
    
    override fun visitPropertyDeclaration(property: KSPropertyDeclaration, data: Unit) {
        if (property.type.resolve().isAssignableFrom(intType)) {
            val name = property.simpleName.asString()
            summables.add(name)
        }
    }
}
复制代码

• 通过 KSClassDeclaration 获取了className， packageName，以及 Properties 并将其存入 summables
• visitPropertyDeclaration 中确保 Property 必须是 Int 类型，这里用到了前面提到的 intType

代码生成

收集完 Class 信息后，着手代码生成。 我们引入 KotlinPoet 帮助我们生成 Kotlin 代码

dependencies {
    implementation("com.squareup:kotlinpoet:1.8.0")
}
复制代码

override fun visitClassDeclaration(classDeclaration: KSClassDeclaration, data: Unit) {

    //1. 合法性检查
    //...
      

    //2. 解析Class信息
    //...
        
        
    //3. 代码生成
    if (summables.isEmpty()) {
        return
    }

    val fileSpec = FileSpec.builder(
        packageName = packageName,
        fileName = classDeclaration.simpleName.asString()
    ).apply {
        addFunction(
            FunSpec.builder("sumInts")
                .receiver(ClassName.bestGuess(className))
                .returns(Int::class)
                .addStatement("val sum = ${summables.joinToString(" + ")}")
                .addStatement("return sum")
                .build()
        )
    }.build()

    codeGenerator.createNewFile(
        dependencies = Dependencies(aggregating = false),
        packageName = packageName,
        fileName = classDeclaration.simpleName.asString()
    ).use { outputStream ->
        outputStream.writer()
            .use {
                fileSpec.writeTo(it)
            }
    }
}
复制代码

• 使用 KotlinPoet 的 FunSpec 生成 function 代码
• 前面SymbolProcessorEnvironment 提供的CodeGenerator用来创建文件，并写入生成的FileSpec代码

总结

通过 IntSummable 的例子可以看到 KSP 完全可以替代 APT/KAPT 进行注解处理，且性能更出色。

目前，已有不少使用 APT 的三方库增加了对 KSP 的支持

Library	Status	Tracking issue for KSP
Room	Experimentally supported
Moshi	Experimentally supported
Kotshi	Experimentally supported
Lyricist	Experimentally supported
Auto Factory	Not yet supported	Link
Dagger	Not yet supported	Link
Hilt	Not yet supported	Link
Glide	Not yet supported	Link
DeeplinkDispatch	Not yet supported	Link

将 KAPT 替换为 KSP 也非常简单，以 Moshi 为例

当然，也可以在项目中同时使用 KAPT 和 KSP ，他们互不影响。KSP 取代 KAPT 的趋势越来越明显，果你的项目也处理注解的需求，不妨试试 KSP ？

github.com/google/ksp