Zenn
🐶

ASTでKotlinに入門する

2024/12/06に公開
Kotlin
静的解析
tech

この記事はKotlin Advent Calendar 2024 6日目の記事です。

おはようございます。今年の10月にKotlinに入門したinouehiといいます、よろしくお願いします！

Kotlinに慣れるためにとにかく手を動かしてみようということで、ASTを操作する簡単なツールを作ってみることにしました。ASTを観察することで言語仕様が垣間見えることも期待しています。

想定読者

  • KotlinでASTを操作したい方
  • 静的解析に興味がある方

実装方法の調査・検討

右も左もわからない、KotlinもGradleもあんまりわからない状態だったのでひとまずChatGPTに聞いてみることにしました。

そこで出てきたキーワードがKotlinPoet, kotlinx.serialization, Kotlin Symbol Processing(KSP)やKotlin Compilerでした。

KotlinPoet, kotlinx.serializationはドキュメントに簡単に目を通して用途に合わなさそうだったので深入りは控えておきました。Kotlin Symbol Processingは簡単なcompiler pluginsを作るのに適しているようで、なんとなくコンパイラならいけそうかもという期待と、面白そうでもあったので試してみました。結果としてはKSPでASTをいじることが（僕には）できませんでした。しかしながら、KSPよりも自由度の高いKotlin Compiler Pluginというものがあることがわかり、試してみることにしました。

ドキュメントがほとんど見つけられずかなり苦しい中お世話になったのがこの記事でした。この記事を起点にして調査し、試行錯誤した結果、Kotlin Compiler PluginでASTをいじることに成功しました。また、ここに至る調査の過程でIntelliJプラグインとして実装することが可能であることもわかりました。

ChatGPTの助言は一通り試したので、次に静的解析ツールを参考にしてみようと思いました。そこでDetektに目をつけました。

https://github.com/detekt/detekt

Detektのコードを読み、参考にして書いたのが以下のコードです。

実験的実装

import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
import org.jetbrains.kotlin.com.intellij.openapi.Disposable
import org.jetbrains.kotlin.com.intellij.openapi.util.Disposer
import org.jetbrains.kotlin.com.intellij.psi.util.PsiUtilCore
import org.jetbrains.kotlin.config.CompilerConfiguration
import org.jetbrains.kotlin.psi.KtFile
import kotlin.system.exitProcess

fun main() {
    // 1. KotlinのコードからKtFileを生成する
    val environment = createKotlinCoreEnvironment()
    val virtualFile = requireNotNull(environment.findLocalFile("/path/to/IWannaBeAst.kt"))
    val psiFile = PsiUtilCore.getPsiFile(environment.project, virtualFile)
    if (psiFile !is KtFile) {
        exitProcess(1)
    }

    // 2. KtFileをトラバースする
    val visitor = CustomKtVisitor()
    psiFile.accept(visitor)
}

fun createKotlinCoreEnvironment(): KotlinCoreEnvironment {
    val configuration = CompilerConfiguration()
    val disposable: Disposable = Disposer.newDisposable()
    return KotlinCoreEnvironment.createForProduction(
        disposable,
        configuration,
        EnvironmentConfigFiles.JVM_CONFIG_FILES
    )
}

要点は、KtFile型のオブジェクトを入手することと、KtTreeVisitorVoidを拡張してvisitorを作ることです。visitorについては後述します。
KotlinのコードからKtFileを入手できます。
KtFileはKotlinのコードをノードの木として表現したものと見做せそうです。
KtFileにはacceptというメソッドがあり、これにvisitorを渡すことでノードの木を辿ることができました。
visitorにはノードに対して行いたい処理を宣言します。例えば、ノードのクラス名を表示する場合はこんな感じです。

import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
import org.jetbrains.kotlin.psi.KtElement

class CustomKtVisitor : KtTreeVisitorVoid() {
    private var indent = 0

    override fun visitKtElement(element: KtElement) {
        printIndent()
        println(element::class.simpleName)
        indent++
        element.children.forEach { it.accept(this) }
        indent--
    }

    private fun printIndent() {
        repeat(indent) { print("  ") }
    }
}
build.gradle.kts
plugins {
    kotlin("jvm") version "2.0.0"
    id("application")
}

group = "org.example"
version = "1.0-SNAPSHOT"

repositories {
    mavenCentral()
}

dependencies {
    api("org.jetbrains.kotlin:kotlin-compiler-embeddable")
    testImplementation(kotlin("test"))
}

tasks.test {
    useJUnitPlatform()
}
kotlin {
    jvmToolchain(17)
}

application {
    mainClass.set("org.example.MainKt")
}

それでは実行してみます。

まず、解析対象のコードはこちら↓（GPTに雑にお願いして作ってもらいました）

package org.example

import kotlin.math.sqrt

fun calculateSquareRoot(number: Double): Double {
    if (number < 0) {
        throw IllegalArgumentException("invalid")
    }
    return sqrt(number)
}

fun main() {
    val number = 16.0
    val result = calculateSquareRoot(number)
    println("The square root of $number is $result")
}
結果はこちら
KtPackageDirective
  KtDotQualifiedExpression
    KtNameReferenceExpression
    KtNameReferenceExpression
KtImportList
  KtImportDirective
    KtDotQualifiedExpression
      KtDotQualifiedExpression
        KtNameReferenceExpression
        KtNameReferenceExpression
      KtNameReferenceExpression
KtNamedFunction
  KtParameterList
    KtParameter
      KtTypeReference
        KtUserType
          KtNameReferenceExpression
  KtTypeReference
    KtUserType
      KtNameReferenceExpression
  KtBlockExpression
    KtIfExpression
      KtContainerNode
        KtBinaryExpression
          KtNameReferenceExpression
          KtOperationReferenceExpression
          KtConstantExpression
      KtContainerNodeForControlStructureBody
        KtBlockExpression
          KtThrowExpression
            KtCallExpression
              KtNameReferenceExpression
              KtValueArgumentList
                KtValueArgument
                  KtStringTemplateExpression
                    KtLiteralStringTemplateEntry
    KtReturnExpression
      KtCallExpression
        KtNameReferenceExpression
        KtValueArgumentList
          KtValueArgument
            KtNameReferenceExpression
KtNamedFunction
  KtParameterList
  KtBlockExpression
    KtProperty
      KtConstantExpression
    KtProperty
      KtCallExpression
        KtNameReferenceExpression
        KtValueArgumentList
          KtValueArgument
            KtNameReferenceExpression
    KtCallExpression
      KtNameReferenceExpression
      KtValueArgumentList
        KtValueArgument
          KtStringTemplateExpression
            KtLiteralStringTemplateEntry
            KtSimpleNameStringTemplateEntry
              KtNameReferenceExpression
            KtLiteralStringTemplateEntry
            KtSimpleNameStringTemplateEntry
              KtNameReferenceExpression

整然と並ぶノードがとても壮観ですね。

ただ、コードとの対応関係が少しわかりにくいです。 もう少しわかり易くしてみましょう。

コードとの対応関係を示すようにvisitorを作り変えて実行すると、、、

結果はこちら
KtPackageDirective
  KtDotQualifiedExpression
    KtNameReferenceExpression
    └ org
    KtNameReferenceExpression
    └ example
KtImportList
  KtImportDirective
    KtDotQualifiedExpression
      KtDotQualifiedExpression
        KtNameReferenceExpression
        └ kotlin
        KtNameReferenceExpression
        └ math
      KtNameReferenceExpression
      └ sqrt
KtNamedFunction
└ calculateSquareRoot
  KtParameterList
    KtParameter
    └ number
      KtTypeReference
        KtUserType
          KtNameReferenceExpression
          └ Double
  KtTypeReference
    KtUserType
      KtNameReferenceExpression
      └ Double
  KtBlockExpression
    KtIfExpression
      KtContainerNode
        KtBinaryExpression
          KtNameReferenceExpression
          └ number
          KtOperationReferenceExpression
          └ PsiElement(LT)
          KtConstantExpression
          └ 0
      KtContainerNodeForControlStructureBody
        KtBlockExpression
          KtThrowExpression
            KtCallExpression
              KtNameReferenceExpression
              └ IllegalArgumentException
              KtValueArgumentList
                KtValueArgument
                  KtStringTemplateExpression
                    KtLiteralStringTemplateEntry
                    └ invalid
    KtReturnExpression
      KtCallExpression
        KtNameReferenceExpression
        └ sqrt
        KtValueArgumentList
          KtValueArgument
            KtNameReferenceExpression
            └ number
KtNamedFunction
└ main
  KtParameterList
  KtBlockExpression
    KtProperty
    └ number
      KtConstantExpression
      └ 16.0
    KtProperty
    └ result
      KtCallExpression
        KtNameReferenceExpression
        └ calculateSquareRoot
        KtValueArgumentList
          KtValueArgument
            KtNameReferenceExpression
            └ number
    KtCallExpression
      KtNameReferenceExpression
      └ println
      KtValueArgumentList
        KtValueArgument
          KtStringTemplateExpression
            KtLiteralStringTemplateEntry
            └ The square root of 
            KtSimpleNameStringTemplateEntry
              KtNameReferenceExpression
              └ number
            KtLiteralStringTemplateEntry
            └  is 
            KtSimpleNameStringTemplateEntry
              KtNameReferenceExpression
              └ result

まとめ

Kotlin Compiler PluginやDetektではorg.jetbrains.kotlin.psi.*というKotlin用のASTを操作するためのクラスが使われるようでした。
一方、IntelliJプラグインでは com.intellij.psi.*やcom.intellij.lang.ASTNodeというKotlinに限らず汎用的に利用できるクラスが使われるようでした。

始め、ASTを簡単にいじってみたいだけなのにKotlin Compiler Pluginを実装しなければならないのかと、なんて重厚なんだと圧倒されかけましたが、最終的にはとてもシンプルなコードに落ち着きました。

以上、お読みいただきありがとうございましたmm

Discussion