【Kotlin】動的プロキシでsuspend関数を実装する

この記事はKotlin Advent Calendar 2022の5日目の記事になりました。

https://qiita.com/advent-calendar/2022/kotlin

# TL;DR
- インターフェース定義が`suspend`関数かつ`InvocationHandler.invoke`関数内でも`suspend`関数を呼び出したい場合、素直に書くと`invoke`関数が非`suspend`関数なためコンパイルエラーになる
- `suspend () -> ...`が`(Continuation<*>) -> ...`にキャストでき、`Continuation`も`args`から取得できることを利用すれば、この問題は回避できる

# 前置き
動的プロキシを知っている方は読み飛ばして頂いて大丈夫です。

## 動的プロキシとは
動的プロキシとは、指定されたインタフェースを実行時に実装するような機能です。

例えば以下のようなことができます。

```kotlin
interface Foo {
    fun foo(): String
}

object FooProxyFactory {
    fun create(): Foo {
        /* 動的にFooを実装して返却 */
    }
}

// 動的に実装したFooを取得
val foo: Foo = FooProxyFactory.create()
```

これができると、例えば一定の規則に従ったインターフェースを作成するだけで、内部処理は自動実装するようなことができます。

詳しい内容は以下のドキュメントを参照下さい。

https://docs.oracle.com/javase/jp/8/docs/technotes/guides/reflection/proxy.html

## 動的プロキシを使ってみる
先ほどの`FooProxyFactory`を実装してみた内容が以下です。

```kotlin
import java.lang.reflect.InvocationHandler
import java.lang.reflect.Method
import java.lang.reflect.Proxy

class FooProxy(private val text: String) : InvocationHandler {
    override fun invoke(proxy: Any, method: Method, args: Array<out Any>?): Any? {
        // 対象がfoo関数であればtextを返す
        if (method.name == "foo") return text

        // それ以外の場合は一旦エラーにする
        TODO("Not yet implemented")
    }
}

object FooProxyFactory {
    fun create(): Foo {
        val fooClass = Foo::class.java
        val proxy = FooProxy("Hello from dynamic proxy.")

        @Suppress("UNCHECKED_CAST")
        return Proxy.newProxyInstance(
            fooClass.classLoader,
            arrayOf(fooClass),
            proxy
        ) as Foo
    }
}
```

これを実行すると以下のようになります。
実装クラスを書かずに、動的に実装できていることが分かります。

```kotlin
// 動的に実装したFooを取得
val foo: Foo = FooProxyFactory.create()
println(foo.foo()) // -> Hello from dynamic proxy.がprintされる
```

### 軽い解説
`InvocationHandler`は、動的プロキシに対する関数呼び出しが有った際に、それに対する処理を実装するクラスです。
`FooProxy`は、`InvocationHandler`を実装しています。

実際に何かを作る際には、`invoke`関数に処理を書きます。
このサンプルでは、「`foo`関数が実行された場合、設定された文字列を返却する」という挙動を実装しています。

`invoke`関数の引数は以下の通りです。

- `proxy`: メソッドが呼び出されるプロキシ・インスタンス
- `method`: 呼び出し対象の関数
- `args`: 引数配列（引数が与えられなかった場合は`null`になる(空配列にはならない)）

# 本題
## suspend関数が絡んだ場合に問題になること
インターフェースに定義された関数が`suspend`関数で、`invoke`関数での実装も`suspend`関数としたい場合を考えます。
この場合、`invoke`関数は非`suspend`関数なため、素直に呼び出すとコンパイルエラーが発生します。

```diff_kotlin
import java.lang.reflect.InvocationHandler

interface Foo {
-   fun foo(): String
+   suspend fun foo(): String // suspend化
}

class FooProxy(private val text: String) : InvocationHandler {
+   // suspend関数で何らかの処理を呼び出したい
+   private suspend fun calcFoo(text: String): String {
+       return text + " suspend."
+   }

    override fun invoke(proxy: Any, method: Method, args: Array<out Any>?): Any? {
-       if (method.name == "foo") return text
+       // calcFooはsuspend関数なため、コンパイルエラーになる！
+       if (method.name == "foo") return calcFoo(text)

        // それ以外の場合は一旦エラーにする
        TODO("Not yet implemented")
    }
}
```

## 対策
`FooProxy`を以下のように変更することで対応できます[^bikou]。

[^bikou]: 一応自分のユースケースでは動いたのでOKとしていますが、本当にこれで全パターンで大丈夫なのかは確認しきれていません。

```kotlin
import kotlin.coroutines.Continuation

class FooProxy(private val text: String) : InvocationHandler {
    private suspend fun calcFoo(text: String): String {
        return text + " suspend."
    }

    // suspend関数を非suspend関数から呼び出すための関数
    private fun <T> invokeSuspendFunction(
        continuation: Continuation<*>,
        block: suspend () -> T
    ): T = @Suppress("UNCHECKED_CAST") (block as (Continuation<*>) -> T)(continuation)

    override fun invoke(proxy: Any, method: Method, args: Array<out Any>?): Any? {
        if (method.name == "foo")  {
            // args.last()はsuspend関数へ暗黙的に渡されるContinuationの取得
            val continuation = args!!.last() as Continuation<*>
            return invokeSuspendFunction(continuation) {
                calcFoo(text)
            }
        }

        // それ以外の場合は一旦エラーにする
        TODO("Not yet implemented")
    }
}
```

実行すると以下のようになります。

```kotlin
import kotlinx.coroutines.runBlocking

val foo: Foo = FooProxyFactory.create()
runBlocking { println(foo.foo()) } // -> Hello from dynamic proxy. suspend.がprintされる
```

### 軽い解説
`Foo`インターフェースを`Java`にデコンパイルすると以下のようになっています。

```java:Fooインターフェースのデコンパイル結果
import kotlin.coroutines.Continuation;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

public interface Foo {
   @Nullable
   Object foo(@NotNull Continuation var1);
}
```

`foo`関数は元々引数無し関数でしたが、`Java`上では`Continuation`というパラメータが存在していることが分かります。
`Continuation`は`suspend`関数の制御に利用される値です。
このように、`suspend`関数は、`Java`上では`Continuation`を引数に取る形になります。

このため、`suspend () -> T`は`(Continuation<*>) -> T`にキャストすることができます。

```kotlin:suspend () -> Tから(Continuation<*>) -> Tにキャストする様子
private fun <T> invokeSuspendFunction(
    continuation: Continuation<*>,
    block: suspend () -> T
): T = @Suppress("UNCHECKED_CAST") (block as (Continuation<*>) -> T)(continuation)
```

この値は関数の実行時に渡されるため、`invoke`関数に渡される`args`から取得することができます。

```kotlin:argsからContinuationを取得する様子
val continuation = args!!.last() as Continuation<*>
```

これらを組み合わせることで、`invoke`関数内からの`suspend`関数呼び出しを実現しています。

これらの内容は以下のコードを参考に実装しました。

https://github.com/softappeal/yass2/blob/00de0edd7786f474226be524e66933e837af29be/yass2-reflect/src/jvmMain/kotlin/reflect/ReflectionInterceptor.kt#L24-L29