Kotlin IR
Compose Compiler 在后端的工作主要是在 Composable 的函数签名或者函数体中生成必要的代码，这样才能在运行时智能地执行或者跳过重组。这部分代码生成是通过 ComposeIrGenerationExtension 实现，它继承自 IrGenerationExtension，需要基于 IR 生成代码，而非传统的 ASM 方式。IR 不像 ASM 那样被熟知，所以在介绍具体代码生成逻辑之前，有必要对 IR 先做一个基本介绍。
IR ，全称 Intermediate Representation，是 Kotlin Compiler 为了多目标平台之间能够共享编译逻辑引入的中间代码，各平台的编译器后端基于 IR 生成目标代码，面向 IR 的代码生成逻辑以此实现多平台复用。Compose 代码生成基于 IR 实现，也为其成为一个跨平台框架奠定了基础。
IR 树的结构
在 《深入浅出 Compose Compiler（1） Kotlin Compiler & KCP》 一文中我们介绍过，IR 跟 PSI 一样也是一棵树。我们可以通过 KCP 来分析 Kt 源码的 IR 树。自定义一个 IrGenerationExtension 并获取源文件的 IrModuleFragment。 IrModuleFragment 是 IR 树都跟节点，调用其 dump 方法可以打印出整颗 IR 树。
本文不详细介绍如何实现 KCP ，有兴趣的可以基于 github.com/bnorm/kotli… 来实验 IrGenerationExtension 与 dump
我们还是拿之前文章中分析 PSI 的例子：
fun main() { println("Hello, World!") }
我们可以自定义 KCP ，在 IrGenerationExtension 中获取对应源文件的 IrModuleFragment，调用其 dump 方法可以打印出整颗 IR 树，如下：
MODULE_FRAGMENT name:<main> FILE fqName:<root> fileName:/var/folders/b1/0fd1b6hs7lz0fm_mh346lybm0000gn/T/Kotlin-Compilation16792700918799505160/sources/main.kt 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!"
IrModuleFragment 以及树上的各个节点 IrFile、IrFunction 等都继承自 IrElement。与 PsiElement 对比，IrElement 不在携带 WhiteSpace 这样的源码级别的信息，更多得是承载了例如 IrType 这样的语义信息。
IR 树的遍历
网上介绍 IR 参考资料很少，但是不影响我们研究相关源码。因为 IR 与其他 AST 在遍历方式是大致类似，都是基于访问者模式。
IrElement 的定义如下：
interface IrElement { fun <R, D> accept(visitor: IrElementVisitor<R, D>, data: D): R fun <D> transform(transformer: IrElementTransformer<D>, data: D): IrElement fun <D> acceptChildren(visitor: IrElementVisitor<Unit, D>, data: D) fun <D> transformChildren(transformer: IrElementTransformer<D>, data: D) }
大家如果有 Gradle Transfrom 中使用 ASM 代码插桩的经验，对上面这些签名应该比较熟悉。IrElemnet 可以 accept 一个 visitor 遍历自身节点，acceptChildren 可以用来递归遍历子节点。
以 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) { typeParameters.forEach { it.accept(visitor, data) } declarations.forEach { it.accept(visitor, data) } thisReceiver?.accept(visitor, data) }
我们可以自定义 IrElementVisitor ，通过调用 IrClass 的 accept 或 acceptChildren 访问 IR 信息。
accept 同时可以接受一个 data，这可以用来传输一些遍历中所需的上下文。 transform 和 transformChildren 用于遍历 IrElement 的同时对其进行变换， Compose 中基于 IR 的代码生成就需要使用到 IrElementTransformer。
IrElementTransformer 是一种特殊的 IrElementVisitor：
interface IrElementTransformer<in D> : IrElementVisitor<IrElement, D>
IrElementTransformer 在被 transform 调用后，会返回更新后的 IrElement。仍然以 IrClass 为例：
override fun <D> transformChildren(transformer: IrElementTransformer<D>, data: D) { typeParameters = typeParameters.transformIfNeeded(transformer, data) declarations.transformInPlace(transformer, data) thisReceiver = thisReceiver?.transform(transformer, data) }
IrClass 的 typeParameters 等成员都是 var 类型，在 transform 之后被更新为新的值
基于 IR 生成代码
当我们想在编译期插入代码时，可以自定义 IrElementTransformer 并在其中完成对 IrElement 的更新。IrElement 的更新需要借助 IrFactory 。IrFacoty 用来创建 IrClass, IrSimpleFunction, IrProperty 等各类 IrElement。
另外，各 IrElement 都有一些扩展函数，封装了对 IrFactory 的调用，可以更加方便的更新自身。
例如，我们可以通过 IrFunction 的扩展函数 addValueParameter 为函数添加参数：
//org.jetbrains.kotlin.ir.builders.declarations.declarationBuilders inline fun IrFunction.addValueParameter(builder: IrValueParameterBuilder.() -> Unit): IrValueParameter = IrValueParameterBuilder().run { builder() if (index == UNDEFINED_PARAMETER_INDEX) { index = valueParameters.size } //调用 factory.buildValueParameter 创建新的 ValueParameter，添加为新的参数 factory.buildValueParameter(this, this@addValueParameter).also { valueParameter -> valueParameters = valueParameters + valueParameter } } fun IrFactory.buildValueParameter(builder: IrValueParameterBuilder, parent: IrDeclarationParent): IrValueParameter = with(builder) { //调用IrFactory#createValueParameter 基于 IrValueParameterBuilder 信息创建 parameters return createValueParameter( startOffset, endOffset, origin, IrValueParameterSymbolImpl(), name, index, type, varargElementType, isCrossInline, isNoinline, isHidden, isAssignable ).also { it.parent = parent } }
其内部通过调用 IrFactory 的 createValueParameter 创建新的参数。创建参数的信息来自 IrValueParameterBuilder。
我们知道 Compose Compiler 会为 Composable 函数增加 $composer 参数，这就是调用 IrFunction#addValueParameter 来实现的。
//androidx.compose.compiler.plugins.kotlin.lower.ComposerParamTransformer // $composer val composerParam = fn.addValueParameter { name = KtxNameConventions.COMPOSER_PARAMETER type = composerType.makeNullable() origin = IrDeclarationOrigin.DEFINED isAssignable = true }
在 IrValueParameterBuilder 的尾 lambda 中，添加了参数信息:
- name : Name.identifier("$composer")
- type : composerIrClass.defaultType.replaceArgumentsWithStarProjections()
- origin: IrDeclarationOrigin.DEFINED
IrDeclarationOrigin 是所有 IrDeclaration 都会有的一个属性，表示当前 IrDeclaration 的来源，DEFINED 表示这个 IrValueParameter 来自参数定义。
IrFactory 可以比较方便地构建 IrDeclaration 实例，例如 IrClass，IrFunction，IrValueParameter 等，但是在构建一些语句或者表达式时还需要有IrBuilder 或者 IrBuilderWithScope 的辅助， 它们可以提供更多生成指定 IrElemnet 的扩展方法。
Comopse Compiler 会为 Composable 函数增加 startRestartGroup/endRestartGroup ，服务于当前 Composable 的重组
// Before compiler (sources) @Composable fun RestartGroupTest(str: String) { Text(str) } // After compiler @Composable fun RestartGroupTest(str: String, $composer: Composer<*>, $changed: Int) { $composer.startRestartGroup(-846332013) // ... Text(str) $composer.endRestartGroup()?.updateScope { next -> RestartGroupTest(str, next, $changed or 0b1) } }
重组是通过递归调用 RestartGroupTest 自身实现的，这里需要构建一个递归调用的 lambda 传入 updateScope，这个 lambda 的构建就使用到了 IrBuilderWithScope。
//androidx.compose.compiler.plugins.kotlin.lower.ComposableFunctionBodyTransformer val localIrBuilder = DeclarationIrBuilder(context, fn.symbol) //... fn.body = localIrBuilder.irBlockBody { // Call the function again with the same parameters +irReturn( irCall(function.symbol).apply { symbol.owner .valueParameters .forEachIndexed { index, param -> if (param.isVararg) { putValueArgument( index, IrVarargImpl( UNDEFINED_OFFSET, UNDEFINED_OFFSET, param.type, param.varargElementType!!, elements = listOf( IrSpreadElementImpl( UNDEFINED_OFFSET, UNDEFINED_OFFSET, irGet(param) ) ) ) ) } else { // NOTE(lmr): should we be using the parameter here, // with the default value? putValueArgument(index, irGet(param)) } } // 更新 composer 参数 putValueArgument( composerIndex, irGet(fn.valueParameters[0]) ) //... } ) }
fn 是待构建的 lambda。DeclarationIrBuilder 创建一个 IrBuilderWithScope，然后通过其扩展方法 irBlockBody 为 lambda 创建函数体实现。
- 将构建的 IrElemnet 插入到当前函数体中，irCall(function.symbol) 生成递归调用当前 function 的代码。这里面使用 putValueArgument 添加递归调用的参数。 IrBuilderWithScope#irBlockBody 等扩展方法会简化一些 codegen 过程中的样板代码，让我们更加聚焦在关键代码的生成上。
ComposeIrGenerationExtension
Compose Compiler 中基于 IR 的代码生成主要是依靠 ComposeIrGenerationExtension 实现的。
class ComposeIrGenerationExtension( //... ) : IrGenerationExtension { override fun generate( moduleFragment: IrModuleFragment, pluginContext: IrPluginContext ) { ClassStabilityTransformer( pluginContext, symbolRemapper, metrics ).lower(moduleFragment) ComposerParamTransformer( pluginContext, symbolRemapper, decoysEnabled, metrics, ).lower(moduleFragment) //... 更多 Lower 调用 } }
在插件回调中我们拿到 IR 的根 moduleFragment，然后使用不同的 Transformer 对其进行 Lowering。
Lowering 是编译原理中的概念，在编译过程的每一步，使用不同的 Lower 来来逐步删除高级特征，并生成更低级别特征的 IR。这样做足够多次，最终你会得到一个足够简单的IR，可以编译到汇编级。
在这里我们可以简单理解为不断调用 Transformer 对 IR 进行更新，直至生成最终的代码。 所有的 Transformer 都存在于 androidx.compose.compiler.plugins.kotlin.lower 包下，其中主要的几个如下，这些也是后续对 Compose 代码生成的重点研究对象。
- ClassStabilityTransformer：判断类型稳定性并针对稳定类型生成跳过重组的对应代码
- ComposerParamTransformer: 为 Composable 函数增加 $compsoer 等参数
- ComposableFunctionBodyTransformer：Composable 函数体内生成 startXXXGroup/endXXXGroup 等相关代码
