LoginSignup
14
3

More than 1 year has passed since last update.

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

Last updated at Posted at 2022-12-21

こんにちは!tkhsktです。

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

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

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

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

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

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

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

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

ゴール

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

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

@Shiirudo
sealed class Event {
    object ShowModal : Event()
    object DismissModal : Event()
    data class ShowToast(val message: String) : Event()
}
fun handleEvent(event: Event) {
    eventShiirudo {
        event
    }.isShowModal { showModal ->
        // 何か
    }.isDismissModal { dismissModal ->
        // 何か
    }.isShowToast { showToast ->
        // 何か
    }.execute()
}

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

生成するコード

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

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のサブクラスに対応する処理を設定できるようにします。

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

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

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

コードを生成する

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

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

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

となります。

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

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

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

Processorの作成

プロジェクトにKSPを導入するための設定は公式のガイドに詳しく記載されているのでそちらをご参照ください。

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

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というライブラリを使用します。KotlinPoetはKotlinのコードを生成するためのライブラリで、KSPと組み合わせるための相互運用APIも提供しています。

Executorクラスの生成

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

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

長くなるので省略していますが、実際のコード生成処理は下記のように行います。(詳しくはKotlinPoetのドキュメントを参照してください)

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を利用するためのメソッドの作成

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

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式をフラットにする方法を紹介しました。正直、ここまでやっても得られるメリットは少ない(むしろ後から剥がしにくいなどのデメリットの方が多そう)ですが、気になった方はぜひ試してみてください。

冒頭に紹介したライブラリでは、この記事で紹介した以外のアプローチでもsealed classやsealed interfaceを扱えるようにしているので、興味のある方は是非見てみてください(ニッコリ)

14
3
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
14
3