孫小逗

一壶浊酒空对月

0%

Sword - 为 Kotlin 函数增加代理功能(四) - Kotlin IR

发表于 更新于 分类于
本文字数： 18k 阅读时长 ≈ 31 分钟
Sword - 为 Kotlin 函数增加代理功能(四) - Kotlin IR

Sword - 为 Kotlin 函数增加代理功能(四) - Kotlin IR

简介

Sword：一个可以给 Kotlin 函数增加代理的第三方库，基于 KCP 实现。

  1. Sword - 为 Kotlin 函数增加代理功能(一)
  2. Sword - 为 Kotlin 函数增加代理功能(二)
  3. Sword - 为 Kotlin 函数增加代理功能(三)

前面三篇文章笔者记录了 Sword 的实现过程，如何使用 Sword 以及如何通过 KSP 为 InvocationHandler 生成 FqName 索引类 HandlerFqName。

在第三篇文章的最后笔者有一个新的想法：通过 Kotlin IR 重新实现 Sword 的功能。经过最近几天晚上和早晨的努力，笔者初步实现了 Sword 的功能，可能还有一些问题，但是效果已经达到了笔者的预期，遂本篇文章记录下笔者的实现过程。

Kotlin IR 是什么以及可以做什么，本文不再赘述，网上有不少资料，读者可以自行参考。

预期效果

假设有以下类(GetTextNoArgInvocationHandler)和函数(getTextNoArg())：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
@ProxyHandler("GET_TEXT_NO_ARG")
class GetTextNoArgInvocationHandler : InvocationHandler {

    private val TAG = GetTextNoArgInvocationHandler::class.java.simpleName

    override fun invoke(className: String, methodName: String, args: Array<Any?>): Any? {
        Log.e(TAG, "invoke: className = $className, methodName = $methodName, args(${args.size}) = ${args.joinToString()}")
        return "guodongAndroid-Debug"
    }
}

// -----------------------------------------------------------------------------------------------------------------------

@Proxy(
    enable = true,
    handler = HandlerFqName.GET_TEXT_NO_ARG
)
fun getTextNoArg() = "guodongAndroid"

通过 Kotlin IR 编译和 Sword 代理后，笔者期望 getTextNoArg 函数转换成类似下面的伪代码：

1
2
3
4
5
6
7
@Proxy(
    enable = true,
    handler = HandlerFqName.GET_TEXT_NO_ARG
)
fun getTextNoArg(): String {
    return GetTextNoArgInvocationHandler().invoke("Test", "getTextNoArg", emptyArray()) as String
}

SwordComponentRegistrar

要使用 IR 首先需要注册 IrGenerationExtension 扩展，修改之前 SwordComponentRegistrar 中的代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class SwordComponentRegistrar : ComponentRegistrar {
    override fun registerProjectComponents(
        project: MockProject,
        configuration: CompilerConfiguration
    ) {
        val messageCollector =
	        configuration.get(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, MessageCollector.NONE)

        /*ClassBuilderInterceptorExtension.registerExtension(
            project,
            SwordClassGenerationInterceptor(messageCollector)
        )*/

        IrGenerationExtension.registerExtension(
            project,
            SwordIrGenerationExtension(messageCollector)
        )
    }
}

在上面的代码中：

  1. 注释掉之前通过 ASM 修改字节码的 ClassBuilderInterceptorExtension 扩展，
  2. 新增 IrGenerationExtension 扩展。

Dump

在 IR 语法树中，所有的节点都实现了 IrElement 接口，这些节点可以是模块，包，文件，类，属性，函数，参数，表达式，函数调用、函数体等等。

那么这些节点是什么样子的呢？实现我们的 SwordIrGenerationExtension：

1
2
3
4
5
6
7
8
9
10
class SwordIrGenerationExtension(
    private val messageCollector: MessageCollector,
) : IrGenerationExtension {
    override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
        messageCollector.report(
            CompilerMessageSeverity.WARNING,
            moduleFragment.dump()
        )
    }
}

使用 IrElement 的扩展函数 dump 可以输出这些节点的语法树信息。

在上面的代码中，我们输出了整个模块节点的语法树信息，如果模块中有许多文件，类，那么这些信息是相当庞大的，Sword 的目标是 Kotlin 函数，所以此处笔者不再贴出模块节点的语法树信息，等到我们转换函数时，再看看函数节点的语法树信息。

假设我们不知道如何编写 IR 编译器插件的代码，我们可以先写出要实现效果的 Kotlin 代码，再借助 dump 函数输出 IR 语法树信息，参考且对比语法树信息进行开发 IR 编译器插件，所以 dump 在开发 IR 编译器插件时非常有用。

IrElement

前面说了很多 IrElement 接口，目前我们还不知道它的真面目，接下来让我们看看它吧：

1
2
3
4
5
6
7
8
9
interface IrElement {
    ...

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

    fun <D> transform(transformer: IrElementTransformer<D>, data: D): IrElement
    fun <D> transformChildren(transformer: IrElementTransformer<D>, data: D)
}

在 IrElement 中有四个接口函数，其中 accept 函数基于访问者模式访问各个节点，transform 函数又是基于 accept 函数提供修改节点语法树的能力。

accept 函数中的参数 IrElementVisitor 接口提供了访问各个节点的函数：

1
2
3
4
5
6
7
8
9
10
11
12
interface IrElementVisitor<out R, in D> {
    
    fun visitElement(element: IrElement, data: D): R
    fun visitModuleFragment(declaration: IrModuleFragment, data: D) = visitElement(declaration, data)
    fun visitPackageFragment(declaration: IrPackageFragment, data: D) = visitElement(declaration, data)
    fun visitFile(declaration: IrFile, data: D) = visitPackageFragment(declaration, data)
    
    fun visitDeclaration(declaration: IrDeclarationBase, data: D) = visitElement(declaration, data)
    fun visitClass(declaration: IrClass, data: D) = visitDeclaration(declaration, data)
    fun visitFunction(declaration: IrFunction, data: D) = visitDeclaration(declaration, data)
    ...
}

IrElementVisitor 中有非常多接口函数，上面代码片段列举了一些接口函数，这些接口函数大多数都有默认实现，仔细观察这些函数的默认实现，最后都直接或间接的调用到 visitElement 函数。

transform 函数中的参数 IrElementTransformer 接口继承自 IrElementVisitor 接口，IrElementTransformer 接口主要是实现了 IrElementVisitor 的接口函数并调用 IrElement.transformChildren 函数遍历节点修改节点的语法树：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
interface IrElementTransformer<in D> : IrElementVisitor<IrElement, D> {
    override fun visitElement(element: IrElement, data: D): IrElement {
        element.transformChildren(this, data)
        return element
    }

    override fun visitModuleFragment(declaration: IrModuleFragment, data: D): IrModuleFragment {
        declaration.transformChildren(this, data)
        return declaration
    }

    override fun visitFile(declaration: IrFile, data: D): IrFile {
        declaration.transformChildren(this, data)
        return declaration
    }
}

在 Sword 中，我们主要利用 transform 函数修改节点语法树的能力来实现为 Kotlin 函数增加代理功能。

IrType & IrSymbol

在 IR 中不仅有 IrElement 还有 IrType 和 IrSymbol 。那么这两个有什么作用呢？

IrType

IrType 可以说是 KotlinType 在 IR 中的另一种表现形式，表示 Kotlin 中的各种类型，比如 Any，Boolean，Int，String 等等。IrType 常用在比较函数参数类型，调用函数时传入参数类型等，举个 Sword 中的栗子：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
private val anyNType = pluginContext.irBuiltIns.anyNType
private val stringType = pluginContext.irBuiltIns.stringType
private val arrayAnyNType = pluginContext.irBuiltIns.arrayClass.typeWith(anyNType)

// `InvocationHandler.invoke(className: String, methodName: String, args: Array<Any?>): Any?`
val invokeSymbol =
	pluginContext.referenceFunctions(FqName("${param.handler}.$INVOKE_METHOD_NAME"))
		.single {
    		val valueParameters = it.owner.valueParameters
    		valueParameters.size == 3 &&
    		valueParameters[0].type == stringType &&
    		valueParameters[1].type == stringType &&
    		valueParameters[2].type == arrayAnyNType
		}

在上面的代码片段中 pluginContext 是 IrGenerationExtension.generate 函数中的第二个参数，通过它的 irBuiltIns 字段我们获取一些 Kotlin 内置的 IrType。

  • anyNType 表示 Kotlin 中的 Any? 类型，
  • stringType 表示 String 类型，如果想获取 String? 类型，则需要调用 stringType.makeNullable() ，
  • arrayAnyNType 表示 Array<Any?> 类型，Array 没有对应的 IrType 表示，我们需要先获取 Array 对应的 IrSymbol，再调用 typeWith 扩展函数传入所需泛型的 IrType 即可获取 Array<Any?> 的 IrType。

在上面代码片段的最后，笔者的目标是获取 InvocationHandler.invoke 函数的 IrSymbol，考虑到开发者可能会重载 invoke 函数，所以笔者增加了以下判断逻辑：

  1. 函数中有且仅有三个参数，
  2. 第一个参数的类型必须是 String 类型，
  3. 第二个参数的类型必须是 String 类型，
  4. 第三个参数的类型必须是 Array<Any?> 类型。

满足以上几个条件的函数笔者才认为是 InvocationHandler.invoke 函数。

IrSymbol

IrSymbol 在 Kotlin IR 中可以算是比较重要的一个接口了。它以「符号」的形式描述了 Kotlin 的包、文件、类、函数、属性、字段等，笔者把它理解为 Java 字节码中的描述符，所以 IrSymbol 常用在创建类、函数、属性，函数调用等，举个 Sword 中的栗子：

1
2
3
4
5
6
private val anyNType = pluginContext.irBuiltIns.anyNType
private val emptyArraySymbol = pluginContext.referenceFunctions(FqName("kotlin.emptyArray")).first()

irCall(emptyArraySymbol).also {
    it.putTypeArgument(0, anyNType)
}
  • anyNType 在上节中出现过，它表示 Any? 类型，
  • emptyArraySymbol 是 emptyArray() 函数在 IR 中的符号，我们同样可以通过 irBuiltIns 获取一些 Kotlin 内置的 IrSymbol，其他的符号可以通过 pluginContext.referenceXXX() 的一系列函数查找。
  • 接下来调用 irCall 函数并传入 emptyArraySymbol，最后调用 putTypeArgument() 函数设置 emptyArray() 函数的泛型。

所以上面的代码片段其实是调用 Kotlin 中的 emptyArray<Any?>() 函数。

Sword

SwordIrGenerationExtension

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class SwordIrGenerationExtension(
    private val messageCollector: MessageCollector,
) : IrGenerationExtension {

    private val proxyAnnotationFqName = FqName("com.guodong.android.sword.api.kt.Proxy")

    override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
        // 输出日志
        messageCollector.report(
            CompilerMessageSeverity.WARNING,
            "Welcome to guodongAndroid sword kcp kotlin ir plugin"
        )

        // 判断当前模块是否存在 `Proxy` 注解
        val proxyAnnotation = pluginContext.referenceClass(proxyAnnotationFqName)
        if (proxyAnnotation == null) {
            messageCollector.report(
                CompilerMessageSeverity.ERROR,
                "Not found `Proxy` annotation, make sure to add the \"sword-api-kt\" library to your dependencies"
            )
            return
        }

        // 开始转换
        moduleFragment.transform(
            SwordTransformer(pluginContext, proxyAnnotation, messageCollector),
            null
        )
    }
}

经过前面的知识铺垫， SwordIrGenerationExtension 中的代码逻辑相信读者应该能理解了，笔者这里就不再赘述了，没有理解的读者可以再回顾下前面的内容。

接下来我们主要看看 SwordTransformer 中的逻辑。

SwordTransformer

因为 Sword 的功能是为 Kotlin 函数增加代理功能，所以在 SwordTransformer 中我们仅关注与函数相关的转换函数，即：visitFunctionNew(declaration: IrFunction) 函数。

首先声明一些变量和常量，其中一些变量前面笔者已经介绍过：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
companion object {
    // `InvocationHandler` 中的 `invoke` 函数名称
    private const val INVOKE_METHOD_NAME = "invoke"
}

private val anyNType = pluginContext.irBuiltIns.anyNType
private val stringType = pluginContext.irBuiltIns.stringType
private val arrayAnyNType = pluginContext.irBuiltIns.arrayClass.typeWith(anyNType)
private val emptyArraySymbol = pluginContext.referenceFunctions(FqName("kotlin.emptyArray")).first()
private val arrayOfSymbol = pluginContext.irBuiltIns.arrayOf

// @JvmName 注解完全限定名
private val jvmNameAnnotationFqName = FqName("kotlin.jvm.JvmName")

// @Proxy 注解完全限定名
private val proxyAnnotationFqName = FqName("com.guodong.android.sword.api.kt.Proxy")

下面我们就覆写 visitFunctionNew() 函数，在覆写的 visitFunctionNew() 函数中有较多的代码逻辑，还是老习惯，笔者先描述下自己的实现思路，然后再根据实现思路依次进行代码实现：

  1. 过滤一些笔者认为不需要处理的函数，读者可以自行斟酌；过滤不包含 Proxy 注解的函数，
  2. 获取 Proxy 注解里的数据存储于第二篇文章中的 SwordParam 中，
  3. 获取当前类名和函数名，判断当前函数是否启用代理，如果启用了代理但是 handler 为空则抛出异常，
  4. 启用代理后，我们直接抛弃原函数体，生成新的代理函数体。

1.过滤

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
override fun visitFunctionNew(declaration: IrFunction): IrStatement {
    // 过滤挂起函数，内联函数，多平台声明函数，外部函数(JNI)
    if (declaration.isSuspend || declaration.isInline || declaration.isExpect || declaration.isExternal) {
        return super.visitFunctionNew(declaration)
    }

    if (declaration is IrSimpleFunction) {
        // 过滤中缀函数，尾递归函数，操作符函数
        if (declaration.isInfix || declaration.isTailrec || declaration.isOperator) {
            return super.visitFunctionNew(declaration)
        }
    }

    // 过滤函数体为空的函数，不包含 `Proxy` 注解的函数
    if (declaration.body == null || !declaration.hasAnnotation(annotationClass)) {
        return super.visitFunctionNew(declaration)
    }
    
    ......
}
  1. 过滤挂起函数，内联函数，多平台声明函数，外部函数(JNI)，
  2. 过滤中缀函数，尾递归函数，操作符函数，
  3. 过滤函数体为空的函数，不包含 Proxy 注解的函数。

2.获取

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
override fun visitFunctionNew(declaration: IrFunction): IrStatement {
    ......

    val param = SwordParam()
    param.hasProxyAnnotation = true

    // 获取函数上 `Proxy` 注解的 `IrElement`, 返回 `IrConstructorCall`
    val irProxyConstructorCall = declaration.annotations.filter {
        it.isAnnotation(proxyAnnotationFqName)
    }.toList().single()

    // 获取 `Proxy.enable` 属性的值并存储
    val enableParam = irProxyConstructorCall.getValueArgument(0)
    enableParam?.let {
        if (it is IrConst<*>) {
            param.enable = it.value as Boolean
        }
    }

    // 获取 `Proxy.handler` 属性的值并存储
    val handlerParam = irProxyConstructorCall.getValueArgument(1)
    handlerParam?.let {
        if (it is IrConst<*>) {
            param.handler = it.value as String
        }
    }

    ......
}

获取 Proxy 注解中的数据，比较麻烦一些，一开始笔者认为获取到的注解可能类似于 IrAnnotation，然而发现却是 IrConstructorCall，后面仔细想来注解不就是通过构造函数构建一个注解实例么？我们在注解中传入的参数都是赋值给了其构造函数的属性。

通过 getValueArgument() 函数根据注解中声明属性的顺序获取其对应的属性，因为属性可能有默认值，我们在使用时可以不传入某个属性，比如：

1
2
3
4
5
@Proxy(
    // enable = true,
    handler = HandlerFqName.GET_TEXT_NO_ARG
)
fun getTextNoArg() = "guodongAndroid"

在 Proxy 注解中 enable 默认为 True，所以我们在使用时可以不传入 enable 的值，使其使用默认值，但是对于 kotlin IR 来说，没有明确使用的属性，通过 getValueArgument() 函数获取到的为 null，因为在 Kotlin IR 语法树中找不到这个属性。

由于注解中的属性值必须是编译期常量，所以我们可以把 handlerParam 转换为 IrConst 并获取它的值。

3.校验

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
override fun visitFunctionNew(declaration: IrFunction): IrStatement {
    ......

    // 获取 ClassName
    val className: String = getClassName(declaration)
    
    // 获取 MethodName
    val methodName = declaration.name.asString()

    // 校验 开启代理后是否注入了 `handler`
    if (param.enable && (param.handler.isEmpty() || param.handler.isBlank())) {
        messageCollector.report(
            CompilerMessageSeverity.ERROR,
            "[$className.$methodName]启用代理后请注入`handler`",
        )
    }

    ......
}

private fun getClassName(
    declaration: IrFunction,
): String {
    val parentClassOrNull = declaration.parentClassOrNull
    val fileOrNull = declaration.fileOrNull

    return when {
        declaration.isLocal -> {
            // 本地方法: 类名.函数名.<anonymous>
            // 源码中有此逻辑, 逻辑较为繁琐，且不是 `Sword` 的核心逻辑，本文就不记录了
        }
        parentClassOrNull != null -> {
            // 获取类名
            parentClassOrNull.name.asString()
        }
        fileOrNull != null -> {
            // 如果是顶级函数，获取文件名或`JvmName`注解指定的名字
            val annotations = fileOrNull.annotations
            if (annotations.hasAnnotation(jvmNameAnnotationFqName)) {
                val annotation = annotations.findAnnotation(jvmNameAnnotationFqName)!!
                val expression = annotation.getValueArgument(0)
                if (expression != null && expression is IrConst<*>) {
                    expression.value as String
                } else {
                    fileOrNull.name
                }
            } else {
                fileOrNull.name
            }
        }
        else -> "Unknown"
    }
}

获取 ClassName 时有以下几点考虑：

  1. 首先判断是否是本地方法，如果是本地方法则获取 类名.方法名.[<anonymous>]，
  2. 其次获取当前函数所在的父级 IrClass，如果不为 null，则使用类名，
  3. 最后获取函数所在的 IrFile，如果不为 null，再判断文件上是否有 JvmName 注解，有的话使用 JvmName 注解指定的名字，否则使用文件名。

获取 MethodName 时直接使用了函数名称，此处没有判断函数上是否有 JvmName 注解逻辑，相信读者可以自行扩展。

下面就是一个开启代理后必须注入 handler 的校验逻辑。

4.转换

在前面介绍 dump 函数时我们并没有实际上看看 dump 函数的输出内容，接下来让我们看看它输出的语法树信息。

以下面的函数为例：

1
2
3
4
5
6
7
@Proxy(
    enable = true,
    handler = HandlerFqName.GetTextArgInvocationHandler
)
fun testHandler(): User {
    return GetTextArgInvocationHandler().invoke("Test", "testHandler", emptyArray()) as User
}

dump 函数的输出结果如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
FUN name:testHandler visibility:public modality:FINAL <> ($this:com.guodong.android.sword.app.Test) returnType:com.guodong.android.sword.app.User
  annotations:
    Proxy(enable = 'true', handler = 'com.guodong.android.sword.app.GetTextArgInvocationHandler')
  $this: VALUE_PARAMETER name:<this> type:com.guodong.android.sword.app.Test
  BLOCK_BODY
    RETURN type=kotlin.Nothing from='public final fun testHandler (): com.guodong.android.sword.app.User declared in com.guodong.android.sword.app.Test'
      TYPE_OP type=com.guodong.android.sword.app.User origin=CAST typeOperand=com.guodong.android.sword.app.User
        CALL 'public open fun invoke (className: kotlin.String, methodName: kotlin.String, args: kotlin.Array<kotlin.Any?>): kotlin.Any? declared in com.guodong.android.sword.app.GetTextArgInvocationHandler' type=kotlin.Any? origin=null
          $this: CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in com.guodong.android.sword.app.GetTextArgInvocationHandler' type=com.guodong.android.sword.app.GetTextArgInvocationHandler origin=null
          className: CONST String type=kotlin.String value="Test"
          methodName: CONST String type=kotlin.String value="testHandler"
          args: CALL 'public final fun emptyArray <T> (): kotlin.Array<T of kotlin.ArrayIntrinsicsKt.emptyArray> [inline] declared in kotlin.ArrayIntrinsicsKt' type=kotlin.Array<kotlin.Any?> origin=null
            <T>: kotlin.Any?

fun testHandler(): User

1
FUN name:testHandler visibility:public modality:FINAL <> ($this:com.guodong.android.sword.app.Test) returnType:com.guodong.android.sword.app.User

这是 testHandler 函数的定义。它定义了函数的名称，可见性，模态以及类型签名。我们可以清楚的看到它是一个 public 和 final 名为 testHandler 的函数，并且它有一个隐含的参数 this（()），但是没有类型参数（<>），最后返回值为 com.guodong.android.sword.app.User。

类中的非静态函数(构造函数除外)都有一个隐含的 this 参数：

1
$this: VALUE_PARAMETER name:<this> type:com.guodong.android.sword.app.Test

GetTextArgInvocationHandler()

1
$this: CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in com.guodong.android.sword.app.GetTextArgInvocationHandler' type=com.guodong.android.sword.app.GetTextArgInvocationHandler origin=null

调用 GetTextArgInvocationHandler 的无参构造方法。

Test & testHandler

1
2
className: CONST String type=kotlin.String value="Test"
methodName: CONST String type=kotlin.String value="testHandler"

Test 和 testHandler 作为一个常量字符串。

emptyArray()

1
2
args: CALL 'public final fun emptyArray <T> (): kotlin.Array<T of kotlin.ArrayIntrinsicsKt.emptyArray> [inline] declared in kotlin.ArrayIntrinsicsKt' type=kotlin.Array<kotlin.Any?> origin=null
            <T>: kotlin.Any?

emptyArray() 函数是有类型参数的，因编译器推导我们在编写代码时可以省略，但是在 Kotlin IR 中明确显示了它的类型参数（<T>: kotlin.Any?），所以我们在转换代码时需要注意此处细节，否则转换代码在编译期不会报错，在运行时会抛出异常。

invoke("Test", "testHandler", emptyArray())

1
CALL 'public open fun invoke (className: kotlin.String, methodName: kotlin.String, args: kotlin.Array<kotlin.Any?>): kotlin.Any? declared in com.guodong.android.sword.app.GetTextArgInvocationHandler' type=kotlin.Any? origin=null

这是 invoke 函数的调用，清楚表示调用 invoke 函数需要三个参数以及参数的类型，同时函数返回值为 kotlin.Any?

as String

1
TYPE_OP type=com.guodong.android.sword.app.User origin=CAST typeOperand=com.guodong.android.sword.app.User

这是强转操作，TYPE_OP 表示一种类型操作，origin=CAST 表示类型操作符。

return

1
RETURN type=kotlin.Nothing from='public final fun testHandler (): com.guodong.android.sword.app.User declared in com.guodong.android.sword.app.Test'

在 Kotlin 中 return 其实是一个表达式，所以此处 return 的 type 是 kotlin.Nothing 。

通过上面 dump 函数输出语法树的分析，我们已经知道转换后函数体的语法信息，接下来让我们根据上面的分析依次来实现吧。

1
2
3
4
val handlerConstructorSymbol =
	pluginContext.referenceConstructors(FqName(param.handler)).single {
	    it.owner.valueParameters.isEmpty()
	}

首先通过 pluginContext.referenceConstructors() 查找 Handler 类的无参构造函数的符号。

1
2
3
4
5
6
7
8
9
val invokeSymbol =
	pluginContext.referenceFunctions(FqName("${param.handler}.$INVOKE_METHOD_NAME"))
		.single {
			val valueParameters = it.owner.valueParameters
			valueParameters.size == 3 &&
					valueParameters[0].type == stringType &&
					valueParameters[1].type == stringType &&
					valueParameters[2].type == arrayAnyNType
		}

接下来通过 pluginContext.referenceFunctions() 查找 Handler 类的 invoke 函数，可能存在函数重载，需要通过 single 函数确定我们需要的函数，这个在前面 IrType 举过例子。

1
2
3
4
5
6
7
8
9
10
11
12
13
private fun IrBuilderWithScope.irSwordArrayParams(function: IrFunction): IrCall {
    val parameters = function.valueParameters
    return if (parameters.isEmpty()) {
        irCall(emptyArraySymbol).also {
            it.putTypeArgument(0, anyNType)
        }
    } else {
        irCall(arrayOfSymbol).also {
            val expressions = parameters.map { parameter -> irGet(parameter) }
            it.putValueArgument(0, irVararg(anyNType, expressions))
        }
    }
}

上面代码是组装 invoke 函数的第三个参数 args: Array<Any?>，当前函数没有参数时使用 emptyArray<Any?>()，有参数时使用 arrayOf(vararg elements: T)。

1
2
3
4
5
6
val invokeCall = irCall(invokeSymbol).apply {
    dispatchReceiver = irCallConstructor(handlerConstructorSymbol, emptyList())
    putValueArgument(0, irString(className))
    putValueArgument(1, irString(methodName))
    putValueArgument(2, irSwordArrayParams(function))
}

构造 Handler 实例并调用它的 invoke 函数。

1
2
3
4
5
6
7
8
val irReturn = irReturn(
    typeOperator(
        resultType = function.returnType,
        argument = invokeCall,
        typeOperator = IrTypeOperator.CAST,
        typeOperand = function.returnType
    )
)

typeOperator 是语法树中的 TYPE_OP，就是强转操作，irReturn 表示 Kotlin 中的 return 表达式。

1
2
3
4
5
6
7
8
9
10
11
12
private fun irSword(
    function: IrFunction,
    param: SwordParam,
    className: String,
    methodName: String,
): IrBlockBody {
    return DeclarationIrBuilder(pluginContext, function.symbol).irBlockBody {
        ......
        
        +irReturn
    }
}

通过 + 操作符（unaryPlus）链接整个返回表达式到函数体，最后用新的函数体替换原函数体达到代理的功能：

1
2
3
4
5
6
7
8
9
override fun visitFunctionNew(declaration: IrFunction): IrStatement {
    ......

    if (param.enable) {
        declaration.body = irSword(declaration, param, className, methodName)
    }

    return super.visitFunctionNew(declaration)
}

至此，使用 Kotlin IR 为 Kotlin 函数增加代理功能完成。

总结

本文简单记录了通过 Kotlin IR 实现 Sword 代理功能的过程，同时简单介绍了一些 Kotlin IR 的 API 以及笔者对这些 API 的个人理解。

希望可以帮您开发自己的 Kotlin 编译器插件，happy~

相关文章

欢迎关注我的其它发布渠道