開發一個支持跨平台的 Kotlin 編譯器插件

語言: CN / TW / HK

theme: smartblue highlight: a11y-dark


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

前言

前面簡單介紹了一下Kotlin編譯器插件是什麼,以及如何一步一步開發一個Kotlin編譯器插件,但是之前開發的編譯器插件是通過修改字節碼的方式來修改產物的,只支持JVM平台。今天主要在此基礎上,介紹一下如何通過修改IR的方式來修改Kotlin編譯器的產物,如何開發一個支持跨平台的Kotlin編譯器插件

本文主要包括以下內容:
1. Kotlin IR是什麼? 2. 如何遍歷Kotlin IR? 3. 如何創建Kotlin IR元素? 4. 如何修改Kotlin IR? 5. 修改Kotlin IR實戰

Kotlin IR是什麼?

Kotlin IRKotlin編譯器中間表示,它從數據結構上來説也是一個抽象語法樹。

因為Kotlin是支持跨平台的,因此有着JVMNativeJS三個不同的編譯器後端,為了在不同的後端之間共享邏輯,以及簡化支持新的語言特性所需的工作,Kotiln編譯器引入IR的概念,如下圖所示:

在前文開發你的第一個 Kotlin 編譯器插件中主要使用了ClassBuilderInterceptorExtension在生成字節碼的時機來修改產物

但是通過這種開發的插件是不支持Kotlin跨平台的,很顯然,NativeJS平台並不會生成字節碼

這就是修改IR的意義,讓我們的編譯器插件支持跨平台

正是為了支持跨平台,官方開發的很多插件,比如Compose編譯器插件,都是基於IrGenerationExtension

IrElement.dump()使用

在從概念上理解了Kotlin IR是什麼樣之後,我們再來看下Kotlin IR在代碼中到底長什麼樣?

Kotlin IR 語法樹中的每個節點都實現了 IrElement。語法樹的元素包括模塊、包、文件、類、屬性、函數、參數、if語句、函數調用等等,我們可以通過IrElement.dump方法來看下它們在代碼中的樣子

```kotlin // 1. 註冊IrGenerationExtension IrGenerationExtension.registerExtension(project, TemplateIrGenerationExtension(messageCollector))

// 2. IrGenerationExtension具體實現 class TemplateIrGenerationExtension() : IrGenerationExtension { override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) { println(moduleFragment.dump()) } } ```

通過以上方式就可以實現IrGenerationExtension的註冊與具體實現,如果我們的源代碼如下所示:

```kotlin fun main() { println(debug()) }

fun debug(name: String = "World") = "Hello, $name!" ```

編譯器插件將輸出以下內容:

kotlin MODULE_FRAGMENT name:<main> FILE fqName:<root> fileName:/var/folders/dk/9hdq9xms3tv916dk90l98c01p961px/T/Kotlin-Compilation3223338783072974845/sources/main.kt // 1 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: CALL 'public final fun debug (name: kotlin.String): kotlin.String declared in <root>' type=kotlin.String origin=null // 2 FUN name:debug visibility:public modality:FINAL <> (name:kotlin.String) returnType:kotlin.String VALUE_PARAMETER name:name index:0 type:kotlin.String EXPRESSION_BODY CONST String type=kotlin.String value="World" // 3 BLOCK_BODY RETURN type=kotlin.Nothing from='public final fun debug (name: kotlin.String): kotlin.String declared in <root>' STRING_CONCATENATION type=kotlin.String CONST String type=kotlin.String value="Hello, " GET_VAR 'name: kotlin.String declared in <root>.debug' type=kotlin.String origin=null CONST String type=kotlin.String value="!"

這就是Kotlin IR在代碼中的樣子,可以看出它包括以下內容:

  1. main函數的聲明,可見性,可變性,參數與返回值,可以看出這是一個名為publicfinal函數main,它不接受任何參數,並返回Unit
  2. debug函數則有一個參數name,該參數具有類型String並且還返回一個String,參數的默認值通過VALUE_PARAMETER表示
  3. debug函數的函數體通過BLOCK_BODY表示,返回的內容是一個String

IrElement.dump()Kotlin編譯器插件的開發過程中是很實用的,通過它我們可以Dump任意代碼並查看其結構

如何遍歷Kotlin IR

如前文所説,Kotlin IR是一個抽象語法樹,這意味着我們可以利用處理樹結構的方式來處理Kotlin IR

我們可以利用訪問者模式來遍歷Kotlin IR,我們知道,Kotlin IR中的每個節點都實現了IrElement接口,而IrElement接口中正好有 2 個訪問者模式相關的函數。

kotlin fun <R, D> accept(visitor: IrElementVisitor<R, D>, data: D): R fun <D> acceptChildren(visitor: IrElementVisitor<Unit, D>, data: D): Unit

接下來我們看一下在IrClass中兩個方法的實現:

```kotlin override fun accept(visitor: IrElementVisitor, data: D): R = visitor.visitClass(this, data)

override fun acceptChildren(visitor: IrElementVisitor, data: D) { thisReceiver?.accept(visitor, data) typeParameters.forEach { it.accept(visitor, data) } declarations.forEach { it.accept(visitor, data) } } ```

可以看出:acceptChildren方法會調用所有childrenaccept方法,而accept方法則是調用對應elementvisit方法,最後都會調用到visitElement方法

因此我們可以通過以下方式遍歷IR

```kotlin // 1. 註冊 visitor moduleFragment.accept(RecursiveVisitor(), null)

// 2. visitor 實現 class RecursiveVisitor : IrElementVisitor { override fun visitElement(element: IrElement, data: Nothing?) { element.acceptChildren(this, data) } } ```

數據輸入與輸出

上文當我們使用IrElementVisitor時,忽略了它的兩個泛型參數,一個用於定義data(每個visit函數接受的參數類型,另一個用於定義每個visitor函數的返回類型。

輸入值data可用於在整個 IR 樹中傳遞上下文信息。例如,可以用於打印出元素細節時使用的當前縮進間距。

kotlin class StringIndentVisitor : IrElementVisitor<Unit, String> { override fun visitElement(element: IrElement, data: String) { println("$data${render(element)} {") element.acceptChildren(this, " $data") println("$data}") } }

而輸出類型可用於返回調用訪問者的結果,在後面的IR轉換中非常實用,這點我們後面再介紹。

如何創建Kotlin IR元素?

上文我們已經介紹瞭如何遍歷Kotlin IR,接下來我們來看下如何創建Kotlin IR元素

kotlin class TemplateIrGenerationExtension() : IrGenerationExtension { override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) { println(moduleFragment.dump()) } }

當定義IrGenerationExtension時,我們之前使用了moduleFragment參數,接下來我們來看下另一個參數:IrPluginContext,它可以為插件提供有關正在編譯的當前模塊之外的內容的上下文信息

IrPluginContext實例中,我們可以獲得IrFactory的實例。這個工廠類是 Kotlin 編譯器插件創建自己的 IR 元素的方式。它包含許多用於構建IrClassIrSimpleFunctionIrProperty等實例的函數。

IrFactory在構建聲明時很有用:比如類、函數、屬性等,但是在構建語句和表達式時,您將需要IrBuilder的實例。更重要的是,您將需要一個IrBuilderWithScope實例。有了這個構建器實例,IR 表達式可以使用更多的擴展函數。

在瞭解了這些基礎內容後,我們來看一個實例

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

這段代碼很簡單,我們來看下如何實用IR構建如上內容:

```kotlin override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) { // 1. 通過pluginContext.irBuiltIns獲取Kotlin語言中的內置內容,在這裏我們獲取了any與unit類型 val typeNullableAny = pluginContext.irBuiltIns.anyNType val typeUnit = pluginContext.irBuiltIns.unitType // 2. 如果您需要的不是語言本身內置的,而是來自依賴項(如標準庫),您可以使用這些IrPluginContext.reference*()函數來查找所需的IrSymbol val funPrintln = pluginContext.referenceFunctions(FqName("kotlin.io.println")) .single { val parameters = it.owner.valueParameters parameters.size == 1 && parameters[0].type == typeNullableAny }

// 3. 使用irFactory構建一個函數 val funMain = pluginContext.irFactory.buildFun { name = Name.identifier("main") visibility = DescriptorVisibilities.PUBLIC // default modality = Modality.FINAL // default returnType = typeUnit }.also { function -> // 4. 設置function.body,構建函數體 function.body = DeclarationIrBuilder(pluginContext, function.symbol).irBlockBody { 通過+號將此表達式添加到塊中 +irCall(funPrintln).also { call -> call.putValueArgument(0, irString("Hello, World!")) } } }

println(funMain.dump()) } ```

以上代碼主要做了這麼幾件事:

  1. 通過pluginContext.irBuiltIns獲取Kotlin語言中的內置內容,在這裏我們獲取了anyunit類型
  2. 如果您需要的不是語言本身內置的,而是來自依賴項(如標準庫),您可以使用這些IrPluginContext.reference*()函數來查找所需的IrSymbol,同時由於函數支持重載,我們這裏需要通過single方法過濾出具有所需簽名的單個函數
  3. 使用irFactory構建一個函數,可以設置各種屬性,如名稱、可見性、可變性和返回類型等
  4. 通過設置function.body構建函數體,irBlockBody會創建出一個IrBuilderWithScope,在其中可以可以調用各種擴展方法創建IR,比如調用irCall。同時需要通過IrCall上的+運算符將此函數調用添加到塊中

如上所示,通過這段代碼就可以構建出我們想要的代碼了

如何修改Kotlin IR?

在瞭解瞭如何遍歷與創建IR之後,接下來就是修改了。

與之前遍歷IR樹類似,IrElement接口也包含兩個與變換相關的接口

kotlin fun <D> transform(transformer: IrElementTransformer<D>, data: D): IrElement = accept(transformer, data) fun <D> transformChildren(transformer: IrElementTransformer<D>, data: D): Unit

transform函數默認委託訪問者函數accept,子類中的函數覆蓋通常只需要將函數的返回類型覆蓋為更具體的類型。例如,IrFile中的transform函數如下所示。

kotlin override fun <D> transform(transformer: IrElementTransformer<D>, data: D): IrFile = accept(transformer, data) as IrFile

transformChildren方法與遍歷訪問每個元素的所有子元素一樣,transformChildren 函數允許對每個子元素進行變換。例如,讓我們看看IrClass的實現。

kotlin override fun <D> transformChildren(transformer: IrElementTransformer<D>, data: D) { thisReceiver = thisReceiver?.transform(transformer, data) typeParameters = typeParameters.transformIfNeeded(transformer, data) declarations.transformInPlace(transformer, data) } 總得來説,transformtransformChildren方法最後也會調用到各種visit方法,我們可以在其中修改IR的內容

在瞭解了這些基礎之後,我們可以開始修改Kotlin IR的實戰了

修改Kotlin IR實戰

目標代碼

接下來我們來看一個修改Kotlin IR的實例,比如以下代碼

kotlin @DebugLog fun greet(greeting: String = "Hello", name: String = "World"): String { return "${'$'}greeting, ${'$'}name!" }

我們希望添加了@DebugLog註解的方法,在函數的入口與出口都通過println打印信息,在變換後代碼如下所示:

kotlin @DebugLog fun greet(greeting: String = "Hello", name: String = "World"): String { println("⇢ greet(greeting=$greeting, name=$name)") val startTime = TimeSource.Monotonic.markNow() try { val result = "${'$'}greeting, ${'$'}name!" println("⇠ greet [${startTime.elapsedNow()}] = $result") return result } catch (t: Throwable) { println("⇠ greet [${startTime.elapsedNow()}] = $t") throw t } }

接下來我們就一步一步來實現這個目標

註冊與定義Transformer

```kotlin // 1. 註冊Transformer moduleFragment.transform(DebugLogTransformer(pluginContext, debugLogAnnotation, funPrintln), null)

// 2. 定義Transformer class DebugLogTransformer( private val pluginContext: IrPluginContext, private val debugLogAnnotation: IrClassSymbol, private val logFunction: IrSimpleFunctionSymbol, ) : IrElementTransformerVoidWithContext() { private val typeUnit = pluginContext.irBuiltIns.unitType

private val classMonotonic = pluginContext.referenceClass(FqName("kotlin.time.TimeSource.Monotonic"))!!

override fun visitFunctionNew(declaration: IrFunction): IrStatement { val body = declaration.body if (body != null && declaration.hasAnnotation(debugLogAnnotation)) { declaration.body = irDebug(declaration, body) } return super.visitFunctionNew(declaration) } ```

這一步主要做了這麼幾件事:

  1. 註冊Transfomer,傳入用於構建 IR 元素的IrPluginContext,需要處理的註解符號IrClassSymbol,用於記錄調試消息的函數的IrSimpleFunctionSymbol
  2. 基於IrElementTransformerVoidWithContext類擴展,自定義Transformer,這個Transformer不接受輸入數據,並維護一個它訪問過的各種 IR 元素的內部堆棧
  3. 定義一些本地屬性來引用已知類型、類和函數,比如typeUnitclassMonotonic
  4. 重寫visitFunctionNew函數來攔截函數語句的轉換,我們需要檢查它是否有body並且擁有目標註解@DebugLog

方法進入打點

進入函數時,我們需要調用println顯示函數名稱和函數入參。

```kotlin private fun IrBuilderWithScope.irDebugEnter( function: IrFunction ): IrCall { val concat = irConcat() concat.addArgument(irString("⇢ ${function.name}(")) for ((index, valueParameter) in function.valueParameters.withIndex()) { if (index > 0) concat.addArgument(irString(", ")) concat.addArgument(irString("${valueParameter.name}=")) concat.addArgument(irGet(valueParameter)) } concat.addArgument(irString(")"))

return irCall(logFunction).also { call -> call.putValueArgument(0, concat) } } ```

在這裏我們主要用到了irConcat來拼接字符串,以及irGet來獲取參數值,這些參數經過concat拼接後通過println一起輸出

方法結束時打點

方法退出時,我們要記錄結果或拋出的異常。如果函數返回 Unit 我們可以跳過顯示結果,因為已知它什麼都沒有。

```kotlin private fun IrBuilderWithScope.irDebugExit( function: IrFunction, startTime: IrValueDeclaration, result: IrExpression? = null ): IrCall { val concat = irConcat() concat.addArgument(irString("⇠ ${function.name} [")) concat.addArgument(irCall(funElapsedNow).also { call -> call.dispatchReceiver = irGet(startTime) }) if (result != null) { concat.addArgument(irString("] = ")) concat.addArgument(result) } else { concat.addArgument(irString("]")) }

return irCall(logFunction).also { call -> call.putValueArgument(0, concat) } } ```

這裏我們需要記錄方法執行時間,startTime通過IrValueDeclaration傳入,這是一個局部變量,可以通過irGet讀取

為了調用kotlin.time.TimeMark.elapsedNow方法,我們可以調用funElapsedNow符號,並將startTime作為dispatcherReceiver,這樣就能計算出方法耗時

result參數是可選的表達式,它可以是函數的返回值,或者是拋出的異常,這些參數經過concat拼接後通過println一起輸出

組裝函數體

```kotlin private fun irDebug( function: IrFunction, body: IrBody ): IrBlockBody { return DeclarationIrBuilder(pluginContext, function.symbol).irBlockBody { +irDebugEnter(function) // ... val tryBlock = irBlock(resultType = function.returnType) { if (function.returnType == typeUnit) +irDebugExit(function, startTime) }.transform(DebugLogReturnTransformer(function, startTime), null)

+IrTryImpl(startOffset, endOffset, tryBlock.type).also { irTry ->
  irTry.tryResult = tryBlock
  irTry.catches += irCatch(throwable, irBlock {
    +irDebugExit(function, startTime, irGet(throwable))
    +irThrow(irGet(throwable))
  })
}

} } ```

這一部分主要是構建出try,catch塊,並在相關部分調用irDebugEnterirDebugExit方法,完整代碼就不在這裏綴述了

經過這一步,支持跨平台的Kotlin編譯器插件就開發完成了,添加了@DebugLog註解的方法在進入與退出時,都會打印相關的日誌

總結

本文主要介紹了Kotlin IR是什麼,如何對Kotlin IR進行增刪改查,如何一步步開發一個支持跨平台的Kotlin編譯器插件,希望對你有所幫助~

示例代碼

本文所有源碼可見:https://github.com/bnorm/debuglog

參考資料

Writing Your Second Kotlin Compiler Plugin