深入淺出 Compose Compiler(1) Kotlin Compiler & KCP

語言: CN / TW / HK

theme: vuepress highlight: androidstudio


本文為稀土掘金技術社群首發簽約文章,14天內禁止轉載,14天后未獲授權禁止轉載,侵權必究!

前言

Compose 的語法簡潔、程式碼效率非常高,這主要得益於 Compose Compiler 的一系列編譯期魔法,幫開發者生成了很多樣板程式碼。但編譯期插樁也阻礙了我們對於 Compose 執行原理的認知,想要真正讀懂 Compose 就必須先了解它的 Compiler。本系列文章將帶大家揭開 Compose Compiler 的神祕面紗。

Compose 是一個 Kotlin Only 框架,所以 Compose Compiler 的本質是一個 KCP(Kotlin Compiler Plugin)。在研究 Compose Compiler 原始碼之前,先要鋪墊一些 Kotlin Compiler 以及 KCP 的基礎知識

Kotlin 編譯流程

Kotlin 是一門跨平臺語言,Kotlin Compiler 可以將 Kt 原始碼編譯成多個平臺的目的碼:JS、JVM 位元組碼,甚至 LLVM 機器碼。但無論編譯成何種目的碼,其編譯過程都可以分為兩個階段: - Frontend(編譯器前端):對原始碼分析得到 AST (抽象語法樹)以及符號表,並完成靜態檢查 - Backend(編譯器後端):基於 AST 等前端產物,生成平臺目的碼

簡而言之:前端負責原始碼的解析和檢查,後端負責目的碼的生成

如上,以 Kotlin/JVM 為例:

  • Frontend 處理中,Kt 原始檔經過詞法、語法和語義分析(Lexer&Paser)生成 PSI 以及對應的 BindingContext。
  • Backend 處理中,基於 PSI 和 BindingContext 先生成 JVM 位元組碼,然後通過 ASM 將位元組碼二進位制化生成 class 檔案

不同目標平臺的編譯流程中 Frontend 的處理流程都一樣,只是在 Backend 中生成不同的目的碼

K1 編譯器:PSI & BindingContext

PSI 全稱 Program Structure Interface, 可以將它理解為 JetBrains 專用的 AST(標準 AST 之上有一些擴充套件)。PSI 可以用於編譯過程中的語法靜態檢查,PSI 也用於 IntelliJ 系列 IDE 的靜態檢查,我們在編寫程式碼過程中能實時提示語法錯誤就是靠它。因此 PSI 有助於編譯和編寫階段複用靜態檢查邏輯。我們在開發 IDE Plugin 或者編寫 Detekt 靜態檢查用例時都有機會使用到 PSI。

  • PSI: https://plugins.jetbrains.com/docs/intellij/psi-elements.html
  • Detekt: https://github.com/detekt/detekt

在 IDE 中通過 PsiViewer 外掛可以實時看到原始碼對應的 PSI,以下面程式碼為例:

kotlin fun main() { println("Hello, World!") }

上圖是 PsiViewer 中的輸出結果,可以看到它體現了以下樹形結構:

PSI 樹的節點是原始碼經分析後的語法元素,例如一個特殊符號,一個字串等,這都是一個個 PsiElement。PsiElement 仍然缺少了基於上下文的語義資訊,比如對於一個 KtFunction,它的引數資訊,修飾符資訊等等,這就需要 BindingContext 的輔助了。

BindingContext 相當於 PSI 配套的符號表,PsiElement 經語義分析後得到對應的 Descriptor (描述符)並記錄到 BindingContext 中,BindingContext 可以快速索引到 PSI 節點對應的 Descriptor。Descriptor 包含我們需要的語義資訊,例如 FunctionDescriptor 可以獲取 TypeParameters,isInline 等資訊。

BindingContext 結構類似一個 Map<Type, Map<key, Descriptor> ,第一個 Map 的 key 代表 PSI 節點型別,第二個 Map 的 key 是 PsiElement 例項,Value 是其對應的 Descriptor。KtFunction 為 key 可以獲取對應的 FunctionDescriptor;KtCallExpression 獲取對應的 ResolvedCall,這裡麵包含了呼叫方法的 FunctionDescriptor 以及傳入的 Parameters。

K2 編譯器:FIR & IR

通過上面的介紹我們知道,Kotlin Compiler 的 Frotend 產物是 PSI 以及 BindingContext,Backend 將基於它們直接輸出目的碼。由於 Backend 耦合了目的碼生成邏輯,一些編譯期的處理和優化邏輯難以多平臺複用。例如我們都知道的 suspend 函式在編譯期會生成額外的程式碼,而我們希望這些 codegen 邏輯得以複用,為此 Kotlin 開發了新一代編譯器,取名為 K2 。

K2: https://blog.jetbrains.com/zh-hans/kotlin/2021/10/the-road-to-the-k2-compiler/

K2 編譯器的最大特點是引入了 IR(Intermediate Representation,中間表達)。IR 是連線前後端的中間產物, 它與平臺無關,類似 suspend 這類編譯期優化可以面向 IR 實現並跨平臺複用。

K2 中使用新的基於 IR 的 Backend 替代舊有的基於 PSI 和 BindingContext 的 Backend。Kotlin 1.5 開始 Kotlin/JVM 預設啟用新的 IR Backend,1.6 開始 Kotin/JS IR Backend 成了標配。下圖是引入 IR Backend 的編譯流程。

IR 也是一顆樹形資料結構,但它的抽象表達更加“低階”,更貼近 CPU 架構。IrElement 帶有多種語義資訊,例如 FUN 的 visibility,modality 以及 returnType 等等,不必像 PsiElement 那樣需要通過查詢 BindingContext 獲取這些資訊。

前面 Hello World 的例子,其對應的 IR 樹列印如下:

txt FUN name:main visibility:public modality:FINAL <> () returnType:kotlin.Unit BLOCK_BODY CALL 'public final fun println (message: kotlin.Any?): kotlin.Unit [inline] declared in kotlin.io.ConsoleKt' type=kotlin.Unit origin=null message: CONST String type=kotlin.String value="Hello, World!"

除了新的 IR Backend,K2 也更新了 Frontend,主要變化是使用 FIR (Frontend IR)替代了 PSI 與 BindingContext。1.7.0 起我們可以使用到 K2 的新前端。

綜上可見: K2 相對於 K1 的主要變化引入了 FIR Frontend 和 IR Backend

IR 可以由 FIR 轉化而來,它們都是樹型結構,那麼這兩者又有什麼區別呢?可以從以下三個方面進行區分:

||FIR|IR| |--|--|--| |目標不同|FIR 整合了 PSI 與 BindingContext 資訊,更快速地查詢描述符資訊,它的首要目標是提升前端靜態分析以及檢查的效能|效能不是 IR 的考慮,它的資料結構的出發點不是為了提升後端編譯速度,而是服務於不同後端之間的編譯邏輯共享,降低不同平臺支援新語言特性的成本| |結構不同|FIR 仍然是一顆 AST,只是增強了一些符號資訊,加速靜態分析|IR 不僅是一顆 AST,它提供了更豐富的基於上下文的語義資訊,比如我可以知道某個程式碼塊中的某個變數是臨時變數還是成員變數,而 FIR 難以做到| |能力不同|雖然 FIR 也可以處理一些簡單的脫糖和程式碼生成工作,但整體上仍然是服務於前端,不能對 AST 大幅度修改| IR 具有豐富的 Godegen API,可以更加靈活地對樹形結構進行 add/remove/update,實現任意編譯期的魔改需求|

KCP(Kotlin Compiler Plugin)

KCP 允許我們在上述 Kotlin 編譯過程中,通過增加擴充套件點以實現各種編譯期魔改。Kotlin 的不少語法糖都是基於 KCP 實現的,比如大家熟知的 No-arg、All-open、kotlinx-serialization 等等。

KCP 也可以像 KAPT 那樣在編譯期進行註解處理,但它相對於 KATP 更具優勢:

  1. KCP 在 Kotlin 編譯過程中進行,而 KAPT 需要在正式編譯之前增加額外的預編譯環節,因此 KCP 的效能更好。KSP(Kotlin Symbol Processing)也是基於 KCP 實現的,這也是為什麼 KSP 的效能更好的原因

  2. KAPT 主要是用來生成新程式碼,難以針對原有程式碼邏輯做修改。KCP 可以針對 Bytecode 或者 IR 做任意修改,能力更強大。

KCP 的開發步驟

KCP 雖然功能強大但是開發難度較高,開發一個完整的 KCP 要涉及多個步驟:

  • Gradle Plugin:
  • Plugin:KCP 是通過 Gradle 配置的,需要定義一個 Gradle 外掛,並在 Gradle 中配置 KCP 所需的編譯引數。
  • Subplugin: 建立從 Gradle Plugin 到 Kotlin Plugin 的連線,並將 Gradle 中配置的引數傳遞給 Kotlin Plugin

  • Kotlin Plugin:

  • CommandLineProcessor:KCP 的入口,定義 KCP 的 id、解析命令列引數等
  • ComponentRegister:註冊 KCP 中的 Extension 擴充套件點。它與 CommandLineProcessor 一樣都是通過 SPI 呼叫,需要新增 auto-service 註解
  • XXExtension:這是實現 KCP 邏輯的地方。Kotlin 提供了許多型別的 Extension 供我們實現。編譯器會在前端、後端的各個編譯環節中呼叫 KCP 註冊的對應型別的 Extension。例如 ExpressionCodegenExtension 可用來修改 Class 的 Body;ClassBuilderInterceptorExtension 可以修改 Class 的 Definition 等等

隨著 Kotlin Compiler 從 K1 升級到 K2,KCP 也提供了面向 K2 的 Extension。

以 No-arg 為例 ,No-arg 通過為 Class 添加註解自動生成無參建構函式。No-arg 原始碼中存在 K1、K2 兩套 Extension,可以相容不同 Kotlin 版本的使用:

  • No-arg: https://kotlinlang.org/docs/no-arg-plugin.html
  • source:https://cs.android.com/android-studio/kotlin/+/master:plugins/noarg/
  • NoArg K1:
  • CliNoArgDeclarationChecker:NoArg 不能作用於 Inner Class,這裡使用基於 PSI 的前端檢查邏輯檢查是否是 Inner Class
  • CliNoArgExpressionCodegenExtension:繼承自 ExpressionCodegenExtension,基於 PSI 和對應的 Descriptor 以 JVM 位元組碼的形式在 Class Body 中新增無參建構函式

  • NoArg K2:

  • FirNoArgDeclarationChecker:新的 K2 前端,可基於 FIR 檢查 InnerClass
  • NoArgIrGenerationExtension:繼承自 IrGenerationExtension ,基於 IR 新增無參建構函式

以 Backend Extension 為例,體會以下具體實現上的區別:

  • CliNoArgExpressionCodegenExtension 中的處理:

```kotlin // 1. 基於 descriptor 獲取 class 資訊 val superClassInternalName = typeMapper.mapClass(descriptor.getSuperClassOrAny()).internalName val constructorDescriptor = createNoArgConstructorDescriptor(descriptor) val superClass = descriptor.getSuperClassOrAny()

// 2. 通過 Codegen 直接生成無參建構函式對應的位元組碼 functionCodegen.generateMethod(JvmDeclarationOrigin.NO_ORIGIN, constructorDescriptor, object : CodegenBased(state) { override fun doGenerateBody(codegen: ExpressionCodegen, signature: JvmMethodSignature) { codegen.v.load(0, AsmTypes.OBJECT_TYPE)

    if (isParentASealedClassWithDefaultConstructor) {
        codegen.v.aconst(null)
        codegen.v.visitMethodInsn(
            Opcodes.INVOKESPECIAL, superClassInternalName, "<init>",
            "(Lkotlin/jvm/internal/DefaultConstructorMarker;)V", false
        )
    } else {
        codegen.v.visitMethodInsn(Opcodes.INVOKESPECIAL, superClassInternalName, "<init>", "()V", false)
    }

    if (invokeInitializers) {
        generateInitializers(codegen)
    }
    codegen.v.visitInsn(Opcodes.RETURN)
}

}) ```

  • NoArgIrGenerationExtension 中的處理:

```kotlin // 1. 基於 IrClass 獲取 Class 資訊 val superClass = klass.superTypes.mapNotNull(IrType::getClass).singleOrNull { it.kind == ClassKind.CLASS } ?: context.irBuiltIns.anyClass.owner val superConstructor = if (needsNoargConstructor(superClass)) getOrGenerateNoArgConstructor(superClass) else superClass.constructors.singleOrNull { it.isZeroParameterConstructor() } ?: error("No noarg super constructor for ${klass.render()}:\n" + superClass.constructors.joinToString("\n") { it.render() })

// 2. 基於 irFactory 等 IR API 建立建構函式 context.irFactory.buildConstructor { startOffset = SYNTHETIC_OFFSET endOffset = SYNTHETIC_OFFSET returnType = klass.defaultType }.also { ctor -> ctor.parent = klass ctor.body = context.irFactory.createBlockBody( ctor.startOffset, ctor.endOffset, listOfNotNull( IrDelegatingConstructorCallImpl( ctor.startOffset, ctor.endOffset, context.irBuiltIns.unitType, superConstructor.symbol, 0, superConstructor.valueParameters.size ), IrInstanceInitializerCallImpl( ctor.startOffset, ctor.endOffset, klass.symbol, context.irBuiltIns.unitType ).takeIf { invokeInitializers } ) ) } ```

NoArgIrGenerationExtension 是一個 IrGenerationExtension,這是專門用來更新 Ir 的擴充套件點,可以看到裡面已經沒有了對位元組碼的操作,取而代之使用 IR 中的各種 buildXXX API。

Compose Compiler 的程式碼生成也是依靠 IrGenerationExtension 實現的,所以:即使最早版本的 Compose 也要求 Kotlin 版本大於 1.5.10,就是因其 Compiler 只支援 IR Backend Extension

Compose Compiler

Compose Compiler 本質上是一個 KCP,在瞭解了 KCP 的基本構成之後,我們知道 Compose Compiler 的核心在於 Extension

Compose Compiler: https://cs.android.com/androidx/platform/frameworks/support/+/androidx-main:compose/compiler/compiler-hosted/

直接找到 ComposeComponentRegistrar,檢視註冊了哪些 Extension:

```kotlin class ComposeComponentRegistrar : ComponentRegistrar { //...

StorageComponentContainerContributor.registerExtensio
    project,
    ComposableCallChecker()
)
StorageComponentContainerContributor.registerExtensio
    project,
    ComposableDeclarationChecker()
)
StorageComponentContainerContributor.registerExtensio
    project,
    ComposableTargetChecker()
)
ComposeDiagnosticSuppressor.registerExtension(
    project,
    ComposeDiagnosticSuppressor()
)
@Suppress("OPT_IN_USAGE_ERROR")
TypeResolutionInterceptor.registerExtension(
    project,
    @Suppress("IllegalExperimentalApiUsage")
    ComposeTypeResolutionInterceptorExtension()
)
IrGenerationExtension.registerExtension(
    project,
    ComposeIrGenerationExtension(
        configuration = configuration,
        liveLiteralsEnabled = liveLiteralsEnabled,
        liveLiteralsV2Enabled = liveLiteralsV2Enabled
        generateFunctionKeyMetaClasses = generateFunc
        sourceInformationEnabled = sourceInformationE
        intrinsicRememberEnabled = intrinsicRememberE
        decoysEnabled = decoysEnabled,
        metricsDestination = metricsDestination,
        reportsDestination = reportsDestination,
    )
)
DescriptorSerializerPlugin.registerExtension(
    project,
    ClassStabilityFieldSerializationPlugin()
)

//...

} ```

  • ComposableCallChecker:檢查是否可以呼叫 @Composable 函式
  • ComposableDeclarationChecker:檢查 @Composable 的位置是否正確
  • ComposeDiagnosticSuppressor:遮蔽不必要的編譯診斷錯誤
  • ComposeIrGenerationExtension:負責 Composable 函式的程式碼生成
  • ClassStabilityFieldSerializationPlugin:分析 Class 是否穩定,並新增穩定性資訊

這裡的各種 Checker 是 Frontend Extension ,目前仍然是基於 K1 實現的,而位於 Backend 的 ComposeIrGenerationExtension 則面向 K2,這也是 Compose 程式碼生成的核心,會在本系列的後續文章中重點介紹。

參考

  • Writing Your First Kotlin Compiler Plugin
    https://resources.jetbrains.com/storage/products/kotlinconf2018/slides/5_Writing%20Your%20First%20Kotlin%20Compiler%20Plugin.pdf

  • Kotlin Compiler Internals In 1.4 and beyond
    https://docs.google.com/presentation/d/e/2PACX-1vTzajwYJfmUi_Nn2nJBULi9bszNmjbO3c8K8dHRnK7vgz3AELunB6J7sfBodC2sKoaKAHibgEt_XjaQ/pub?slide=id.g955e8c1462_0_190