ASTでKotlinに入門する②

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

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

ASTでKotlinに入門するの続編です。前回はASTを参照してみましたが、今回はASTを改変してみたいと思います。自動リファクタリングツールのPoCを示す試みです。

想定読者

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

前置き

• IntelliJ Platform Pluginの実装の適切さやリファクタツールとしての実用性を度外視しています。
• 簡略化したコードでASTを改変する操作が可能であることを観察します。

PoCの実装

メソッドを根こそぎprivateに変換するという乱暴なリファクタリングを行うコードです。

import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.LocalFileSystem
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiManager
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiElementVisitor
import com.intellij.psi.PsiFile
import org.jetbrains.kotlin.lexer.KtTokens
import org.jetbrains.kotlin.psi.KtFile
import org.jetbrains.kotlin.psi.KtModifierListOwner
import org.jetbrains.kotlin.psi.KtNamedFunction

class RefactorAction : AnAction() {

    override fun actionPerformed(event: AnActionEvent) {
        val project = event.project ?: return

        // プロジェクト中のファイルをpsiFileとして読み込む
        val psiFiles = getAllPsiFiles(project)

        for (psiFile in psiFiles) {
            // Kotlinのファイルのみをトラバースしたい
            if (psiFile !is KtFile) {
                continue
            }

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

            // 簡易的な検証用にprintlnする
            // ファイルに書き出せばリファクタできる
            println(psiFile.text)
        }
    }
}

class CustomKtVisitor(private val project: Project) : PsiElementVisitor() {

    override fun visitElement(element: PsiElement) {
        if (element is KtNamedFunction) {
            WriteCommandAction.runWriteCommandAction(project) {
                // ASTを改変する
                (element as KtModifierListOwner).removeModifier(KtTokens.PUBLIC_KEYWORD)
                element.addModifier(KtTokens.PRIVATE_KEYWORD)
            }
        }
        element.children.forEach { it.accept(this) }
    }
}

getAllPsiFiles()

fun getAllPsiFiles(project: Project): List<PsiFile> {
    val basePath = project.basePath ?: return emptyList()
    val rootDirectory = LocalFileSystem.getInstance().findFileByPath(basePath)
    val files = collectFilesIteratively(rootDirectory!!)

    val psiManager = PsiManager.getInstance(project)
    return files.mapNotNull { psiManager.findFile(it) }
}

fun collectFilesIteratively(root: VirtualFile): List<VirtualFile> {
    val files = mutableListOf<VirtualFile>()
    val stack = ArrayDeque<VirtualFile>()
    stack.addFirst(root)

    while (stack.isNotEmpty()) {
        val file = stack.removeFirst()
        if (file.isDirectory) {
            file.children.forEach { stack.addFirst(it) }
        } else {
            files.add(file)
        }
    }
    return files
}

build.gradle.kts

build.gradle.ktsはKotlin Analysis APIを触ってみるを参考にさせていただきました。

plugins {
    id("java")
    id("org.jetbrains.kotlin.jvm") version "2.0.21"
    id("org.jetbrains.intellij.platform") version "2.1.0"
}

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

repositories {
    mavenCentral()

    intellijPlatform {
        defaultRepositories()
    }
}

dependencies {
    intellijPlatform {
        intellijIdeaCommunity("2024.3")

        instrumentationTools()

        bundledPlugin("org.jetbrains.kotlin")
    }
}

tasks {
    withType<JavaCompile> {
        sourceCompatibility = "21"
        targetCompatibility = "21"
    }
    withType<org.jetbrains.kotlin.gradle.tasks.KotlinCompile> {
        kotlinOptions.jvmTarget = "21"
    }

    patchPluginXml {
        sinceBuild.set("243")
        untilBuild.set("243.*")
    }

    signPlugin {
        certificateChain.set(System.getenv("CERTIFICATE_CHAIN"))
        privateKey.set(System.getenv("PRIVATE_KEY"))
        password.set(System.getenv("PRIVATE_KEY_PASSWORD"))
    }

    publishPlugin {
        token.set(System.getenv("PUBLISH_TOKEN"))
    }
}

src/main/resources/META-INF/plugin.xml

メニューバーのTools > Refactorから実行できるようにしています。

<idea-plugin>
    <actions>
        <action id="RefactorAction" class="org.example.astplugin.RefactorAction" text="Refactor" description="Refactor">
            <add-to-group group-id="ToolsMenu" anchor="last"/>
        </action>
    </actions>
</idea-plugin>

実装のポイント

要点は以下の通りです。

• 参照と違って改変したい場合はcom.intellij.psi.*を使う必要があるようでした。
• そのためにIntelliJ Platform Pluginとして実装しました。
• ASTを丸裸で改変するとエラーになりました。WriteCommandAction.runWriteCommandAction(project) {}で包むことで解決しました。

前回のASTでKotlinに入門するの延長でASTの改変を試みたところ、内部でcom.intellij.psi.*に依存しており単純には実現できないことがわかりました。そこでPluginとして実装することにしました。改変のロジックはシンプルでしたがビルドまわりで試行錯誤が必要でした。

動かしてみる

手元で挙動を確認するには./gradlew runIdeを実行してIDEを起動します。リファクタしたいコードを含むプロジェクトを開いて、メニューバーのTools > Refactorをクリックするとリファクタが実行されます。

リファクタ前

class The {
    fun world(): void {}
    public fun hand(): void {}
}

↓

リファクタ後

class The {
    private fun world(): void {}
    private fun hand(): void {}
}

まとめ

ひとまずASTを改変できることがわかりました。これを発展させてリファクタのロジックを作り込むことで、自動リファクタリングツールを作れそうなことがわかりました。

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