コード生成の力でsealed class/interfaceのwhen式のネストを浅くする

こんにちは！tkhsktです。

Kotlinのsealed class/interfaceは便利なんですが、when式でサブクラスの分岐を作るとネストが深くなりますよね...
この記事では、sealed class/interfaceのwhen式をフラットにするコードの生成方法を紹介します。

今回紹介する方法はライブラリとして公開しているので、気になる方はそちらもご参照ください🙏
https://github.com/tkhskt/shiirudo

## 「ネストが深くなる」とは

この記事では下記のようなsealed classについて考えます。

```kotlin
sealed class Event {
    object ShowModal : Event()
    object DismissModal : Event()
    data class ShowToast(val message: String) : Event()
}
```

このsealed classのサブクラスをwhen式で分岐する場合、通常であれば下記のようになります。

```kotlin
fun handleEvent(event: Event) {
    when(event) {
        is Event.ShowModal -> {
            // 何かする
        }
        is Event.DismissModal -> {
            // 何かする
        }
        is Event.ShowToast -> {
            // 何かする
        }
    }
}
```

ご覧の通りwhen式のブロックの中で分岐するため、ネストが1段深くなります。

## ゴール

そんな訳で、この記事ではコードを生成してwhen式の分岐をフラットにすることを試みます。

目指す姿は下記の通りです。

```kotlin
@Shiirudo
sealed class Event {
    object ShowModal : Event()
    object DismissModal : Event()
    data class ShowToast(val message: String) : Event()
}
```

```kotlin
fun handleEvent(event: Event) {
    eventShiirudo {
        event
    }.isShowModal { showModal ->
        // 何か
    }.isDismissModal { dismissModal ->
        // 何か
    }.isShowToast { showToast ->
        // 何か
    }.execute()
}
```

このように対象のsealed class/interfaceにアノテーションを付与することで、分岐処理をフラットに書けるようにします。

### 生成するコード

上記のようにsealed class/interfaceを処理するためには下記のようなコードを生成します。

```kotlin
public class EventShiirudoExecutor(
  private val event: Event,
) {
  private var isDismissModal: ((Event.DismissModal) -> Unit)? = null

  private var isShowModal: ((Event.ShowModal) -> Unit)? = null

  private var isShowToast: ((Event.ShowToast) -> Unit)? = null

  private var isElse: (Event) -> Unit = {}

  public fun isDismissModal(f: (Event.DismissModal) -> Unit): EventShiirudoExecutor {
    this.isDismissModal = f
    return this
  }

  public fun isShowModal(f: (Event.ShowModal) -> Unit): EventShiirudoExecutor {
    this.isShowModal = f
    return this
  }

  public fun isShowToast(f: (Event.ShowToast) -> Unit): EventShiirudoExecutor {
    this.isShowToast = f
    return this
  }

  public fun isElse(f: (Event) -> Unit): EventShiirudoExecutor {
    this.isElse = f
    return this
  }

  public fun execute(): Unit {
    when(event) {
      is com.tkhskt.shiirudo.sample.qiita.Event.DismissModal -> {
        val f = this.isDismissModal ?: this.isElse
        f.invoke(event)
      }
      is com.tkhskt.shiirudo.sample.qiita.Event.ShowModal -> {
        val f = this.isShowModal ?: this.isElse
        f.invoke(event)
      }
      is com.tkhskt.shiirudo.sample.qiita.Event.ShowToast -> {
        val f = this.isShowToast ?: this.isElse
        f.invoke(event)
      }
    }
  }
}
```

Builderのようなクラス(ここではExecutorという名前にしています)を生成し、sealed class/interfaceのサブクラスに対応する処理を設定できるようにします。

そして、上記のクラスを利用するためのメソッドも生成します。

```kotlin
public fun eventShiirudo(block: () -> Event): EventShiirudoExecutor =
    EventShiirudoExecutor(block.invoke())
```

Note: ここまで挙げたコードは生成せずに自前で書くこともできます（面倒ですが）

## コードを生成する

ここからは具体的なコード生成の方法について説明します。

コード生成の大まかな流れは、

1. プロジェクト内からアノテーションが付けられているコードを収集する
1. 収集したコードの情報を元にコード/ファイルを生成する

となります。

### プロジェクト内からアノテーションが付けられているコードを収集する

ではまず、step1の「プロジェクト内からアノテーションが付けられているコードを収集する」方法について説明します。Kotlinのプロジェクトでは、kaptや[Kotlin Symbol Processing (KSP)](https://kotlinlang.org/docs/ksp-overview.html)といったアノテーション処理の技術を利用することができます。

この記事ではKSPを使用します。

#### Processorの作成

プロジェクトにKSPを導入するための設定は[公式のガイド](https://kotlinlang.org/docs/ksp-quickstart.html#make-ide-aware-of-generated-code)に詳しく記載されているのでそちらをご参照ください。

KSPではProcessorというクラスでアノテーションが付与されているコードを収集することができます。`@Shiirudo`というアノテーションが付与されたコードを収集するためのProcessorは下記のようになります。

```kotlin
class ShiirudoProcessor(
    private val codeGenerator: CodeGenerator,
    private val logger: KSPLogger,
) : SymbolProcessor {

    override fun process(resolver: Resolver): List<KSAnnotated> {
        val symbols =
            resolver.getSymbolsWithAnnotation(ClassName("com.tkhskt.shiirudo.annotation", "Shiirudo").canonicalName)
        val ret = symbols.filter { !it.validate() }.toList()
        symbols
            .filter { it is KSClassDeclaration && it.validate() }
            .forEach { it.accept(ShiirudoVisitor(), Unit) }
        return ret
    }

    inner class ShiirudoVisitor : KSVisitorVoid() {
        override fun visitClassDeclaration(classDeclaration: KSClassDeclaration, data: Unit) {
            /*
            Shiirudoアノテーションが付与されたコードについて下記をチェック
              - sealed classもしくはsealed interfaceか
              - サブクラスが存在するか
             */
            ShiirudoTargetValidator.validate(classDeclaration, logger)

            // コード生成処理
            val shiirudoExecutorGenerator =
                ShiirudoExecutorGenerator(codeGenerator, logger)
            shiirudoExecutorGenerator.generate(classDeclaration)
            val shiirudoExecutorMethodGenerator =
                ShiirudoExecutorMethodGenerator(codeGenerator, logger)
            shiirudoExecutorMethodGenerator.generate(classDeclaration)
        }
    }
}
```

ここでは、`getSymbolsWithAnnotation`というメソッドを使って`@Shiirudo`アノテーションが付与されたコードのリストを取得し、リストの各要素についてバリデーションの実行とコード生成処理を開始しています。

### 収集したコードの情報を元にコード/ファイルを生成する

コード生成には[KotlinPoet](https://github.com/square/kotlinpoet)というライブラリを使用します。KotlinPoetはKotlinのコードを生成するためのライブラリで、KSPと組み合わせるための相互運用APIも提供しています。

#### Executorクラスの生成

KotlinPoetを利用したコード（クラス）生成の大まかな流れは、

1. KSPのオブジェクトからアノテーションが付与されたクラスの情報（パッケージ名やクラス名）を取得
1. 生成するクラスのコンストラクタなどを設定
1. 生成するクラスにメソッドなどを追加
1. ファイルを作成して書き込む

長くなるので省略していますが、実際のコード生成処理は下記のように行います。（詳しくは[KotlinPoetのドキュメント](https://square.github.io/kotlinpoet/)を参照してください）

```kotlin
internal class ShiirudoExecutorGenerator(
    private val codeGenerator: CodeGenerator,
) {

    private lateinit var annotatedClassDeclaration: KSClassDeclaration
    private lateinit var annotatedClassName: ClassName
    private lateinit var subclasses: Sequence<KSClassDeclaration>

    
    fun generate(classDeclaration: KSClassDeclaration) {
        // KSPのオブジェクトから、アノテーションが付与されているクラスが存在するパッケージ名やクラス名などを取得
        val packageName = classDeclaration.packageName.asString()
        annotatedClassDeclaration = classDeclaration
        annotatedClassName = classDeclaration.toClassName()
        subclasses = classDeclaration.getSealedSubclasses()
        generateShiirudoExecutor(
            packageName = packageName,
            containingFile = classDeclaration.containingFile!!
        )
    }

    // コード生成処理
    private fun generateShiirudoExecutor(
        packageName: String,
        containingFile: KSFile,
    ) {
        val executorClassNameString = getExecutorClassName(annotatedClassDeclaration)
        val constructorFunSpec = FunSpec.constructorBuilder()
            .addConstructorParameters()
            .build()
        val typeSpec = TypeSpec.classBuilder(executorClassNameString)
            .primaryConstructor(constructorFunSpec)
            .addConstructorProperties()
            .addProperties()
            .addFunctions()
            .addExecuteFunction()
            .build()

        // ファイルの作成とコードの書き込み
        val file = FileSpec
            .builder(packageName, executorClassNameString.simpleName)
            .addType(typeSpec)
            .build()
        file.writeTo(codeGenerator, Dependencies(false, containingFile))
    }
    ...

    private fun TypeSpec.Builder.addConstructorProperties(): TypeSpec.Builder {
        val property = PropertySpec.builder(
            name = "event",
            type = annotatedClassName,
            modifiers = listOf(KModifier.PRIVATE)
        ).initializer("event").build()
        addProperty(property)
        return this
    }
    ...

    // クラスにメソッドを追加する
    private fun TypeSpec.Builder.addExecuteFunction(): TypeSpec.Builder {
        val branches = subclasses.map { subclass ->
            val nameSuffix = NameResolver.createPropertyName(
                rootDeclaration = annotatedClassDeclaration,
                classDeclaration = subclass,
            )
            subclass.toClassName().canonicalName to nameSuffix
        }.joinToString("\n") {
            """
            |  is ${it.first} -> {
            |    val f = this.is${it.second} ?: this.isElse
            |    f.invoke(event)
            |  }
            """.trimMargin()
        }
        addFunction(
            FunSpec.builder("execute")
                .addCode(
                    """
                    |when(event) {
                    |$branches
                    |}
                    """.trimMargin()
                )
                .build()
        )
        return this
    }
}
```

### Executorを利用するためのメソッドの作成

メソッドを生成する処理も大まかにはクラスの生成と同じです。

```kotlin
internal class ShiirudoExecutorMethodGenerator(
    private val codeGenerator: CodeGenerator,
    private val logger: KSPLogger,
) {

    private lateinit var annotatedClassDeclaration: KSClassDeclaration
    private lateinit var annotatedClassName: ClassName
    private lateinit var subclasses: Sequence<KSClassDeclaration>

    fun generate(classDeclaration: KSClassDeclaration) {
        val packageName = classDeclaration.packageName.asString()
        annotatedClassDeclaration = classDeclaration
        annotatedClassName = classDeclaration.toClassName()
        subclasses = classDeclaration.getSealedSubclasses()
        generateFunction(
            packageName = packageName,
            containingFile = classDeclaration.containingFile!!
        )
    }

    private fun generateFunction(
        packageName: String,
        containingFile: KSFile,
    ) {
        val namePrefix = ShiirudoExecutorGenerator.getExecutorClassName(annotatedClassDeclaration)
        val fileName = "${namePrefix.simpleName}Extension"
        val file = FileSpec
            .builder(packageName, fileName)
            .addExecuteFunction()
            .addLambdaExecuteFunction()
            .build()
        file.writeTo(codeGenerator, Dependencies(false, containingFile))
    }

    private fun FileSpec.Builder.addExecuteFunction(): FileSpec.Builder {
        val executorClassName =
            ShiirudoExecutorGenerator.getExecutorClassName(annotatedClassDeclaration)
        addFunction(
            FunSpec.builder("shiirudo")
                .returns(executorClassName)
                .receiver(annotatedClassName)
                .addCode(
                    """
                    |return ${executorClassName.simpleName}(this)
                    """.trimMargin()
                )
                .build()
        )
        return this
    }

    private fun FileSpec.Builder.addLambdaExecuteFunction(): FileSpec.Builder {
        val executorClassName =
            ShiirudoExecutorGenerator.getExecutorClassName(annotatedClassDeclaration)
        val lambdaTypeName = LambdaTypeName.get(
            receiver = null,
            returnType = annotatedClassName
        )
        val shiirudoClassNamePrefix =
            NameResolver.createPropertyName(
                rootDeclaration = null,
                classDeclaration = annotatedClassDeclaration,
                includeRoot = true
            ).lowerCamelCase()
        addFunction(
            FunSpec.builder("${shiirudoClassNamePrefix}Shiirudo")
                .returns(executorClassName)
                .addParameter("block", lambdaTypeName)
                .addCode(
                    """
                    |return ${executorClassName.simpleName}(block.invoke())
                    """.trimMargin()
                )
                .build()
        )
        return this
    }

    private fun String.lowerCamelCase(): String {
        return when (this.length) {
            0 -> ""
            1 -> this.lowercase()
            else -> this[0].lowercase() + this.substring(1)
        }
    }
}
```

長いですね...

## まとめ

この記事では、KSPとKotlinPoetを使ったコード生成によって、sealed classやsealed interfaceのサブクラスに対するwhen式をフラットにする方法を紹介しました。正直、ここまでやっても得られるメリットは少ない（むしろ後から剥がしにくいなどのデメリットの方が多そう）ですが、気になった方はぜひ試してみてください。

[冒頭に紹介したライブラリ](https://github.com/tkhskt/shiirudo)では、この記事で紹介した以外のアプローチでもsealed classやsealed interfaceを扱えるようにしているので、興味のある方は是非見てみてください(ﾆｯｺﾘ)