前言

前面简单介绍了一下Kotlin编译器插件是什么，以及如何一步一步开发一个Kotlin编译器插件，但是之前开发的编译器插件是通过修改字节码的方式来修改产物的，只支持JVM平台。今天主要在此基础上，介绍一下如何通过修改IR的方式来修改Kotlin编译器的产物，如何开发一个支持跨平台的Kotlin编译器插件

本文主要包括以下内容：

1. Kotlin IR是什么？
2. 如何遍历Kotlin IR？
3. 如何创建Kotlin IR元素?
4. 如何修改Kotlin IR?
5. 修改Kotlin IR实战

Kotlin IR是什么？

Kotlin IR即Kotlin编译器中间表示，它从数据结构上来说也是一个抽象语法树。

因为Kotlin是支持跨平台的，因此有着JVM，Native，JS三个不同的编译器后端，为了在不同的后端之间共享逻辑，以及简化支持新的语言特性所需的工作，Kotiln编译器引入IR的概念，如下图所示：

在前文开发你的第一个 Kotlin 编译器插件中主要使用了ClassBuilderInterceptorExtension在生成字节码的时机来修改产物

但是通过这种开发的插件是不支持Kotlin跨平台的，很显然，Native与JS平台并不会生成字节码

这就是修改IR的意义，让我们的编译器插件支持跨平台

正是为了支持跨平台，官方开发的很多插件，比如Compose编译器插件，都是基于IrGenerationExtension的

IrElement.dump()使用

在从概念上理解了Kotlin IR是什么样之后，我们再来看下Kotlin IR在代码中到底长什么样？

Kotlin IR 语法树中的每个节点都实现了 IrElement。语法树的元素包括模块、包、文件、类、属性、函数、参数、if语句、函数调用等等，我们可以通过IrElement.dump方法来看下它们在代码中的样子

//  1\. 注册IrGenerationExtension
IrGenerationExtension.registerExtension(project, TemplateIrGenerationExtension(messageCollector))

// 2\. IrGenerationExtension具体实现
class TemplateIrGenerationExtension() : IrGenerationExtension {
    override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
        println(moduleFragment.dump())
    }
}

通过以上方式就可以实现IrGenerationExtension的注册与具体实现，如果我们的源代码如下所示：

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

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

编译器插件将输出以下内容：

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函数的声明，可见性，可变性，参数与返回值，可以看出这是一个名为public的final函数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 个访问者模式相关的函数。

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

接下来我们看一下在IrClass中两个方法的实现：

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

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

可以看出：acceptChildren方法会调用所有children的accept方法，而accept方法则是调用对应element的visit方法，最后都会调用到visitElement方法

因此我们可以通过以下方式遍历IR树

// 1\. 注册 visitor
moduleFragment.accept(RecursiveVisitor(), null)

// 2\. visitor 实现
class RecursiveVisitor : IrElementVisitor<Unit, Nothing?> {
  override fun visitElement(element: IrElement, data: Nothing?) {
    element.acceptChildren(this, data)
  }
}

数据输入与输出

上文当我们使用IrElementVisitor时，忽略了它的两个泛型参数，一个用于定义data(每个visit函数接受的参数类型，另一个用于定义每个visitor函数的返回类型。

输入值data可用于在整个 IR 树中传递上下文信息。例如，可以用于打印出元素细节时使用的当前缩进间距。

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元素

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

当定义IrGenerationExtension时，我们之前使用了moduleFragment参数，接下来我们来看下另一个参数：IrPluginContext，它可以为插件提供有关正在编译的当前模块之外的内容的上下文信息

从IrPluginContext实例中，我们可以获得IrFactory的实例。这个工厂类是 Kotlin 编译器插件创建自己的 IR 元素的方式。它包含许多用于构建IrClass、IrSimpleFunction、IrProperty等实例的函数。

IrFactory在构建声明时很有用：比如类、函数、属性等，但是在构建语句和表达式时，您将需要IrBuilder的实例。更重要的是，您将需要一个IrBuilderWithScope实例。有了这个构建器实例，IR 表达式可以使用更多的扩展函数。

在了解了这些基础内容后，我们来看一个实例

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

这段代码很简单，我们来看下如何实用IR构建如上内容：

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语言中的内置内容，在这里我们获取了any与unit类型
2. 如果您需要的不是语言本身内置的，而是来自依赖项（如标准库），您可以使用这些IrPluginContext.reference*()函数来查找所需的IrSymbol，同时由于函数支持重载，我们这里需要通过single方法过滤出具有所需签名的单个函数
3. 使用irFactory构建一个函数，可以设置各种属性，如名称、可见性、可变性和返回类型等
4. 通过设置function.body构建函数体，irBlockBody会创建出一个IrBuilderWithScope，在其中可以可以调用各种扩展方法创建IR，比如调用irCall。同时需要通过IrCall上的+运算符将此函数调用添加到块中

如上所示，通过这段代码就可以构建出我们想要的代码了

如何修改Kotlin IR?

在了解了如何遍历与创建IR之后，接下来就是修改了。

与之前遍历IR树类似，IrElement接口也包含两个与变换相关的接口

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函数如下所示。

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

而transformChildren方法与遍历访问每个元素的所有子元素一样，transformChildren 函数允许对每个子元素进行变换。例如，让我们看看IrClass的实现。

override fun <D> transformChildren(transformer: IrElementTransformer<D>, data: D) {
  thisReceiver = thisReceiver?.transform(transformer, data)
  typeParameters = typeParameters.transformIfNeeded(transformer, data)
  declarations.transformInPlace(transformer, data)
}

总得来说，transform与transformChildren方法最后也会调用到各种visit方法，我们可以在其中修改IR的内容

在了解了这些基础之后，我们可以开始修改Kotlin IR的实战了

修改Kotlin IR实战

目标代码

接下来我们来看一个修改Kotlin IR的实例，比如以下代码

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

我们希望添加了@DebugLog注解的方法，在函数的入口与出口都通过println打印信息，在变换后代码如下所示：

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

接下来我们就一步一步来实现这个目标

注册与定义Transformer

// 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. 定义一些本地属性来引用已知类型、类和函数，比如typeUnit与classMonotonic
4. 重写visitFunctionNew函数来拦截函数语句的转换，我们需要检查它是否有body并且拥有目标注解@DebugLog

方法进入打点

进入函数时，我们需要调用println显示函数名称和函数入参。

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 我们可以跳过显示结果，因为已知它什么都没有。

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一起输出

组装函数体

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块，并在相关部分调用irDebugEnter与irDebugExit方法，完整代码就不在这里缀述了

经过这一步，支持跨平台的Kotlin编译器插件就开发完成了，添加了@DebugLog注解的方法在进入与退出时，都会打印相关的日志

总结

本文主要介绍了Kotlin IR是什么，如何对Kotlin IR进行增删改查，如何一步步开发一个支持跨平台的Kotlin编译器插件，希望对你有所帮助~

作者：程序员江同学
链接：https://juejin.cn/post/7147468443413905416
来源：稀土掘金