本文为稀土技术社区首发签约文章，14天内禁止转载，14天后未获授权禁止转载，侵权必究！

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!"

深入浅出 Compose Compiler（3） IR & IrGenerationExtension

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。

深入浅出 Compose Compiler（3） IR & IrGenerationExtension

别的，各 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，能够编译到汇编级。

深入浅出 Compose Compiler（3） IR & IrGenerationExtension

在这里咱们能够简略理解为不断调用 Transformer 对 IR 进行更新，直至生成最终的代码。 一切的 Transformer 都存在于 androidx.compose.compiler.plugins.kotlin.lower 包下，其中主要的几个如下，这些也是后续对 Compose 代码生成的重点研究目标。

  • ClassStabilityTransformer：判别类型安稳性并针对安稳类型生成越过重组的对应代码
  • ComposerParamTransformer: 为 Composable 函数添加 $compsoer 等参数
  • ComposableFunctionBodyTransformer：Composable 函数体内生成 startXXXGroup/endXXXGroup 等相关代码

参阅

  • Writing Your Second Kotlin Compiler Plugin
    blog.bnorm.dev/writing-you…

更多文章

  • 浅显易懂 Compose Compiler（1） Kotlin Compiler & KCP
  • 浅显易懂 Compose Compiler（2） 编译器前端查看